Bonus: Optimizing R code for speed

Reading/working time: ~30 min.

# Preparation: Install and load all necessary packages
# install.packages(c("RcppArmadillo", "future.apply"))

library(RcppArmadillo) # for fast LMs
library(future.apply)  # for parallel processing

# plan() initializes parallel processing, 
# for faster code execution
# By default, it uses all available cores
plan(multisession)

Optimizing code for speed can be an art - and you get lost and spend/waste hours by micro-optimizing some milliseconds. But the Pareto principle applies here: with 20% effort, you can have quick and substantial gains.

Code profiling means that the code execution is timed, just like you had a stopwatch. Your goal is to make your code snippet as fast as possible. RStudio has a built-in profiler that (in theory) allows to see which code line takes up the longest time. But in my experience, if the computation of each single line is very short (and the duration mostly comes from the many repetitions), it is very inaccurate (i.e., the time spent is allocated to the wrong lines). Therefore, we’ll resort to the simplest way of timing code: We will measure overall execution time by wrapping our code in a system.time({ ... }) call. Longer code blocks need to be wrapped in curly braces {...}. The function returns multiple timings; the relevant number for us is the “elapsed” time. This is also called the “wall clock” time - the time you actually have to wait until computation finished.

First, naive version

Here is a first version of the power simulation code for a simple LM. Let’s see how long it takes:

t0 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)
result <- data.frame()

for (n in ns) {
  p_values <- c()
  
  for (i in 1:iterations) {
    treatment <- c(rep(0, n/2), rep(1, n/2))
    BDI <- 23 - 3*treatment + rnorm(n, mean=0, sd=sqrt(117))
    df <- data.frame(treatment, BDI)
    res <- lm(BDI ~ treatment, data=df)
    p_values <- c(p_values, summary(res)$coefficients["treatment", "Pr(>|t|)"])
  }
  
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})
t0

       User      System verstrichen 
      9.627       0.218       9.888 

result

    n  power
1 300 0.3348
2 350 0.4088
3 400 0.4914
4 450 0.5358
5 500 0.6146

This first version takes 9.888 seconds. Of course we have sampling error here as well; if you run this code multiple times, you will always get slightly different timings. But, again, we refrain from micro-optimizing in the millisecond range, so a single run is generally good enough. You should only tune your simulation in a way that it takes at least a few seconds; if you are in the millisecond range, the timings are imprecise and you won’t see speed improvements very well.

Rule 1: No growing vectors/data frames

This is one of the most common bottlenecks: You start with an empty vector (or even worse: data frame), and grow it by rbind-ing new rows to it in each iteration.

t1 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  # print(n)   # uncomment to see progress
  
  # CHANGE: Preallocate vector with the final size, initialize with NAs
  p_values <- rep(NA, iterations)
  
  for (i in 1:iterations) {
    treatment <- c(rep(0, n/2), rep(1, n/2))
    BDI <- 23 - 6*treatment + rnorm(n, mean=0, sd=sqrt(117))
    df <- data.frame(treatment, BDI)
    res <- lm(BDI ~ treatment, data=df)
    
    # CHANGE: assign resulting p-value to specific slot in vector
    p_values[i] <- summary(res)$coefficients["treatment", "Pr(>|t|)"]
  }
  
  # Here we stil have a growing data.frame - but as this is only done 5 times, it does not matter.
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations)) 
}

})

# Combine the different timings in a data frame
timings <- rbind(t0[3], t1[3]) |> data.frame()

# compute the absolute and relative difference of consecutive rows:
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3)

timings

  elapsed  diff rel_diff
1   9.888    NA       NA
2  10.579 0.691     0.07

OK, this didn’t really change anything here. But in general (in particular with data frames) this is worth looking at.

Rule 2: Avoid data frames as far as possible

Use matrizes instead of data frames wherever possible; or avoid them at all (as we do in the code below).

t2 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  treatment <- c(rep(0, n/2), rep(1, n/2))
  p_values <- rep(NA, iterations)
  
  for (i in 1:iterations) {
    BDI <- 23 - 3*treatment + rnorm(n, mean=0, sd=sqrt(117))
    
    # CHANGE: We don't need the data frame - just create the two variables
    # in the environment and lm() takes them from there.
    #df <- data.frame(treatment, BDI)
    res <- lm(BDI ~ treatment)
    
    p_values[i] <- summary(res)$coefficients["treatment", "Pr(>|t|)"]
  }
  
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})

