venpopov / chkptstanr

Checkpoint Stan R
https://venpopov.github.io/chkptstanr/
Other
6 stars 1 forks source link

Simplifying chkpt_brms #11

Open frankhezemans opened 6 months ago

frankhezemans commented 6 months ago

This is just a minor suggestion to help simplify chkpt_brms.R (perhaps @venpopov you had already thought of it). You don't need a big if-else loop to separately run brms::make_standata and brms::make_stancode for the cases with or without threading (currently these lines).

Instead, you can declare the number of threads, num_threads, to be NULL if the input argument threads_per was equal to one, and otherwise set num_threads to threads_per. After that, you can do one call to brms::make_standata and brms::make_stancode, both including the input argument threads = brms::threading(threads = num_threads). This is because calling brms::threading with threads = NULL is the default value, which will lead to no threading, as desired (see these and these lines).

So something like this should do:

num_threads <- ifelse(threads_per > 1L, threads_per, NULL)
threads_list <- brms::threading(threads = num_threads)

stan_data <- brms::make_standata(
  formula = formula,
  data = data,
  threads = threads_list,
  ...
)
stan_code <- brms::make_stancode(
  formula = formula,
  data = data,
  threads = threads_list,
  ...
)
venpopov commented 6 months ago

Already done almost exactly like that as part of a big refactoring of the code I'm doing. This is legacy code from the original package, and my first goal was to get it to work ASAP, now onto doing exactly things like you suggested. There are many unnecessary if else clauses like that one, and also a lot of the code between chkpt_brms and chkp_stan is the same and can be abstracted away. Stay tuned for the next release :) and do post more if you have any other suggestions for features, I appreciate all feedback and requests!

frankhezemans commented 6 months ago

Excellent, I'm looking forward to the next release!