Saving & reusing adaptation in cmdstanr

Interfaces
51

I’ve been using this function, including in a fairly complicated model.
You wrote you didn’t test whether 2d-arrays work correctly. I found they do not, the elements are not sorted correctly (adding byrow=TRUE in the call to matrix() fixes it, or alternatively removing the line starting with “arrange” probably does too).

I’ve taken the liberty to:

  • fix the bug
  • wrap your function into a closure to make it pure,
  • add support for arrays of any dimensionalty (including arrays of vectors or matrices)
  • add support for initialising from previous sampling runs, nut just from warmup runs
  • some minor refactoring.

Here’s the code:

createStanInitsPreviousRun=
    # function to initialize one stan run (via cmdstanr) from the endpoint of another
    # returns a function that returns inits.
    # this function is discussed on stan forums here: https://discourse.mc-stan.org/t/saving-reusing-adaptation-in-cmdstanr/19166/51
    # known bugs:
    #   errors (complaining about missing argument) when used inside a dplyr pipeline (e.g. "fit %>% createStanInitsPreviousRun()") because non standard evaluation somehow breaks the creation of the promise because aparrently the arguments dont get evaulated at the right moment and so the data cant be captured in the environment.
    
    # # useage examples:
    # 
    # # initial pure warmup run
    # lastRunResults=mod$sample(
    #     data=data_for_stan
    #     ,chains=parallel_chains
    #     ,parallel_chains=parallel_chains
    #     # initialisation and adaptation settings
    #     ,iter_warmup=iter_warmup # run automatic initialisation and adaptation
    #     # sampling settings
    #     ,iter_sampling=0
    #     # settings to allow chains to be initialized from this one
    #     ,save_warmup=T # for getting inits from a pure warmup run
    #     ,sig_figs=18 # in crease output bitdepth to match the one used internally by stan. prevents rounding errors when initializing a run from this one (e.g. erroneously initializing at zero causing infinite density)
    #     ,output_dir="stanSamples/" # makes stan save its output to a non-temporary directory (since the r-object returned doen't contain all the results, but instead links to the file). Often needed to be able to continue sampling a model in a separate r-session or after its associated r-object has been removed (to be precise: the standard temporary output file is deleten when the r-object representing the sampling results (and thus linking to this file) has bee ngrbage collected, or the R-session holding it closed.)
    # )
    # 
    # # followup pure sampling runs
    # lastRunResults=mod$sample(
    #     data=data_for_stan
    #     ,chains=parallel_chains
    #     ,parallel_chains=parallel_chains
    #     # initialisation and adaptation settings
    #     ,iter_warmup=0 # use prespecified adaptation and inits
    #     ,adapt_engaged=FALSE # use prespecified adaptation and inits
    #     ,inv_metric=lastRunResults$inv_metric(matrix=F)
    #     ,step_size=lastRunResults$metadata()$step_size_adaptation
    #     ,init=createInitsFunction(lastRunResults)
    #     # sampling settings
    #     ,iter_sampling=iter_sample # difference
    #     # settings to allow chains to be initialized from this one
    #     ,save_warmup=T # needs to be set to TRUE even if no warmup is executed, because otherwise the function that extracts draws from the run complains since it is set to include warmup draws.
    #     ,sig_figs=18 # increase output bitdepth to match the one used internally by stan. prevents rounding errors when initializing a run from this one (e.g. erroneously initializing at zero causing infinite density)
    #     ,output_dir="stanSamples/" # makes stan save its output to a non-temporary directory (since the r-object returned doen't contain all the results, but instead links to the file). Often needed to be able to continue sampling a model in a separate r-session or after its associated r-object has been removed (to be precise: the standard temporary output file is deleten when the r-object representing the sampling results (and thus linking to this file) has bee ngrbage collected, or the R-session holding it closed.)
    # )
function(
    lastRunResults # cmdstanr fit which was fitted with include_warmup=T (this is required even if the number of warmup draws was zero, because of an API limitation)
)
{
    initsFunction = function(chain_id){ # modified from https://discourse.mc-stan.org/t/saving-reusing-adaptation-in-cmdstanr/19166/44
        #   note: this function is not pure on its own. It captures lastRunResults from the enclosing environment, forming a pure closure.
        lastRun_draws = lastRunResults$draws(inc_warmup=T) # note that this row errors when used on a model which wasnt run with save_warmup=TRUE.
        final_warmup_value = lastRun_draws[dim(lastRun_draws)[1],chain_id,2:(dim(lastRun_draws)[3])]
        final_warmup_value %>% 
            tibble::as_tibble(
                .name_repair = function(names){
                    dimnames(final_warmup_value)$variable
                }
            ) %>% 
            tidyr::pivot_longer(cols=dplyr::everything()) %>% 
            tidyr::separate(
                name
                , into = c('variable','index')
                , sep = "\\["
                , fill = 'right'
            ) %>% 
            dplyr::mutate(
                index = 
                    stringr::str_replace(index,']','') %>% 
                    {ifelse(!is.na(.),.,"")}
            ) %>% 
            dplyr::group_split(variable) %>% 
            purrr::map(
                .f = function(x){
                    out= 
                        x %>%
                        dplyr::mutate(
                            index=strsplit(index,",",fixed=TRUE)
                            ,maxIndexes=
                                do.call(rbind,index) %>%
                                {array(as.numeric(.),dim=dim(.))} %>%
                                plyr::aaply(2,max) %>%
                                list()
                        ) %>%
                        {
                            array(
                                .$value
                                ,dim=
                                    if(length(.$maxIndexes[[1]])>0){
                                        .$maxIndexes[[1]]
                                    } else {
                                        1
                                    }
                            )
                        } %>% 
                        list()
                    names(out) = x$variable[1]
                    out
                }
            ) %>% 
            unlist(recursive=F)
    }
    force(lastRunResults) # force evaluation of lastRunResults (i could've just written "lastRunResults=lastRunResults") to change its promise to an actual value that can get captured by the closure (since R has lazy evaluation objects dont get passed by reference or by copy, but by passing a promise to the object, if the object changes before the promise is evaluated then the promise returns the changed object instead of the "promised" one. Thus in that case the closure captures the changed object instead of the one originally passed (by promise) to its environment.)
    initsFunction # return the closure
}

Please note that this implementation throws an error when used as an elemen of a pipe (see comments under “known bugs” in the code.

Edit:
removed some unnecessary parantheses in the code.
clarified one comment in code
removed mention of Mike-Lawrence in the code because I didn’t ask for his permission to include him and the post clearly references his contribution.

4 Likes