Parallelized calculation of the cross-product, computed with as much memory efficiency as possible.

parallel_crossprod.R

library(BiocParallel)
library(DelayedArray)

#' @importFrom stats start end
.grid_iterate <- function(x, grid = NULL) {
    b <- 0L
    function() {
        if (b == length(grid)) {
            return(NULL)
        }
        b <<- b + 1L

        viewport <- grid[[b]]
        rows <- start(viewport)[1]:end(viewport)[1]
        cols <- start(viewport)[2]:end(viewport)[2]
        return(x[rows,cols,drop=FALSE])
    }
}

.fun_crossprod <- function(i, ...) {
    crossprod(as.matrix(i))
}

#' @importFrom DelayedArray rowGrid 
#' @importFrom BiocParallel bpiterate SerialParam
.delayed_crossprod <- function(X, BPPARAM=SerialParam())
# DelayedMatrix crossprod by row-based blocks.
{
    nrows.per.core <- as.integer(ceiling(nrow(X)/bpnworkers(BPPARAM)))
    G <- RegularArrayGrid(dim(X), c(nrows.per.core, ncol(X)))
    bpiterate(.grid_iterate(X, G), FUN=.fun_crossprod,
        REDUCE = function(x, y) { x + y },
        BPPARAM = BPPARAM)
}

.fun_mult <- function(X, Y, ...) {
    as.matrix(X %*% as.matrix(Y))
}

#' @importFrom DelayedArray RegularArrayGrid
#' @importFrom BiocParallel bpiterate SerialParam
.delayed_mult <- function(X, Y, BPPARAM=SerialParam())
# DelayedMatrix multiplication by column.
{
    ncols.per.core <- as.integer(ceiling(ncol(Y)/bpnworkers(BPPARAM)))
    G <- RegularArrayGrid(dim(Y), c(nrow(Y), ncols.per.core))
    out <- bpiterate(.grid_iterate(Y, G), FUN=.fun_mult, X=X, BPPARAM = BPPARAM)
    do.call(cbind, out)
}