R/arr_3d.R

In rtrend: Trend Estimating Tools

#' apply function for 3d array
#'
#' NA values will be removed automatically
#'
#' @param array A 3d array
#' @param dim giving the subscripts to split up data by.
#' @param FUN function, should only be row applied function, e.g. `matrixStats::rowMeans2`,
#' `matrixStats::rowMins`, `matrixStats::rowRanges`.
#' Because 3d array will be convert to matrix first, with the aggregated dim in
#' the last dimension.
#' @param by
#' * If not provided (`NULL`), the aggregated dim will be disappear.
#'   For example, daily precipitation `[nrow, ncol, 31-days]` aggregate into
#'   monthly `[nrow, ncol]`.
#' * If provided, `by` should be equal to the aggregated `dim`.
#'   For example, daily precipitation `[nrow, ncol, 365-days]` aggregate into
#'   monthly `[nrow, ncol, 12-months]`. In that situation, `by` should be equal
#'   to 365, and be `format(date, '%Y%m')`.
#'
#' @inheritParams set_dim
#' @inheritParams apply_row
#'
#' @seealso [apply_row] [matrixStats::rowRanges]
#'
#' @keywords internal
#' @examples
#' set.seed(1)
#' size <- c(10, 8, 31)
#' arr <- array(rnorm(10 * 8 * 31), dim = size)
#'
#' by <- c(rep(1, 10), rep(2, 21))
#' r2 <- apply_3d(arr, 3, by = by, FUN = rowMeans)
#'
#' \dontrun{
#' arr_yearly <- apply_3d(arr, by = year(dates), scale = days_in_month(dates))
#' }
#' @importFrom matrixStats rowMeans2 rowMins rowMaxs
#' @export
apply_3d <- function(array, dim = 3, FUN = rowMeans2, by = NULL, scale = 1, na.rm = TRUE, ...) {
  # TODO: add by at here
  dims <- dim(array)
  ndim <- length(dims) # dimensions

  I_dims <- setdiff(1:ndim, dim) # dimensions order
  dims_head <- dims[I_dims] # header dimensions

  # move grouped dim to the last
  if (dim != ndim) {
    array %<>% aperm(c(I_dims, dim))
  }

  mat <- array_3dTo2d(array)

  if (is.null(by)) {
    ans <- FUN(mat, ..., na.rm = na.rm)
    dim_new <- dims_head
  } else {
    dim_new <- c(dims_head, length(unique(by)))
    ans <- apply_row(mat, by, FUN, scale = scale)
  }
  dim(ans) <- dim_new
  ans
}

#' apply_col
#'
#' * `apply_col`: aggregate by col, return a `[ngrp, ncol]` matrix
#' * `apply_row`: aggregate by row, return a `[nrow, ngrp]` matrix
#'
#' For example, setting the dimension of `mat` is `[ngrid, ntime]`,
#' if you want to aggregate by time, `apply_row` should be used here;
#' if you want to aggregate by region (grids), `apply_col` should be used.
#'
#' @param mat matrix, `[nrow, ncol]`
#' @param by integer vector, with the dim of `[ntime]`
#' @param scale in the same length of `by`, or a const value,
#' `value_returned` = `FUN(x)*scale`. This parameter is designed for converting
#' monthly to yearly, meanwhile multiply days in month.
#' Currently, same group should have the same scale factor. Otherwise, only the
#' first is used.
#' @note This function also suits for big.matrix object.
#'
#' @examples
#' mat <- matrix(rnorm(4 * 6), 4, 6)
#' mat_bycol <- apply_col(mat, c(1, 1, 2, 2), colMeans)
#' mat_byrow <- apply_row(mat, c(1, 1, 2, 2, 3, 3), rowMeans)
#'
#' @importFrom matrixStats colMeans2 rowMeans2 colMins colMaxs rowMins rowMaxs
#' @keywords internal
#' @export
apply_col <- function(mat, by, FUN = colMeans2, scale = 1, ...) {
  if (length(by) != nrow(mat)) {
    stop("Length of by is not equal to nrow of mat")
  }
  if (length(scale) == 1) scale <- rep(scale, length(by))
  grps <- unique(by) %>% sort()

  ans <- lapply(grps, function(grp) {
    I <- which(by == grp)
    factor <- scale[I][1]
    FUN(mat[I, , drop = FALSE] * factor, na.rm = TRUE, ...)
  }) %>% do.call(rbind, .)

  if (!is.matrix(ans)) ans <- as.matrix(ans)
  ans %>%
    set_rownames(grps) %>%
    set_colnames(colnames(mat))
}


#' @rdname apply_col
#' @export
apply_row <- function(mat, by, FUN = rowMeans2, scale = 1, ...) {
  if (length(by) != ncol(mat)) {
    stop("Length of by is not equal to ncol of mat")
  }
  if (length(scale) == 1) scale <- rep(scale, length(by))
  grps <- unique(by) %>% sort()

  ans <- lapply(grps, function(grp) {
    I <- which(by == grp)
    factor <- scale[I][1]
    FUN(mat[, I, drop = FALSE] * factor, na.rm = TRUE, ...)
  }) %>% do.call(cbind, .)

  if (!is.matrix(ans)) ans <- as.matrix(ans)
  ans %>%
    set_colnames(grps) %>%
    set_rownames(rownames(mat))
}

#' array_3dTo2d
#'
#' @param array array with the dimension of `[nlon, nlat, ntime]`
#' @param I_grid subindex of `[nrow, ncol]`
#'
#' @return
#' `[nlat*nlon, ntime]`
#'
#' @keywords internal
#' @export
array_3dTo2d <- function(array, I_grid = NULL) {
  dim <- dim(array)
  names_last <- dimnames(array) %>% last()

  if (length(dim) >= 3) {
    dim(array) <- c(prod(dim[1:2]), dim[3])
  }
  if (!is.null(I_grid)) {
    array <- array[I_grid, ]
  }
  if (!is.null(names_last)) array %<>% set_dimnames(list(NULL, names_last))
  return(array)
}

#' @param dim `[nrow, ncol]`
#'
#' @keywords internal
#' @rdname array_3dTo2d
#' @export
array_2dTo3d <- function(array, I_grid = NULL, dim) {
  ntime <- dim(array) %>% last()
  dim <- c(dim, ntime)

  temp <- array(NA * array[1], dim = dim) %>%
    array_3dTo2d()
  if (is.null(I_grid)) {
    temp <- array
  } else {
    temp[I_grid, ] <- array
  }
  ans <- set_dim(temp, dim)
  names_last <- dimnames(array) %>% last()
  if (!is.null(names_last)) ans %<>% set_dimnames(list(NULL, NULL, names_last))
  ans
}