Write your deep loops in Stan, not R

deep_loops_in_stan.R

library(tidyverse); library(rstan)

# Some fake data

Individuals <- 1000
Sims <- 1000
Months <- 20


initial_values <- data.frame(customer_id = 1:Individuals, initial_value = rnorm(Individuals))
parameter_draws <- matrix(c(rnorm(Sims, 0.5, 0.1), rnorm(Sims, 0.3, .1)), 2, Sims, byrow = T)
common_shocks <- matrix(rnorm(Months*Sims, 0, .1), Months, Sims)


# OK -- so the model is y_{it} = alpha[t] + beta * y_{it-1} + epsilon_{it}
# where a[t] ~ normal(a, 0.1) is a common shock to all series. 
# and epsilon is n(0,1)

# A function to produce sims in R

sims_in_r <- function(Individuals, Sims, Months, 
                      initial_values, parameter_draws, 
                      common_shocks) {
  out <- matrix(NA, Individuals * Months, Sims)
  count <- 0
  for(i in 1:Individuals) {
    for(m in 1:Months) {
      count <- count + 1
      for(s in 1:Sims) {
        if(m==1) {
          out[count,s] <- rnorm(1, parameter_draws[1,s] + 
                                  common_shocks[m,s] + 
                                  parameter_draws[2,s] * initial_values[i,2], 
                                1)
        } else {
          out[count,s] <- rnorm(1, parameter_draws[1,s] + 
                                  common_shocks[m,s] + 
                                  parameter_draws[2,s] * out[count -1,s], 
                                1)
        }
      }
    }
  }
  return(out)
}

# The same function in Stan -- note that it's almost the same
stan_code <- "
functions {
matrix sims_in_stan_rng(int Individuals, int Sims, int Months, matrix initial_values, matrix parameter_draws, matrix common_shocks) {
matrix[Individuals*Months, Sims] out;

  int count = 0;
  for(i in 1:Individuals) {
     for(m in 1:Months) {
         count = count + 1;
         for(s in 1:Sims) {
           if(m==1) {
             out[count,s] = normal_rng(parameter_draws[1,s] + 
                                        common_shocks[m,s] + 
                                        parameter_draws[2,s] * initial_values[i,2], 
                                        1);
           } else {
             out[count,s] = normal_rng(parameter_draws[1,s] + 
                                       common_shocks[m,s] + 
                                       parameter_draws[2,s] * out[count -1,s], 
                                       1);
            }
          }
      }
  }
  
  return(out);

}
}
data{}
model{}
"

fileConn<-file("demo.stan")
writeLines(stan_code, fileConn)
close(fileConn)

# Let's 
expose_stan_functions("demo.stan")


time_r <- system.time({test_r <- sims_in_r(Individuals, Sims, Months, initial_values, parameter_draws, common_shocks)})
time_stan <- system.time({test_stan <- sims_in_stan_rng(Individuals, Sims, Months, as.matrix(initial_values), parameter_draws, common_shocks)})

time_r/time_stan