Search code examples
rfor-loopmultidimensional-arrayapplynested-loops

Converting multiple for loops to single or nested apply statements R


I'm looking to decrease excessive processing time from running multiple/nested for loops. There seems to be several questions on this topic but not many R answers so here goes another (hopefully more simplistic) example.

Essentially, we can view this example as creating predictions from a logistic regression output. I'd like to convert the multiple/nested loops below into an apply statement.

Example Using Loops:

set.seed(1111)

#Data
b0 <- matrix(c(rnorm(10,-1, 0.1), rnorm(10,-2, 0.1)), c(10,2))
b1 <- matrix(c(rnorm(10,-.4, 0.1), rnorm(10,-.1, 0.1)), c(10,2))
b2 <- matrix(c(rnorm(10, 0, 0.1), rnorm(10,-.1, 0.1)), c(10,2))
x <- rnorm(5, 0.5, 0.01)
x2 <- rnorm(5, 0.25, 0.01)

#Predictions Array
pred.ray = array(NA, dim = c(10, 5, 2))

# Can do for loops but takes waaaay too long on full data set ~ 20hrs
for(i in 1:5){
  for (j in 1:2){
    pred.ray[,i,j] = plogis(b0[,j] + b1[,j]*x[i] + b2[,j]*x2[i])
  }}

pred.ray

Solution

  • len_x <- 5
    ncol_b <- 2
    val <- len_x * ncol_b
    
    result <- mapply(function(i, j) plogis(b0 + b1 * i + b2 * j), x, x2)
    inds <- lapply(0:1, function(i) sort(rep(seq(i, val+i-1, ncol_b) * val, val) + 
    rep(1:val, each=len_x)))
    ans <- array(c(sapply(inds, function(i) result[i])), dim=c(10,5,2))
    identical(pred.ray, ans)
    # TRUE