Vectorizing a complex nested for loop in R (running models on different subsets of a data set, subsetting the data differently for each loop)

Question

This is my first time posting so hopefully I manage to get my question across in a clear and concise manner. I'm not sure this is the right place but someone on StackOverflow suggested I ask here.

I am running a cross-validation experiment which involves training a model on a subset of a data set and then validating it against the rest of the data set. This is done by randomly subsetting the data into "folds" and then running the model for each fold. This is then replicated in order to produce standard errors. This is the code I have written in R:

reps = 10
folds = 10
rep.Res<-list()

for(i in 1:reps){
  #i=1
  phen2<-phen
  phen2$fold<-sample(1:folds,nrow(phen),replace=T)

  fold.Res<-list()

  for(j in 1:folds){
    #j=1

    print(paste(i," ",j))

    phen3<-phen2
    phen3[which(phen3$fold==j),2]<-NA

    ETA<-list(list(X = geno_centred, model = 'BayesA'))

    model <- BGLR(y = phen3[,2], ETA = ETA, nIter = 20000, 
                  burnIn = 1500,  thin=10) 

    model.output<-data.frame(ID = phen3[which(phen2$fold == j), 1],
                             PBV = as.numeric(model$yHat[which(phen2$fold == j)]),
                             phen = phen2[which(phen2$fold == j), 2], rep = i, fold = j)


fold.Res[[j]]<-model.output


  }

  rep.Res[[i]]<-do.call(rbind,fold.Res)


}

"phen" is a data frame with a a column of IDs and a column of values. For each replicate, each sample is assigned to a fold randomly. Then, for each fold individuals in the fold are given an NA value, i.e. excluded from the model.

The model is run using the Bayesian Generalised Linear Regression (BGLR) package. Unfortunately, I can't really give a reproducible example. X is a scaled centred genotype matrix of -1's, 0's and 1's (not critical for my problem here). The output of the model is then written out for each fold, and then finally for each replicate.

My question is, is it possible to vectorize these loops? Using 10 folds and 10 reps means that the model runs 100 time, at 20 000 MCMC iterations each time. This causes the system to slow down significantly before it's even half done. I know that there are many possible approaches using basic R or packages in the tidyverse suite, but I am really a bit lost with such a complex loop. I can imagine that I would try to write at least one of the loops into a function enclosing the other loop but I'm not even sure how that would work...

Again, first time posting so please let me know if I need to add anything to my question! Thanks!

Edit: Here are a small, random centred genotype matrix (X in the call to BGLR) and an accompanying phen file.

X (rows are sample IDs matching those in phen, columns are genomic locations (SNPs)):

    A  B  C  D  E  F  G  H  I  J  K  L  M  N  O  P  Q  R  S  T
1   1  1  1 -1 -1  0 -1 -1  0  1 -1  0  0  0  0 -1  0  0  1  1
2  -1  1  0 -1  0  0  1 -1  1  0 -1  1 -1  0  0  0 -1 -1 -1 -1
3  -1 -1  1 -1  0  0  1  0  0 -1 -1  0 -1  0  0 -1  1 -1  0 -1
4   0  0 -1  1 -1  0 -1  1  1  1 -1  0  0  0  1  1 -1 -1  0  1
5   1  0  1  1  0  0  1  1  0  0 -1  1  1  1  1 -1 -1  0 -1  0
6   1 -1 -1  0  1 -1  1  1  1 -1  0  0  0 -1  0  1  0 -1 -1  0
7   1 -1 -1  1  1  0 -1  1 -1  1  0 -1  1  1 -1 -1 -1  1 -1  1
8   0  1  0  1 -1 -1  0 -1 -1 -1  1 -1  0  0  0  0  0 -1  0 -1
9  -1  1 -1 -1  1  0  0  0  0  0  1  0 -1 -1 -1  0  1 -1  0 -1
10  0  0 -1  1  0  0 -1  0 -1  1  1  1  0  0  1  1  0 -1  0  0

Phen:

ID  phenotype
1   1000
2   1500
3   1200
4   500
5   700
6   2000
7   1500
8   1000
9   1300
10  900

Answer 1

I'm not familiar with BGLR but I think I can give some pointers.

