fast lm f statistics

requirePackages <- c("foreach", "doSNOW", "data.table", "purrr",
                     "pipeR", "parallel", "Rcpp", "RcppArmadillo", "RcppParallel")
if (!all(requirePackages %in% rownames(installed.packages())))
  install.packages(requirePackages[!requirePackages %in% rownames(installed.packages())])

library(foreach)
library(doSNOW)
library(plyr)
library(data.table)
library(purrr)
library(pipeR)
library(Rcpp)
library(RcppArmadillo)

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

// [[Rcpp::export]]
double fast_lm_f_stat(const vec& y, const mat& X)
{
  if (y.n_elem != X.n_rows)
    Rcpp::stop("The size of y must be equal to the number of rows of X." );
  vec beta = solve(X, y);
  vec yhat = X * beta;
  double MSR = sum(square(yhat - mean(y))) / ((double) X.n_cols - 1.0),
         MSE = sum(square(y - yhat)) /
           ((double) y.n_elem - (double) X.n_cols);
  return(MSR / MSE);
}')

gene <- fread('GE_QNbyPopu_qc_CHB.txt') %>>%
  '['(j = TargetID := NULL) %>>% as.matrix
# gene is a matrix with dimensio c(12028, 45)
R <- 999
x <- as.factor(rep(c("AA","AB","BB"), c(19, 15, 11)))
xMat <- model.matrix(~ ., data.frame(x))

lm_f_stat <- function(y, X){
  solve(crossprod(X), crossprod(X, y)) %>>%
    (b ~ X %*% b) %>>%
    (yhat ~ sum((yhat - mean(y))^2) / sum((y - yhat)^2) /
       (ncol(X) - 1) * (length(y) - ncol(X)))
}

# check the result is corerect
summary(lm(gene[1, ]~ xMat))$fstatistic[1] %>>%
  {
    stopifnot(abs(. - lm_f_stat(gene[1, ], xMat)) < 1e-6)
    stopifnot(abs(. - fast_lm_f_stat(gene[1, ], xMat)) < 1e-6)
  }

st <- proc.time()
p.value <- map_dbl(1:nrow(gene), function(i){
  mean(map_dbl(1:R, ~fast_lm_f_stat(sample(gene[i, ]), xMat)) >
         fast_lm_f_stat(gene[i, ], xMat))
})
proc.time() - st
#  user  system elapsed
# 10.05    0.72  277.50

cl <- makeCluster(parallel::detectCores(), type = "SOCK")
registerDoSNOW(cl)

clusterEnv <- clusterEvalQ(cl, {
  library(purrr)
  library(pipeR)
})
clusterExport(cl, c("lm_f_stat", "xMat", "R"))

st <- proc.time()
p.value <- aaply(gene, 1, function(y){
  mean(map_dbl(1:R, ~lm_f_stat(sample(y), xMat)) >
         lm_f_stat(y, xMat))
}, .parallel = TRUE)
proc.time() - st
#   user  system elapsed
# 296.68    0.11  298.29

stopCluster(cl)

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

inline double fast_lm_f_stat(const vec& y, const mat& X)
{
  vec beta = solve(X, y);
  vec yhat = X * beta;
  double MSR = sum(square(yhat - mean(y))) / ((double) X.n_cols - 1.0),
         MSE = sum(square(y - yhat)) /
           ((double) y.n_elem - (double) X.n_cols);
  return(MSR / MSE);
}

struct pvalueWorker: public Worker {
  const double& rep;
  const mat& Y;
  const mat& X;
  vec& pvalues;
  pvalueWorker(const double& rep, const mat& Y, const mat& X, vec& pvalues):
    rep(rep), Y(Y), X(X), pvalues(pvalues) {}
  void operator()(std::size_t begin, std::size_t end)
  {
    for (uword i = begin; i < end; i++)
    {
      vec tmp_y = Y.row(i).t(), tmp_y2 = tmp_y;
      double f0 = fast_lm_f_stat(tmp_y, X);
      vec f_rep(rep);
      for (uword j = 0; j < (uword) rep; j++)
      {
        tmp_y2 = shuffle(tmp_y);
        f_rep(j) = fast_lm_f_stat(tmp_y2, X);
      }
      pvalues(i) = (double) sum(f_rep > f0) / rep;
    }
  }
};

// [[Rcpp::export]]
vec pvalue_f(const double& rep, const mat& Y, const mat& X)
{
  vec pvalues(Y.n_rows);
  pvalueWorker pvalue_worker(rep, Y, X, pvalues);
  parallelFor(0, Y.n_rows, pvalue_worker);
  return(pvalues);
}')

st <- proc.time()
p.value <- pvalue_f(R, gene, xMat)
proc.time() - st
#   user  system elapsed
# 207.96    0.00   26.82