timings <- rbind(t0[3], t1[3], t2[3]) |> data.frame()
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3)
timings

  elapsed   diff rel_diff
1   9.888     NA       NA
2   9.623 -0.265   -0.027
3   7.384 -2.239   -0.233

This showed a substantial improvement of around -3.2 seconds; a relative gain of -23.3%.

Rule 3: Avoid unnecessary computations

What do we actually need? In fact only the p-value for our focal predictor. But the lm function does so many more things, for example parsing the formula BDI ~ treatment.

We could strip away all overhead and do only the necessary steps: Fit the linear model, and retrieve the p-values (see https://stackoverflow.com/q/49732933). This needs some deeper knowledge of the functions and some google-fu. When you do this, you should definitely compare your results with the original result from the lm function and verify that they are identical!

t3 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  # construct the design matrix: first column is all-1 (intercept), second column is the treatment factor
  x <- cbind(
    rep(1, n),
    c(rep(0, n/2), rep(1, n/2))
  )
  
  p_values <- rep(NA, iterations)
  
  for (i in 1:iterations) {
    y <- 23 - 3*x[, 2] + rnorm(n, mean=0, sd=sqrt(117))

    # For comparison - do we get the same results? Yes!
    # res0 <- lm(y ~ x[, 2])
    # summary(res0)
    
    # fit the model:
    m <- .lm.fit(x, y)
    
    # compute p-values based on the residuals:
    rss <- sum(m$residuals^2)
    rdf <- length(y) - ncol(x)
    resvar <- rss/rdf
    R <- chol2inv(m$qr)
    se <- sqrt(diag(R) * resvar)
    ps <- 2*pt(abs(m$coef/se),rdf,lower.tail=FALSE)
    
    p_values[i] <- ps[2]
  }
  
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})
timings <- rbind(t0[3], t1[3], t2[3], t3[3]) |> data.frame()
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3)
timings

  elapsed   diff rel_diff
1   9.888     NA       NA
2   9.623 -0.265   -0.027
3   7.384 -2.239   -0.233
4   0.793 -6.591   -0.893

This step led to a massive improvement of around -6.6 seconds; a relative gain of -89.3%.

Rule 4: Use optimized packages

For many statistical models, there are packages optimized for speed, see for example here: https://stackoverflow.com/q/49732933

t4 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  # construct the design matrix: first column is all-1 (intercept), second column is the treatment factor
  x <- cbind(
    rep(1, n),
    c(rep(0, n/2), rep(1, n/2))
  )
  
  p_values <- rep(NA, iterations)
  
  for (i in 1:iterations) {
    y <- 23 - 3*x[, 2] + rnorm(n, mean=0, sd=sqrt(117))

    # For comparison - do we get the same results? Yes!
    # res0 <- lm(y ~ x[, 2])
    # summary(res0)
    
    mdl <- RcppArmadillo::fastLmPure(x, y)
    
    # compute the p-value - but only for the coefficient of interest!
    p_values[i] <- 2*pt(abs(mdl$coefficients[2]/mdl$stderr[2]), mdl$df.residual, lower.tail=FALSE)
  }
  
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})
timings <- rbind(t0[3], t1[3], t2[3], t3[3], t4[3]) |> data.frame()
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3)
timings

  elapsed   diff rel_diff
1   9.888     NA       NA
2  10.579  0.691    0.070
3   7.384 -3.195   -0.302
4   0.793 -6.591   -0.893
5   0.764 -0.029   -0.037

This step only gave a minor -4% relative increase in speed - but as a bonus, it made our code much easier to read and shorter.

Rule 5: Go parallel

By default, R runs single-threaded. That means, a single CPU core works off all lines of code sequentially. When you optimized this single thread performance, the only way to gain more speed (except buying a faster computer) is to distribute the workload to multiple CPU cores that work in parallel. Every modern CPU comes with multiple cores (also called “workers” in the code); typically 4 to 8 on local computers and laptops.

Preparation: Wrap the simulation in a function

