Questions on array operations in R

I frequently come up with a problem, where I want to modify, say, a column in an array based on information on another array, but I don't know how to do these sort of problems, efficiently avoiding for loops.

To give an example, here is a reproducible R code for a maximization and mass transfer problem using 3D arrays. The code works fine, but the nested loop at the bottom is too slow for my purposes.

The question is, whether there is a way to improve its speed, e.g., by using mapply. I haven't been able to do so this far.


# Initialize parameters


N = 11
I = 3
P = 10


# Initialize some random-ish inputs


prob = runif(P)
prob = prob/sum(prob)

alpha = array(0, c(N,I))

for (i in 1:I) {
  alpha[, i] = (1-seq(0, 1, 1/(N-1)))/i
}

f0 = array(runif(N*I), c(N,I))
f0 = f0/(rep(1,N) %*% t(colSums(f0)))
obj = array(runif(N*N*I*P), c(N,N,I,P))


# Initialize the arrays needed in the solution


maxmat = array(0, c(N, I,P))
inds = array(0, c(I*N, 3, P))
p0 = array(0, c(N,I))
arr = array(0, c(N,I))


# Solve a maximization problem 


maxmat = apply(obj, c(1,3,4), function(x) which.max(x))


# How do I improve the performance of the below code?

for (i in 1:I) {
  for (p in 1:P) {

    p_tmp = 0 
    p_tmp = prob[p]*rowsum(f0[,i], maxmat[, i,p])

    p0[as.numeric(attributes(p_tmp)$dimnames[[1]]),i] = p0[as.numeric(attributes(p_tmp)$dimnames[[1]]),i]+ p_tmp[,1] 
    arr[,i] = arr[,i] + prob[p]*alpha[maxmat[,i, p],i]*f0[,i]
  }
}

Solution

This loop can be fully vectorized using outer and rowSums (with the dims argument) and adjusting the maxmat indices to correspond with i. The vectorized solution is ~18 times faster.

fp <- outer(f0, prob)
mm <- maxmat + rep(rep(seq(0, (I - 1L)*N, N), each = N), P)
p02 <- matrix(rowsum(c(fp, numeric(N*I)), c(mm, 1:(I*N))), N, I)
arr2 <- rowSums(fp*alpha[mm], dims = 2L)

Check that the results are the same:

all.equal(list(p0, arr), list(p02, arr2))
#> [1] TRUE

Benchmark with N = 110, I = 30, and P = 100:

microbenchmark::microbenchmark(
  vectorized = {
    fp <- outer(f0, prob)
    mm <- maxmat + rep(rep(seq(0, (I - 1L)*N, N), each = N), P)
    p02 <- array(rowsum(c(fp, numeric(N*I)), c(mm, 1:(I*N))), c(N, I))
    arr2 <- rowSums(fp*alpha[mm], dims = 2L)
  },
  loops = {
    for (i in 1:I) {
      for (p in 1:P) {
        p_tmp = prob[p]*rowsum(f0[,i], maxmat[, i,p])
        p0[as.numeric(attributes(p_tmp)$dimnames[[1]]),i] = p0[as.numeric(attributes(p_tmp)$dimnames[[1]]),i]+ p_tmp[,1] 
        arr[,i] = arr[,i] + prob[p]*alpha[maxmat[,i, p],i]*f0[,i]
      }
    }
  }
)
#> Unit: milliseconds
#>        expr      min        lq      mean   median        uq      max neval
#>  vectorized  12.7971  17.76585  18.83623  19.4157  19.87045  33.4201   100
#>       loops 213.0616 333.43520 348.48067 366.3358 382.36215 422.4743   100

A little explanation about c(fp, numeric(N*I)) and c(mm, 1:(I*N)): appending N*I zeros (whose groups are 1:(I*N)) to the end of fp ensures that rowsum returns a vector of length N*I so we don't have to initialize p0 or mess with dimnames.