First to your direct question - yes you can vectorise. A simple approach is to use one of the apply functions. lapply usually is a good fit for such loops. However this is not a magic bullet and I would strongly suggest to review the code.

I think the main thing that slows you down - besides the time BGLR needs - is how R deals with memory. The relevant properties are:

• most objects need to be placed continuous in memory
• memory isn't immediately freed - but R waits till it is necessary and does garbage collect then
• the time when R creates a copy of memory is a bit tricky. It is not when one writes an assignment - at this point the variable just holds a reference to the original data. But as soon as a change on the data is performed it creates a full copy.

All this effects don't appear in a "clean" R session, but as memory clogs things get significantly slower - agreeing with your observation.

I see following issues in the code:

Growing lists fold.Res and rep.Res are growing lists. As grow them they need to be written anew. This means first all the previous entries need to be copied (which can take some time) - but more important they can not be written in the same place as the new list takes more space. This will create fragmented memory - meaning whenever R can't fit the new list where a whole from a previous was created it will first need to garbage collect and then maybe defragment memory. I think this is the serious slow down you are experiencing.
Allocating lists unfortunately doesn't help, but with the apply functions like lapply this is also taken care of.

creating modified copies You create phen2 and phen3 as copies - but you only modify and use column 2. Meaning there is some unnecessary copying going on. Also this will lead faster to some extra garbage collects as memory fills up.

possible side effects of BGLR While the first runs are relatively free where to put their data, later may run into troubles finding suitable spots - causing again some superfluous memory cleanups. This is not so easy to deal with. I worked with stan which is quite the hog - and the best solution was to run it in separate R processes and save the results and the kill it - it was way faster than having these cleanups going on. The main scripts would then just collect the files. I described how I do it in my answer to: Using doParallel to cycle through *.rds files (specifying how to do it outside doParallel).

to style It is a bit confusing that you sometimes refer to columns by name and sometimes by column number. This makes things harder to read than necessary. I would suggest always using the name approach.

Those apply functions are awkward to write. Personally I would consider to rewrite it to use dplyr - in particular piping and do. However that may be to much work.

Answer 2

If you are doing only 100 iterations, there is probably very little you can do to optimize your code. Vectorization, preallocation, etc... are all no concern at such few iterations. Still, there are a few general things to consider:

Profiling

Before you optimize your code, the first step is always profiling your code for bottlenecks. Do this before you consider any of my (or other peoples) suggestions. just wrap all your code into profvis::profvis({..}) or use RStudios "Profiling" Menu (which does the same). The output will show you where your code spends most its time/memory in a nice gui.

Store your resample-indices in a matrix

You could create all your resamples at once and store the indices in a matrix, instead of always copying the data.frame and modifying it. This way you can also conveniently save the matrix for later reuse/reproduction if you want.

resamples <- sample(
  seq_len(folds), nrow(phen) * folds, replace = TRUE)
)
resamples <- matrix(resamples, ncol = folds)

#...

model <- BGLR(y = phen3[resamples[, j], 2], ...) 

You should also set the seed for the RNG with set.seed() (this way you can reproduce the output of sample when you run your code again)

Save intermediate steps

If you have R processes like yours that slow down the computer, and might eventually crash its a good idea to regularly save the output. Rewrite your loop in a way so that you can continue from any i and j and save your result every few minutes with saveRDS()(no need to add a timing feature, just figure out a sensible number of iterations). This will barely impact performance and you don't always have to restart from the beginning if something goes amiss.

(don't!) Parallelize the code

edit: This is a very bad idea if your code already slows the system down, but I'll leave that here for people that have similar issues with smaller models.

I would venture your profiling showed you the bottleneck of your code is the model. You can use the foreach package, or rewrite your loop as a lapply() call and use mclapply() from the parallel package, or look into the new and awesome future package. mclapply() and foreach() will also take care of the preallocating, so no need to worry about that sepparately. All of this will require some reading on your part, but its pretty simple to parallelize code in R. The easiest in your case is probably foreach.

You could not reinvent the wheel

the package caret is designed to do just what you are trying in a streamlined manner. There are also other packages with similar functionality like MLR and modelr, though I have not used those. On the other hand learning those packages will take time, your task seems fairly simple, and your code already works.