RcppParallel LOOCV

library(Rcpp)
library(RcppArmadillo)
library(RcppParallel)

sourceCpp(code = '
// [[Rcpp::plugins(cpp11)]]
// [[Rcpp::depends(RcppArmadillo, RcppParallel)]]
#include <RcppArmadillo.h>
#include <RcppParallel.h>
using namespace RcppParallel;
using namespace arma;
using namespace Rcpp;

double fastMSE_RcppArma(vec y, mat X, vec yTest, mat XTest)
{
  vec coef = solve(join_rows(ones<vec>(y.n_elem), X), y);
  return as_scalar(pow(yTest - join_rows(ones<vec>(yTest.n_elem), XTest)*coef, 2));
}

struct MSE_Compute: public Worker {
  mat& X;
  vec& y;
  uvec& index;
  vec& mse;
  MSE_Compute(mat& X, vec& y, uvec& index, vec& mse):
    X(X), y(y), index(index), mse(mse) {}
  void operator()(std::size_t begin, std::size_t end)
  {
    for (uword row_index = begin; row_index < end; row_index++)
    {
      mse(row_index) = fastMSE_RcppArma(y.elem(find(index != row_index)),
        X.rows(find(index != row_index)), y.elem(find(index == row_index)),
        X.rows(find(index == row_index)));
    }
  }
};

// [[Rcpp::export]]
NumericVector fastLOOCV(NumericVector yr, NumericMatrix Xr)
{
  int n = Xr.nrow(), p = Xr.ncol();
  mat X(Xr.begin(), n, p, false);
  if (yr.size() != n)
    Rcpp::stop("The size of y must be equal to the number of rows of X." );
  vec y(yr.begin(), n, false),
  mse = zeros<vec>(n);
  uvec index = linspace<uvec>(0, y.n_elem - 1, y.n_elem);
  MSE_Compute mseResults(X, y, index, mse);
  parallelFor(0, n, mseResults);
  return wrap(mse);
}

// [[Rcpp::export]]
NumericVector fastLOOCV_noRP(NumericVector yr, NumericMatrix Xr)
{
  int n = Xr.nrow(), p = Xr.ncol();
  mat X(Xr.begin(), n, p, false);
  if (yr.size() != n)
    Rcpp::stop("The size of y must be equal to the number of rows of X." );
  vec y(yr.begin(), n, false),
  mse = zeros<vec>(n);
  uvec index = linspace<uvec>(0, y.n_elem - 1, y.n_elem);
  for (uword row_index = 0; row_index < y.n_elem; row_index++)
  {
    mse(row_index) = fastMSE_RcppArma(y.elem(find(index != row_index)),
        X.rows(find(index != row_index)), y.elem(find(index == row_index)),
        X.rows(find(index == row_index)));
  }
  return wrap(mse);
}')

LOOCV_R <- function(y, X){
  mses <- vector('numeric', length(y))
  for (i in seq_along(y))
  {
    lmCoefs <- coef(lm(y[-i] ~ X[-i, ]))
    mses[i] <- mean((y[i] - c(1, X[i, ]) %*% lmCoefs)**2)
  }
  return(mses)
}

set.seed(10)
N <- 2000
p <- N/5
X <- matrix(rnorm(N*p), N)
lm_coef <- replicate(p, rnorm(1, rnorm(1, 2, 3), rgamma(1,10,2)))
addingCol = matrix(sample(p, p, TRUE), p/2, 2)
addingCol = addingCol[addingCol[ ,1] != addingCol[ ,2], ]
for (i in 1:ncol(addingCol))
  X[ , addingCol[i,1]] = X[ , addingCol[i,1]] + X[ , addingCol[i,2]]
y <- 9 + rowSums(sweep(X, 2, lm_coef, '*')) + rnorm(N, 0, 10)

st <- proc.time()
mses_R <- LOOCV_R(y, X)
proc.time() - st
#    user  system elapsed
#  372.58   11.23  388.93

st <- proc.time()
mses_cpp1 <- fastLOOCV_noRP(y, X)
proc.time() - st
#    user  system elapsed
#  180.85   18.54  101.63
all.equal(mses_R, as.vector(mses_cpp1)) # TRUE

st <- proc.time()
mses_cpp2 <- fastLOOCV(y, X)
proc.time() - st
#    user  system elapsed
#  351.96  179.52    73.7
all.equal(mses_R, as.vector(mses_cpp2)) # TRUE

# environment: RRO-3.3.1 on my machine with [email protected]