padding_arrays.R

# Based on https://floss.social/@matt_hinckley/113847180218473193


padme <- function(x, pre_pad = 0, post_pad = 0, padding_value = 0){

  ndims <- length(dim(x))
  if (length(pre_pad) == 1L) {
    pre_pad <- rep(pre_pad, ndims)
  }
  if (length(post_pad) == 1L) {
    post_pad <- rep(post_pad, ndims)
  }

  if (length(pre_pad) != ndims) {
    stop("pre_pad should be of length 1 or length ndims")
  }
  if (length(post_pad) != ndims) {
    stop("post_pad should be of length 1 or length ndims")
  }

  new_dims <- pre_pad + dim(x) + post_pad

  y <- array(padding_value, dim = new_dims)

  index_list <- mapply(
    function(d, pre_pad_i) {
      seq_len(d) + pre_pad_i
    },
    dim(x),
    pre_pad
  )
  
  do.call(`[<-`, c(list(y), index_list, list(x)))
}

padme_mult <- function(x, mult= c(16,16,1)) {
   orig_dim <- dim(x)
   new_dim <- ceiling(orig_dim/mult) * mult
   post_pad <- new_dim - orig_dim
   padme(x, post_pad = post_pad)
}