The first step does not really change a lot: We put the simulation code into a separate function that returns the quantity of interest (in our case: the focal p-value). Different settings of the simulation parameters, such as the sample size or the effect size, can be defined as parameters of the function.

Every single function call sim() now gives you one simulated p-value - try it out!

We then use the replicate function to run the sim function many times and to store the resulting p-values in a vector. Programming the simulation in such a functional style also has the nice side effect that you do not have to pre-allocate the results vector; this is automatically done by the replicate function.

# Wrap the code for a single simulation into a function. It returns the quantity of interest.
sim <- function(n=100) {
  # the "n" is now taken from the function parameter "n"
  x <- cbind(
    rep(1, n),
    c(rep(0, n/2), rep(1, n/2))
  )
  
  y <- 23 - 3*x[, 2] + rnorm(n, mean=0, sd=sqrt(117))
  mdl <- RcppArmadillo::fastLmPure(x, y)
  p_val <- 2*pt(abs(mdl$coefficients[2]/mdl$stderr[2]), mdl$df.residual, lower.tail=FALSE)

  return(p_val)
}

t5 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  p_values <- replicate(n=iterations, sim(n=n))
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})
timings <- rbind(t0[3], t1[3], t2[3], t3[3], t4[3], t5[3]) |> data.frame()
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3)
timings

  elapsed   diff rel_diff
1   9.888     NA       NA
2  10.579  0.691    0.070
3   7.384 -3.195   -0.302
4   0.793 -6.591   -0.893
5   0.764 -0.029   -0.037
6   0.935  0.171    0.224

While this refactoring actually slightly increased computation time, we need this for the last, final optimization where we reap the benefits.

Run on multiple cores

With the use of the replicate function in the previous step, we prepared everything for an easy switch to multi-core processing. You only need to load the future.apply package, start a multi-core session with the plan command, and replace the replicate function call with future_replicate.

# Show how many cores are available on your machine:
availableCores()

# with plan() you enter the parallel mode. Enter the number of workers (aka. CPU cores)
plan(multisession, workers = 4)

t6 <- system.time({

iterations <- 5000
ns <- seq(300, 500, by=50)

result <- data.frame()

for (n in ns) {
  # future.seed = TRUE is needed to set seeds in all parallel processes. Then the computation is reproducible.
  p_values <- future_replicate(n=iterations, sim(n=n), future.seed = TRUE)
  result <- rbind(result, data.frame(n = n, power = sum(p_values < .005)/iterations))
}

})

timings <- rbind(t0[3], t1[3], t2[3], t3[3], t4[3], t5[3], t6[3]) |> data.frame()
timings$diff <- c(NA, timings[2:nrow(timings), 1] - timings[1:(nrow(timings)-1), 1])
timings$rel_diff <- c(NA, timings[2:nrow(timings), "diff"]/timings[1:(nrow(timings)-1), 1]) |> round(3) |> round(2)
timings

  elapsed   diff rel_diff
1   9.888     NA       NA
2  10.579  0.691     0.07
3   7.384 -3.195    -0.30
4   0.793 -6.591    -0.89
5   0.764 -0.029    -0.04
6   0.935  0.171     0.22
7   0.788 -0.147    -0.16

The speed improvement seems only small - with 4 workers, one might expect that the computations only need 1/4th of the previous time. But parallel processing creates some overhead. For example, 4 separate R sessions need to be created and all packages, code (and sometimes data) need to be loaded in each session. Finally, all results must be collected and aggregated from all separate sessions. This can add up to substantial one-time costs. If your (single-core) computations only take a few seconds or less, parallel processing can even take longer.

The final speed test: Burn your machine 🔥

Let’s see if parallel processing has an advantage when we have longer computations. We now expand the simulation by exploring a broad parameter range (n ranging from 100 to 1000) and increasing the iterations to 20,000 for more stable results. (See also: “Bonus: How many Monte-Carlo iterations are necessary?”)

plan(multisession, workers = 4)

iterations <- 20000
ns <- seq(100, 1000, by=50)
result_single <- result_parallel <- data.frame()

# single core
t_single <- system.time({
  for (n in ns) {
    p_values <- replicate(n=iterations, sim(n=n))
    result_single <- rbind(result_single, data.frame(n = n, power = sum(p_values < .005)/iterations))
  }
})

