How to speed up this Rcpp function?

Question

I wish to implement a simple split-apply-combine routine in Rcpp where a dataset (matrix) is split up into groups, and then the groupwise column sums are returned. This is a procedure easily implemented in R, but often takes quite some time. I have managed to implement an Rcpp solution that beats the performance of R, but I wonder if I can further improve upon it. To illustrate, here some code, first for the use of R:

n <- 50000
k <- 50
set.seed(42)
X <- matrix(rnorm(n*k), nrow=n)
g=rep(1:8,length.out=n )

use.for <- function(mat, ind){
  sums <- matrix(NA, nrow=length(unique(ind)), ncol=ncol(mat))
  for(i in seq_along(unique(ind))){
    sums[i,] <- colSums(mat[ind==i,])
  }
  return(sums)
}

use.apply <- function(mat, ind){
  apply(mat,2, function(x) tapply(x, ind, sum))
}

use.dt <- function(mat, ind){ # based on Roland's answer
   dt <- as.data.table(mat)
   dt[, cvar := ind]
   dt2 <- dt[,lapply(.SD, sum), by=cvar]
   as.matrix(dt2[,cvar:=NULL])
}

It turns out that the for-loops is actually quite fast and is the easiest (for me) to implement with Rcpp. It works by creating a submatrix for each group and then calling colSums on the matrix. This is implemented using RcppArmadillo:

#include <RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
using namespace Rcpp;
using namespace arma;

// [[Rcpp::export]]
arma::mat use_arma(arma::mat X, arma::colvec G){

  arma::colvec gr = arma::unique(G);
  int gr_n = gr.n_rows;
  int ncol = X.n_cols;

  arma::mat out = zeros(gr_n, ncol); 

  for(int g=0; g<gr_n; g++){
   int g_id = gr(g);
   arma::uvec subvec = find(G==g_id);
   arma::mat submat = X.rows(subvec);
   arma::rowvec res = sum(submat,0);
   out.row(g) = res;     
  }
 return out;
}

However, based on answers to this question, I learned that creating copies is expensive in C++ (just as in R), but that loops are not as bad as they are in R. Since the arma-solution relies on creating matrixes (submat in the code) for each group, my guess is that avoiding this will speed up the process even further. Hence, here a second implementation based on Rcpp only using a loop:

#include <Rcpp.h>
using namespace Rcpp;

// [[Rcpp::export]]
NumericMatrix use_Rcpp(NumericMatrix X, IntegerVector G){

  IntegerVector gr = unique(G);
  std::sort(gr.begin(), gr.end());
  int gr_n = gr.size();
  int nrow = X.nrow(), ncol = X.ncol();

  NumericMatrix out(gr_n, ncol);

  for(int g=0; g<gr_n; g++){
     int g_id = gr(g);

      for (int j = 0; j < ncol; j++) {
      double total = 0;
        for (int i = 0; i < nrow; i++) {

          if (G(i) != g_id) continue;    // not sure how else to do this
          total += X(i, j);
        }
        out(g,j) = total;
      }
  }
      return out;
}

Benchmarking these solutions, including the use_dt version provided by @Roland (my previous version discriminted unfairly against data.table), as well as the dplyr-solution suggested by @beginneR, yields the following:

 library(rbenchmark)
 benchmark(use.for(X,g), use.apply(X,g), use.dt(X,g), use.dplyr(X,g), use_arma(X,g), use_Rcpp(X,g), 
+           columns = c("test", "replications", "elapsed", "relative"), order = "relative", replications = 1000)
             test replications elapsed relative
# 5  use_arma(X, g)         1000   29.65    1.000
# 4 use.dplyr(X, g)         1000   42.05    1.418
# 3    use.dt(X, g)         1000   56.94    1.920
# 1   use.for(X, g)         1000   60.97    2.056
# 6  use_Rcpp(X, g)         1000  113.96    3.844
# 2 use.apply(X, g)         1000  301.14   10.156

My intution (use_Rcpp better than use_arma) did not turn out right. Having said that, I guess that the line if (G(i) != g_id) continue; in my use_Rcpp function slows down everything. I am happy to learn about alternatives to set this up.

I am happy that I have achieved the same task in half the time it takes R to do it, but maybe the several Rcpp is much faster than R-examples have messed with my expectations, and I am wondering if I can speed this up even more. Does anyone have an idea? I also welcome any programming / coding comments in general since I am relatively new to Rcpp and C++.

Answer 1

No, it's not the for loop that you need to beat:

library(data.table)
#it doesn't seem fair to include calls to library in benchmarks
#you need to do that only once in your session after all

use.dt2 <- function(mat, ind){
  dt <- as.data.table(mat)
  dt[, cvar := ind]
  dt2 <- dt[,lapply(.SD, sum), by=cvar]
  as.matrix(dt2[,cvar:=NULL])
}

all.equal(use.dt(X,g), use.dt2(X,g))
#TRUE

benchmark(use.for(X,g), use.apply(X,g), use.dt(X,g), use.dt2(X,g),
          columns = c("test", "replications", "elapsed", "relative"), 
          order = "relative", replications = 50)

#             test replications elapsed relative
#4   use.dt2(X, g)           50    3.12    1.000
#1   use.for(X, g)           50    4.67    1.497
#3    use.dt(X, g)           50    7.53    2.413
#2 use.apply(X, g)           50   17.46    5.596

Answer 2

Here's my critiques with in-line comments for your Rcpp solution.

#include <Rcpp.h>
using namespace Rcpp;

// [[Rcpp::export]]
NumericMatrix use_Rcpp(NumericMatrix X, IntegerVector G){

  // Rcpp has a sort_unique() function, which combines the
  // sort and unique steps into one, and is often faster than
  // performing the operations separately. Try `sort_unique(G)`
  IntegerVector gr = unique(G);
  std::sort(gr.begin(), gr.end());
  int gr_n = gr.size();
  int nrow = X.nrow(), ncol = X.ncol();

  // This constructor zero-initializes memory (kind of like
  // making a copy). You should use:
  // 
  //     NumericMatrix out = no_init(gr_n, ncol)
  //
  // to ensure the memory is allocated, but not zeroed.
  // 
  // EDIT: We don't have no_init for matrices right now, but you can hack
  // around that with:
  // 
  //     NumericMatrix out(Rf_allocMatrix(REALSXP, gr_n, ncol));
  NumericMatrix out(gr_n, ncol);

  for(int g=0; g<gr_n; g++){

     // subsetting with operator[] is cheaper, so use gr[g] when
     // you can be sure bounds checks are not necessary
     int g_id = gr(g);

      for (int j = 0; j < ncol; j++) {
      double total = 0;
        for (int i = 0; i < nrow; i++) {

          // similarily here
          if (G(i) != g_id) continue;    // not sure how else to do this
          total += X(i, j);
        }
        // IIUC, you are filling the matrice row-wise. This is slower as
        // R matrices are stored in column-major format, and so filling
        // matrices column-wise will be faster.
        out(g,j) = total;
      }
  }
      return out;
}

Answer 3

Maybe you're looking for (the oddly named) rowsum

library(microbenchmark)
use.rowsum = rowsum

and

> all.equal(use.for(X, g), use.rowsum(X, g), check.attributes=FALSE)
[1] TRUE
> microbenchmark(use.for(X, g), use.rowsum(X, g), times=5)
Unit: milliseconds
             expr       min        lq    median        uq       max neval
    use.for(X, g) 126.92876 127.19027 127.51403 127.64082 128.06579     5
 use.rowsum(X, g)  10.56727  10.93942  11.01106  11.38697  11.38918     5