Checkpoint brms with data2 argument

Modeling
1

Hi all,

Just wondering if it’s possible to use chkpt_brms() with the “data2” argument from the normal brms function. I normally use it to specifiy the correlation matrix for a random effect, but I can’t see how I can add it to the chkpt_brms() function.

I’ve added in a simple example of the model formula that I currently use (but takes too long to run).

# Model formula.
model_formula <- paste0("habitat_density ~ seasonality + (1|gr(tree_tip, cov=A))")
 
# brms formula.
brms_formula <- brmsformula(model_formula, family = cumulative(threshold = "equidistant"))

# # Add priors.
normal_priors <- c(prior(normal(0,1), class="Intercept"),
                   prior(normal(0,1), class="b"),
                   prior(gamma(2,1), "sd")) 

  brms_model <- brm(
    brms_formula,
    data = model_data,
    data2 = list(A=model_covar),
    prior = normal_priors,
    iter = 10000,
    warmup = 5000,
    chains = 2,
    thin = 20,
    cores = 32,
    init = 0,
    file = model_pathway,
    normalize = FALSE,
    backend = "cmdstanr",
    threads = threading(16),
  )

In the above code, I specify the co-variance matrix using the data2 arguement with “list(A = model_covar)”, which responds to “(1|gr(tree_tip, cov=A))” in the formula.

If I try and test code on a smaller sample of data using chkpt_brms() i get an error.
This is my checkpoint code:

  # Run brms checkpoint model.
  brms_model <- chkpt_brms(
    brms_formula,
    data = model_data,
    data2 = list(A=model_covar),
    prior = normal_priors,
    iter_sampling  = 1000,
    iter_warmup = 500,
    iter_per_chkpt = 100,
    parallel_chains = 2,
    threads_per = 2,
    path = model_pathway)

The output looks normal to begin with but has an error at the end:

Initial Warmup (Typical Set)
Chkpt: 1 / 15; Iteration: 100 / 1500 (warmup)
Chkpt: 2 / 15; Iteration: 200 / 1500 (warmup)
Chkpt: 3 / 15; Iteration: 300 / 1500 (warmup)
Chkpt: 4 / 15; Iteration: 400 / 1500 (warmup)
Chkpt: 5 / 15; Iteration: 500 / 1500 (warmup)
Chkpt: 6 / 15; Iteration: 600 / 1500 (sample)
Chkpt: 7 / 15; Iteration: 700 / 1500 (sample)
Chkpt: 8 / 15; Iteration: 800 / 1500 (sample)
Chkpt: 9 / 15; Iteration: 900 / 1500 (sample)
Chkpt: 10 / 15; Iteration: 1000 / 1500 (sample)
Chkpt: 11 / 15; Iteration: 1100 / 1500 (sample)
Chkpt: 12 / 15; Iteration: 1200 / 1500 (sample)
Chkpt: 13 / 15; Iteration: 1300 / 1500 (sample)
Chkpt: 14 / 15; Iteration: 1400 / 1500 (sample)
Chkpt: 15 / 15; Iteration: 1500 / 1500 (sample)
Checkpointing complete
Error: Object 'A' was not found in 'data2'.

I’m not sure how to diagnose the problem, so please let me know if there’s extra info I need to supply as well.

Any help would be very appreciated!

2

I have fixed this issue in a fork, as described in this other post with the same issue: Running speed - phylogenetic models - #18 by Ven_Popov

3 Likes
3

PS: would be great if you can test it and share if it works (even if you have moved on from this project - just for the sake of testing). I couldn’t run your code specifically because I don’t have the data, but it ran fine after the fix with this example (before the fix it gave the same error as yours).