How can I speed up my Rcpp code, which only carries out simple operations?

user_14815374 0
Viewed 94

I'm trying to write a function that takes in a matrix and computes a value for every pair of columns. The matrix always has 2000 rows, but can potentially have a very large number of columns (up to 100,000 or so). The R code I started with is as follows:

x_dist <- data.frame(array(0,dim=c(ncol(x),ncol(x))))
cs <- colSums(x)
for (i in 1:ncol(x)) {
  p_i <- x[,i]
  for (j in 1:ncol(x)) {
    p_j <- x[,j]
    s <- p_i+p_j
    fac <- cs[i]/(cs[i]+cs[j])
    N1 <- fac*s
    N2 <- (1-fac)*s
    d1 <- (p_i+1)/(N1+1)
    d2 <- (p_j+1)/(N2+1)
    x_dist[i,j] <- sum(N1+N2-N1*d1-N2*d2+p_i*log(d1)+p_j*log(d2))
  }
}

This function is quite slow. When there are only 400 columns in the matrix x, it takes about 32 seconds, and obviously grows quadratically in the number of columns.

Since I've heard Rcpp is good for speeding up for loops and matrix operations, I decided to give that a try. I am completely new to it, but ended up putting together the following function:

#include <Rcpp.h>
using namespace Rcpp;
// [[Rcpp::export]]
NumericMatrix wdist(NumericMatrix x) {
    int nrow = x.nrow(),ncol=x.ncol();
    NumericMatrix m = no_init_matrix(ncol,ncol);
    NumericVector v1 = no_init_vector(nrow);
    NumericVector v2 = no_init_vector(nrow);
    NumericVector s = no_init_vector(nrow);
    NumericVector N1 = no_init_vector(nrow);
    NumericVector N2 = no_init_vector(nrow);
    NumericVector d1 = no_init_vector(nrow);
    NumericVector d2 = no_init_vector(nrow);
    for(int i=0; i<ncol; ++i){
      v1 = x(_,i);
      for(int j=0; j<i; ++j){
        v2 = x(_,j);
        s = v1+v2;
        N1 = sum(v1)*s/(sum(v1)+sum(v2));
        N2 = s-N1;
        d1 = (v1+1)/(N1+1);
        d2 = (v2+1)/(N2+1);
        m(i,j) = sum(N1+N2-N1*d1-N2*d2+v1*log(d1)+v2*log(d2));
      }
    }
    return m;
}

This certainly makes a big difference. Now with 400 columns, this takes about 8 seconds. I am pleased by the improvement, but this is still intractably slow for my current test case of interest, which is 32,000 columns. I feel like I am doing some relatively simple operations, so it's confusing to me why my code is still this slow. I've tried to do some reading on writing efficient Rcpp code, but haven't found anything that helps address my issue. Please let me know if there is anything I'm doing wrong or any improvements I can look into to make my code faster (or even the R code itself, if that can be made faster than the Rcpp code!)

Some example data could be:

set.seed(121220) 
x <- array(rpois(2000*400,3),dim=c(2000,400))
rcpp dataframe r
1 Answers

I refactored your base R code and hope it could speed up somewhat

f <- function(...) {
  p <- x[, t(...)]
  N <- matrix(rowSums(p), ncol = 1) %*% colSums(p) / sum(p)
  d <- (p + 1) / (N + 1)
  sum(N - N * d + p * log(d))
}
x_dist <- diag(0, ncol(x))
x_dist[lower.tri(x_dist)] <- combn(ncol(x), 2, FUN = f)
x_dist <- pmax(x_dist, t(x_dist))

To speed up your Rcpp code, you can try the following nested for loops after initializing your matrix m as a all-zero matrix:

    for(int i=0; i<ncol-1; ++i){
      v1 = x(_,i);
      for(int j=i+1; j<ncol; ++j){
        v2 = x(_,j);
        s = v1+v2;
        N1 = sum(v1)*s/sum(s);
        N2 = s-N1;
        d1 = (v1+1)/(N1+1);
        d2 = (v2+1)/(N2+1);
        val = sum(N1+N2-N1*d1-N2*d2+v1*log(d1)+v2*log(d2));
        m(i,j) = val;
        m(j,i) = val;
      }
    }

which applies the property that the matrix is symmetry and thus reduce computational complexity by half.

user_121587570
Related