# multi-core
t_parallel <- system.time({
  for (n in ns) {
    p_values <- future_replicate(n=iterations, sim(n=n), future.seed = TRUE)
    result_parallel <- rbind(result_parallel, data.frame(n = n, power = sum(p_values < .005)/iterations))
  }
})

# compare results
cbind(result_single, power.parallel = result_parallel[, 2])

      n   power power.parallel
1   100 0.07075        0.07395
2   150 0.13085        0.12945
3   200 0.19460        0.19165
4   250 0.26615        0.26295
5   300 0.33020        0.33635
6   350 0.41320        0.40745
7   400 0.48625        0.47710
8   450 0.54520        0.55245
9   500 0.61160        0.61090
10  550 0.67060        0.66615
11  600 0.71825        0.72150
12  650 0.75465        0.76470
13  700 0.80650        0.79815
14  750 0.83965        0.83605
15  800 0.86445        0.86475
16  850 0.89125        0.88980
17  900 0.91100        0.90960
18  950 0.92700        0.93055
19 1000 0.94025        0.94065

rbind(t_single, t_parallel) |> data.frame()

           user.self sys.self elapsed user.child sys.child
t_single      15.613    0.784  16.446          0         0
t_parallel     1.071    0.064   6.622          0         0

With this optimized setup, we are running 380000 simulations in just 6.622 seconds. If you try this with the first code version, it takes 165.732 seconds.

With the final, parallelized version we have a 25x speed gain relative to the first version!

Recap

We covered the most important steps for speeding up your code in R:

1. No growing vectors/data frames. Solution: Pre-allocate the results vector.
2. Avoid data.frames. Solution: Use matrices wherever possible, or switch to data.table for more complex data structures (not covered here).
3. Avoid unnecessary computations and/or switch to optimized packages that do the same computations much faster, such as the package Rfast.
4. Switch to parallel processing. Solution: If you already programmed your simulations with the replicate function, it is very easy with the future.apply package.

Some steps, such as avoiding growing vectors, didn’t really help in our current example, but will help a lot in other scenarios.

There are many blog post showing and comparing strategies to increase R performance, e.g.:

• https://www.r-bloggers.com/2016/01/strategies-to-speedup-r-code/
• https://adv-r.hadley.nz/perf-improve.html
• https://csgillespie.github.io/efficientR/performance.html

But always remember:

“We should forget about small efficiencies, say about 97% of the time: premature optimization is the root of all evil. Yet we should not pass up our opportunities in that critical 3%.”

Donald Knuth (Structured Programming with go to Statements, ACM Journal Computing Surveys, Vol 6, No. 4, Dec. 1974. p. 268)

Session Info

These speed measurements have been performed on a 2021 MacBook Pro with M1 processor.

sessionInfo()

R version 4.2.0 (2022-04-22)
Platform: aarch64-apple-darwin20 (64-bit)
Running under: macOS 13.2.1

Matrix products: default
BLAS:   /Library/Frameworks/R.framework/Versions/4.2-arm64/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/4.2-arm64/Resources/lib/libRlapack.dylib

locale:
[1] de_DE.UTF-8/de_DE.UTF-8/de_DE.UTF-8/C/de_DE.UTF-8/de_DE.UTF-8

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] future.apply_1.10.0      future_1.32.0            RcppArmadillo_0.12.0.1.0
[4] prettycode_1.1.0         colorDF_0.1.7           

loaded via a namespace (and not attached):
 [1] Rcpp_1.0.10       parallelly_1.35.0 knitr_1.42        magrittr_2.0.3   
 [5] rlang_1.1.0       fastmap_1.1.1     globals_0.16.2    tools_4.2.0      
 [9] parallel_4.2.0    xfun_0.37         cli_3.6.1         htmltools_0.5.5  
[13] yaml_2.3.7        digest_0.6.31     lifecycle_1.0.3   crayon_1.5.2     
[17] purrr_1.0.1       htmlwidgets_1.6.2 vctrs_0.6.1       codetools_0.2-19 
[21] evaluate_0.20     rmarkdown_2.20    compiler_4.2.0    jsonlite_1.8.4   
[25] listenv_0.9.0