parallel_boot_lm.cpp

// [[Rcpp::depends(RcppParallel)]]
// [[Rcpp::depends(RcppEigen)]]
#include <Rcpp.h>
#include <RcppEigen.h>
#include <RcppParallel.h>

using namespace std;
using namespace Rcpp;
using namespace RcppParallel;

using Eigen::Map;
using Eigen::MatrixXd;
using Eigen::VectorXd;
using Eigen::MatrixXi;
using Eigen::VectorXi;

template<typename Derived>
inline MatrixXd matIdx(const Eigen::MatrixBase<Derived>& X, const VectorXi& index)
{
  MatrixXd X2(index.size(), X.cols());
  for (unsigned int i = 0; i < index.size(); ++i)
    X2.row(i) = X.row(index(i));
  return X2;
}

struct CVLM : public Worker
{
  // design matrix and outcome
  const MatrixXd& X;
  const VectorXd& y;
  const MatrixXi& index;

  // output
  MatrixXd& betamat;

  // initialize with input and output
  CVLM(const MatrixXd& X, const VectorXd& y, const MatrixXi& index, MatrixXd& betamat)
    : X(X), y(y), index(index), betamat(betamat) {}

  void operator()(std::size_t begin, std::size_t end) {
    for(unsigned int i = begin; i < end; ++i)
    {
      MatrixXd Xi = matIdx(X, index.col(i));
      VectorXd yi = matIdx(y, index.col(i));
      betamat.col(i) = Xi.colPivHouseholderQr().solve(yi);
    }
  }
};

// [[Rcpp::export]]
Eigen::MatrixXd parallelFit(Rcpp::NumericMatrix xr,
                            Rcpp::NumericVector dvr,
                            Rcpp::IntegerMatrix indexr) {
  MatrixXd x = MatrixXd::Map(xr.begin(), xr.nrow(), xr.ncol());
  VectorXd dv = VectorXd::Map(dvr.begin(), dvr.size());
  MatrixXi index = MatrixXi::Map(indexr.begin(), indexr.nrow(), indexr.ncol());

  // allocate the output matrix
  MatrixXd betamat = MatrixXd::Zero(x.cols(), index.cols());

  // pass input and output
  CVLM cvLM(x, dv, index, betamat);

  // parallelFor to do it
  parallelFor(0, index.cols(), cvLM);

  return betamat;
}