R/par_plyr.R

In rtrend: Trend Estimating Tools

#' chunk
#' 
#' @param x a vector or list
#' @param nchunk the number of chunks to be splitted
#' 
#' @references https://stackoverflow.com/questions/3318333/split-a-vector-into-chunks-in-r
#' @keywords internal
#' @export
chunk <- function(x, nchunk = 6) {
    split(x, cut(seq_along(x), nchunk, labels = FALSE)) %>% set_names(NULL)
}

#' split_data
#' 
#' @inheritParams chunk
#' @param x a vector, list, or matrix (not support 3d array)
#' @param byrow If `TRUE`, split by row, otherwise by column.
#' 
#' @keywords internal
#' @export
split_data <- function(x, nchunk = 6, byrow = TRUE) {
    n <- length(x)
    if (is.matrix(x)) n <- ifelse(byrow, nrow(x), ncol(x))

    lst_inds <- chunk(1:n, nchunk)
    if (is.matrix(x)) {
        FUN <- if (byrow) . %>% x[., ] else . %>% x[, ., drop = FALSE]
        lapply(lst_inds, FUN)
    } else if (is.list(x) || is.vector(x)) {
        lapply(lst_inds, . %>% x[.])
    } else {
        stop("unsupported type!")
    }
}

#' parallel apply and llply
#' 
#' @keywords internal
#' @importFrom parallel parLapply
#' @export
llply_par <- function(X, FUN, ..., byrow = TRUE, .combine = c) {
    nchunk <- length(getOption("cl"))
    if (nchunk <= 1) {
        ans <- FUN(X, ...)
    } else {
        lst <- split_data(X, nchunk, byrow = byrow)
        ans <- parLapply(getOption("cl"), lst, FUN, ...)
        if (!is.null(.combine)) ans %<>% do.call(.combine, .)
    }
    ans
}

#' @keywords internal
#' @rdname llply_par
#' @export
parLapply2 <- llply_par

#' @keywords internal
#' @rdname llply_par
#' @export
apply_par <- function(X, .margins = 1, FUN, ..., .progress = "text") {
    byrow <- .margins == 1
    res <- llply_par(X, function(x) {
        plyr::aaply(x, .margins, FUN, ..., .progress = .progress)
    }, byrow = byrow, .combine = NULL)

    if (is.list(res)) {
        comb <- if (is.matrix(res[[1]])) rbind else c
        do.call(comb, res)
    } else {
        res
    }
}