Merge pull request #22 from chrisdjscott/rcpp

Rcpp versions of some functions

Author: doddsk

GBS-Chip-Gmatrix.R

@@ -10,6 +10,38 @@ if (!exists("cex.pointsize"))    cex.pointsize    <- 1
 if (!exists("functions.only"))   functions.only   <- FALSE
 if (!exists("outlevel"))         outlevel         <- 9
 
+# function to locate Rcpp file (assume it is in the same directory as this file and this file was 'sourced')
+pathToCppFile = function() {
+    cpp.name <- "GBS-Rcpp-functions.cpp"
+    this.file <- NULL
+    for (i in -(1:sys.nframe())) {
+        if (identical(sys.function(i), base::source)) this.file <- (normalizePath(sys.frame(i)$ofile))
+    }
+    if (!is.null(this.file)) {
+        source.dir <- dirname(this.file)
+        return(file.path(source.dir, cpp.name))
+    }
+    else {
+        # assume it is in the current working directory
+        return(cpp.name)
+    }
+}
+
+# load C++ functions
+# compiling can take ~10 seconds, so we cache the file (under ~/R/RcppCache)
+# it will only be recompiled when the C++ file changes or is moved
+have_rcpp <- FALSE
+if (require(Rcpp)) {
+    # RcppArmadillo is also required, but we don't need to load it here
+    if (is.element("RcppArmadillo", installed.packages()[,"Package"])) {
+        have_rcpp <- TRUE
+        cpp.path <- pathToCppFile()
+        cat("Loading C++ functions:", cpp.path, "\n")
+        sourceCpp(file = cpp.path, showOutput = TRUE,
+                  cacheDir = file.path(path.expand("~"), "R", "RcppCache"))
+    }
+}
+
 readGBS <- function(genofilefn = genofile) {
  if (gform == "chip") readChip(genofilefn)
  if (gform == "ANGSDcounts") readANGSD(genofilefn)
@@ -210,7 +242,12 @@ GBSsummary <- function() {
  havedepth <- exists("depth")  # if depth present, assume it and genon are correct & shouldn't be recalculated (as alleles may be the wrong one)
  if(gform != "chip") {
   if (!havedepth) depth <<- alleles[, seq(1, 2 * nsnps - 1, 2)] + alleles[, seq(2, 2 * nsnps, 2)]
-  sampdepth.max <<- apply(depth, 1, max)
+  if (have_rcpp) {
+   sampdepth.max <<- rcpp_rowMaximums(depth)
+  }
+  else {
+   sampdepth.max <<- apply(depth, 1, max)
+  }
   sampdepth <<- rowMeans(depth)
   u0 <- which(sampdepth.max == 0)
   u1 <- setdiff(which(sampdepth.max == 1 | sampdepth < sampdepth.thresh), u0)
@@ -295,7 +332,15 @@ cat("Analysing", nind, "individuals and", nsnps, "SNPs\n")
  l10LRT <<- -log10(exp(1)) * pchisq(LRT, 1, lower.tail = FALSE, log.p = TRUE)
 
  sampdepth <<- rowMeans(depth)  # recalc after removing SNPs and samples
- if(outlevel > 4) sampdepth.med <<- apply(depth, 1, median)
+ #if(outlevel > 4) sampdepth.med <<- apply(depth, 1, median)
+ if(outlevel > 4) {
+   if (have_rcpp) {
+     sampdepth.med <<- rcpp_rowMedians(depth)
+   }
+   else {
+     sampdepth.med <<- apply(depth, 1, median)
+   }
+ }
  depth0 <- rowSums(depth == 0)
  snpdepth <<- colMeans(depth)
  missrate <- sum(depth == 0)/nrow(depth)/ncol(depth)
@@ -366,19 +411,47 @@ if(!functions.only) {
 
 
 ################## functions
- depth2K <- function(depthvals)  1/2^depthvals   # convert depth to K value assuming binomial 
+ r_depth2K <- function(depthvals)  1/2^depthvals   # convert depth to K value assuming binomial 
+ # select R or Rcpp version of depth2K depending on whether Rcpp is installed
+ if (have_rcpp) {
+     depth2K <- function(depthvals) {
+         # Rcpp version only works with matrix as input, so fallback to R version otherwise
+         if (is.matrix(depthvals)) {
+             result <- rcpp_depth2K(depthvals)
+         } else {
+             result <- r_depth2K(depthvals)
+         }
+         return(result)
+     }
+ } else {
+     depth2K <- r_depth2K
+ }
 
  depth2Kbb <- function(depthvals, alph=Inf) {
   # convert depth to K value assuming beta-binomial with parameters alpha=beta=alph. Inf gives binomial
   if (alph==Inf) 1/2^depthvals else beta(alph,depthvals+alph)/beta(alph,alph)
   }
 
 # convert depth to K value modp model. prob of seeing same allele as last time is modp (usually >= 0.5)
- depth2Kmodp <- function(depthvals, modp=0.5 ) {
+ r_depth2Kmodp <- function(depthvals, modp=0.5 ) {
   Kvals <- 0.5*modp^(depthvals-1)
   Kvals[which(depthvals==0)] <- 1
   Kvals
   }
+ # select R or Rcpp version depending on whether Rcpp is installed
+ if (have_rcpp) {
+    depth2Kmodp <- function(depthvals, modp=0.5) {
+        # Rcpp version only works with matrix as input, so fallback to R version otherwise
+        if (is.matrix(depthvals)) {
+            result <- rcpp_depth2Kmodp(depthvals, modp)
+        } else {
+            result <- r_depth2Kmodp(depthvals, modp)
+        }
+        return(result)
+    }
+ } else {
+     depth2Kmodp <- r_depth2Kmodp
+ }
 
 depth2Kchoose <- function(dmodel="bb",param) {  # function to choose redefine depth2K
  if (!dmodel=="modp") dmodel <- "bb"
@@ -389,7 +462,6 @@ depth2Kchoose <- function(dmodel="bb",param) {  # function to choose redefine de
  depth2K
  }
 
-
 upper.vec <- function(sqMatrix) as.vector(sqMatrix[upper.tri(sqMatrix)])
 #seq2samp <- function(seqIDvec=seqID) read.table(text=seqIDvec,sep="_",stringsAsFactors=FALSE,fill=TRUE)[,1]  # undoc AgR function # might not get number of cols right
 seq2samp <- function(seqIDvec=seqID) sapply(strsplit(seqIDvec,split="_"),"[",1)
@@ -594,7 +666,12 @@ calcG <- function(snpsubset, sfx = "", puse, indsubset, depth.min = 0, depth.max
    hist(upper.vec(cocall)/nsnpsub, breaks = 50, xlab = "Co-call rate (for sample pairs)", main="", col = "grey")
    dev.off()
   lowpairs <- which(cocall/nsnpsub <= cocall.thresh & upper.tri(cocall),arr.ind=TRUE)
-  sampdepth.max <- apply(depthsub, 1, max)
+  if (have_rcpp) {
+    sampdepth.max <- rcpp_rowMaximums(depthsub)
+  }
+  else {
+    sampdepth.max <- apply(depthsub, 1, max)
+  }
   samp.removed <- NULL
   if(cocall.thresh >= 0) {  # remove samples which wont get self-rel
    samp.removed <- which(sampdepth.max < 2)
@@ -655,11 +732,17 @@ calcG <- function(snpsubset, sfx = "", puse, indsubset, depth.min = 0, depth.max
   rm(GGBS4top)
 
   genon01 <- genon0
-  genon01[depth[indsubset, snpsubset] < 2] <- 0
   P0 <- matrix(puse[snpsubset], nrow = nindsub, ncol = nsnpsub, byrow = TRUE)
   P1 <- 1 - P0
-  P0[!usegeno | depth[indsubset, snpsubset] < 2] <- 0
-  P1[!usegeno | depth[indsubset, snpsubset] < 2] <- 0
+  if (have_rcpp) {
+    rcpp_assignP0P1Genon01(P0, P1, genon01, usegeno, depth[indsubset, snpsubset])
+  }
+  else {
+    genon01[depth[indsubset, snpsubset] < 2] <- 0
+    P0[!usegeno | depth[indsubset, snpsubset] < 2] <- 0
+    P1[!usegeno | depth[indsubset, snpsubset] < 2] <- 0
+  }
+
 #  div0 <- 2 * diag(tcrossprod(P0, P1))  # rowSums version faster
   div0 <- 2 * rowSums(P0 * P1)
   Kdepth <- depth2K(depth[indsubset, snpsubset])
@@ -702,10 +785,18 @@ calcG <- function(snpsubset, sfx = "", puse, indsubset, depth.min = 0, depth.max
        #htmlwidgets::saveWidget(temp_p, paste0("Heatmap-G5", sfx, ".html"))
      } else {
        png(paste0("Heatmap-G5", sfx, ".png"), width = 2000, height = 2000, pointsize = cex.pointsize *  18)
+       if (require(parallelDist)) {
+         cat("Using parallelDist function in heatmap\n")
+         distfun <- parDist
+       }
+       else {
+         cat("Using normal dist function in heatmap\n")
+         distfun <- dist
+       }
        if(length(table(pcacolo)) > 1) {
-         hmout <- heatmap(temp, col = rev(heat.colors(50)), ColSideColors=pcacolo, RowSideColors=pcacolo, symm=T, revC=F)
+         hmout <- heatmap(temp, col = rev(heat.colors(50)), ColSideColors=pcacolo, RowSideColors=pcacolo, symm=T, revC=F, distfun=distfun)
        } else {
-         hmout <- heatmap(temp, col = rev(heat.colors(50)), symm=T, revC=F)
+         hmout <- heatmap(temp, col = rev(heat.colors(50)), symm=T, revC=F, distfun=distfun)
        }
        hmdat <- data.frame(rowInd=hmout$rowInd,seqIDInd=indsubset[pcasamps[hmout$rowInd]],seqID=seqID[indsubset[pcasamps[hmout$rowInd]]])
        write.csv(hmdat,paste0("HeatmapOrder", sfx, ".csv"),row.names=FALSE,quote=FALSE)

GBS-Rcpp-functions.cpp

@@ -0,0 +1,115 @@
+// disable all run-time checks in armadillo (e.g. bounds checking, etc.)
+#define ARMA_NO_DEBUG
+
+// we depend on the R package "RcppArmadillo"
+// [[Rcpp::depends(RcppArmadillo)]]
+#include <RcppArmadillo.h>
+
+// we need OpenMP for parallelisation
+// [[Rcpp::plugins(openmp)]]
+
+// function for finding row medians (alternative to apply(depth, 1, median))
+// requires integer type matrix as input, returns list of doubles
+// [[Rcpp::export]]
+std::vector<double> rcpp_rowMedians(const arma::imat &depth) {
+    // number of rows
+    const int nrows = depth.n_rows;
+
+    // vector for storing the result
+    std::vector<double> medians(nrows);
+
+    // loop over the rows
+    #pragma omp parallel for
+    for (int i = 0; i < nrows; i++) {
+        // convert the row to double type, to compute median correctly
+        arma::rowvec row = arma::conv_to<arma::rowvec>::from(depth.row(i));
+
+        // compute the median for this row
+        medians[i] = arma::median(row);
+    }
+
+    return medians;
+}
+
+// function for finding row maximums (alternative to apply(mat, 1, max))
+// requires integer type matrix as input, return list of integers
+// [[Rcpp::export]]
+std::vector<int> rcpp_rowMaximums(const arma::imat &mat) {
+    const int nrows = mat.n_rows;
+
+    // create vector to store the result
+    std::vector<int> maximums(nrows);
+
+    // loop over rows
+    #pragma omp parallel for
+    for (int i = 0; i < nrows; i++) {
+        // find the maximum for this row
+        maximums[i] = mat.row(i).max();
+    }
+
+    return maximums;
+}
+
+// C++ version of depth2K function
+// [[Rcpp::export]]
+Rcpp::NumericMatrix rcpp_depth2K(const Rcpp::NumericMatrix &A) {
+    // create the output matrix (same size as input)
+    Rcpp::NumericMatrix Aout(A.rows(), A.cols());
+
+    // number of elements
+    const long Asize = A.rows() * A.cols();
+
+    // loop over elements in parallel and apply operation
+    #pragma omp parallel for
+    for (long i = 0; i < Asize; i++) {
+        Aout[i] = 1.0 / pow(2.0, A[i]);
+    }
+
+    // return matrix
+    return Aout;
+}
+
+// C++ version of depth2Kmodp function
+// [[Rcpp::export]]
+Rcpp::NumericMatrix rcpp_depth2Kmodp(const Rcpp::NumericMatrix &depthvals, double modp = 0.5) {
+    // create matrix for storing the result
+    Rcpp::NumericMatrix result(depthvals.rows(), depthvals.cols());
+
+    // size of the matrix
+    const long size = depthvals.rows() * depthvals.cols();
+
+    // loop over the elements in parallel
+    #pragma omp parallel for
+    for (long i = 0; i < size; i++) {
+        double value = 0.5 * pow(modp, depthvals[i] - 1.0);
+        result[i] = (value == 0) ? 1.0 : value;
+    }
+    return result;
+}
+
+// function for setting unused values of P0, P1 and genon01 to zero
+// modifies the matrices in-place (i.e. doesn't return anything)
+// [[Rcpp::export]]
+void rcpp_assignP0P1Genon01(Rcpp::NumericMatrix &P0, Rcpp::NumericMatrix &P1, Rcpp::NumericMatrix &genon01,
+        const Rcpp::LogicalMatrix &usegeno, const Rcpp::NumericMatrix &dsub) {
+    // number of elements (assumes all inputs are the same size!)
+    const long size = P0.rows() * P0.cols();
+
+    // loop over elements in parallel
+    #pragma omp parallel for
+    for (long i = 0; i < size; i++) {
+        // set to zero if they match the conditions
+        if (dsub[i] < 2.0) {
+            P0[i] = 0.0;
+            P1[i] = 0.0;
+            genon01[i] = 0.0;
+        }
+        else if (!usegeno[i]) {
+            P0[i] = 0.0;
+            P1[i] = 0.0;
+        }
+    }
+
+    // nothing to return, matrices are modified in-place
+    return;
+}