Using a list-matrix to supply indices to N-dimensional array

Question

I am working on a partial differential equations project where some N-dimensional objects are required. I got stuck in padding one N-dim object with a copy of each of its dimensions. Here is the function:

padreplicate <- function(a, padSize) {
    # A: is a N-dimensional array
    # padSize: is a vector that defines the padding
    # Pad an array by replicating values.
    numDims <- length(padSize)
    idx <- vector("list", numDims)
    for (k in 1:numDims) {
        M <-  dim(a)[k]       # 32
        onesVector <-  matrix(1, 1, padSize[k])
        idx[[k]] <- c(onesVector, 1:M, M * onesVector)
    }
    do.call( `[`, c(list(a), idx)) 
}

The first argument is the N-dim object; it could be a matrix, a 3D array or higher dimensional array.

An example for a 2D object or matrix would be:

# pad a matrix 4x3 with c(1,1)
set.seed(123456)
mx = matrix(sample.int(9, size = 9*100, replace = TRUE), nrow = 4, ncol = 3)
mx
  #       [,1] [,2] [,3]
  # [1,]    8    4    9
  # [2,]    7    2    2
  # [3,]    4    5    8
  # [4,]    4    1    6
padreplicate(mx, c(1,1))

The padded matrix looks like this:

      [,1] [,2] [,3] [,4] [,5]
[1,]    8    8    4    9    9
[2,]    8    8    4    9    9
[3,]    7    7    2    2    2
[4,]    4    4    5    8    8
[5,]    4    4    1    6    6
[6,]    4    4    1    6    6

For a 3-D array the input array and the padded array.

ar = array(sample.int(9, size = 9*100, replace = TRUE), dim = c(3, 3, 1))
ar
padreplicate(ar, c(1,1,1))

# input 3x3x1 array
, , 1

     [,1] [,2] [,3]
[1,]    3    4    6
[2,]    7    2    9
[3,]    3    4    8

   # padded 5x5x3 array
, , 1

     [,1] [,2] [,3] [,4] [,5]
[1,]    3    3    4    6    6
[2,]    3    3    4    6    6
[3,]    7    7    2    9    9
[4,]    3    3    4    8    8
[5,]    3    3    4    8    8

, , 2

     [,1] [,2] [,3] [,4] [,5]
[1,]    3    3    4    6    6
[2,]    3    3    4    6    6
[3,]    7    7    2    9    9
[4,]    3    3    4    8    8
[5,]    3    3    4    8    8

, , 3

     [,1] [,2] [,3] [,4] [,5]
[1,]    3    3    4    6    6
[2,]    3    3    4    6    6
[3,]    7    7    2    9    9
[4,]    3    3    4    8    8
[5,]    3    3    4    8    8

These are all correct results from my unit tests.

My question is this: "is there a better way of doing this N-dim padding operation in R or a package that address N-dimensional arrays?

Answer 1

Using your 2D example, doing mx[c(1, 1, 2, 3, 4, 4), c(1, 1, 2, 3, 3)] is the obvious solution, both elegant and fast. So your way of generalizing the use of [ via do.call to handle any number of dimensions is what I would consider the optimal approach. The rewrite below won't make the code faster but maybe a little more simple and robust. The main improvement is the computation of the idx using Map rather than a for loop. And I added a number of stopifnot checks to verify all the assumptions you are making regarding your inputs.

padreplicate <- function(a, padSize) {
   stopifnot(is.array(a),
             is.vector(padSize),
             is.numeric(padSize),
             length(padSize) == length(dim(a)),
             padSize >= 0L)
   idx <- Map(function(p, n) c(rep(1L, p), seq(n), rep(n, p)),
              padSize, dim(a))
   do.call( `[`, c(list(a), idx)) 
}