Kfold CV in brms

I am trying to use kfold CV as a means of evaluating a model run using brms and I feel like I'm missing something. As a reproducible example, my data are structured as a binary response (0, 1) dependent on the length of an individual. Here is some code to generate and plot data similar to those I am working with:

library(brms)
library(tidyverse)
library(loo)

length <- seq(0, 100, by = 1)
n_fish_per_length <- 10

a0 <- -48
a1 <- 2
a2 <- -0.02

prob <- plogis(a0 + a1 * length + a2 * length^2)

plot(length, prob , type = 'l')

sim_data <-
  expand_grid(fish_id = seq_len(n_fish_per_length),
              length = length) %>%
  mutate(prob_use =  plogis(a0 + a1 * length + a2 * length^2)) %>%
  mutate(is_carp = rbinom(n = n(), size = 1, prob= prob_use))

ggplot(sim_data, aes(x = length, y = is_carp)) +
  geom_jitter(width = 0, height = 0.05) +
  geom_smooth(method = "glm", formula = y ~ x + I(x^2),
              method.args = list(family = binomial(link = "logit")))

I then use brms to run my model.

Bayes_Model_Binary <- brm(formula = is_carp ~ length + I(length^2),  
                          data=sim_data, 
                          family = bernoulli(link = "logit"),
                          warmup = 2500, 
                          iter = 5000, 
                          chains = 4, 
                          inits= "0", 
                          cores=4,
                          seed = 123)

summary(Bayes_Model_Binary)

I'd like to use kfold CV to evaluate the model. I can use something like this:

kfold(Bayes_Model_Binary, K = 10, chains = 1, save_fits = T)

but the response in my data is highly imbalanced (~18% = 1, ~82% = 0) and my reading suggests that I need to used stratified kfold cv to account for this. If I use:

sim_data$fold <- kfold_split_stratified(K = 10, x = sim_data$is_carp)

the data are split the way I would expect but I'm not sure what the best way is to move forward with the CV process from here. I saw this post https://mc-stan.org/loo/articles/loo2-elpd.html, but I'm not sure how to modify this to work with a brmsfit object. Alternatively, it appears that I should be able to use:

kfold(Bayes_Model_Binary, K = 10, folds = 'stratified', group = sim_data$is_carp)

but this throws an error. Likely because is_carp is the response rather than a predictor in the model. What would my group be in this context? Am I missing/misinterpreting something here? I'm assuming that there is a very simple solution here that I am overlooking but appreciate any thoughts.

Solution

After some additional digging and learning how to access information about each fold in the analysis, I was able to determine that the structure of the data (proportion of 0s and 1s in the response) is maintained using the default settings in the kfold() function. To do this I used the following code.

First, save the kfold CV analysis as an object.

kfold1 <- kfold(Bayes_Model_Binary, K = 10, save_fits = T)

kfold1$fits is a list of the model fitting results and the observations used in the test data set (omitted) for each fold.

From this information, I created a loop to print the proportion of observations in each training data set where is_carp = 1 (could also do this for each test data set) with the following code.

for(i in 1:10){
    print(length(which(sim_data$is_carp[-kfold1$fits[i, ]$omitted] == 1)) / 
           nrow(sim_data[-kfold1$fits[i, ]$omitted, ]))
}

[1] 0.1859186
[1] 0.1925193
[1] 0.1991199
[1] 0.1914191
[1] 0.1881188
[1] 0.1848185
[1] 0.1936194
[1] 0.1980198
[1] 0.190319
[1] 0.1870187

and it's easy to then compare these proportions with the proportion of observations where is_carp = 1 from the original data set.

length(which(sim_data$is_carp == 1)) / nrow(sim_data)

[1] 0.1910891