More divergences with composed warmup-then-sample than combined warmup-and-sample

[mike-lawrence]

Take the 8schools_ncp example:

library(cmdstanr)
example_program = 'schools_ncp.stan'
example_data = 'schools_ncp.data.json'
tmp <- file.path(tempdir(), example_program)
if (!file.exists(tmp)) {
  file.copy(system.file(example_program, package = "cmdstanr"), tmp)
}
mod <- cmdstan_model(tmp)
data_file <- system.file(example_data, package = "cmdstanr")

seed = 1
iter_warmup = 1e3
iter_sample = 1e3
parallel_chains = 4

#first the typical combined warmup-and-sample run:
both = mod$sample(
  data = data_file
  , chains = parallel_chains
  , parallel_chains = parallel_chains
  , refresh = 0
  , show_messages = F
  , seed = seed
  , iter_warmup = iter_warmup
  , iter_sampling = iter_sample
)

Which yields no divergences. Now an attempt at doing warmup and sampling separately, using the adapted info from the former in the latter:

warmup = mod$sample(
  data = data_file
  , chains = parallel_chains
  , parallel_chains = parallel_chains
  , refresh = 0
  , show_messages = F
  , seed = seed
  , iter_warmup = iter_warmup
  , save_warmup = T #for inits
  , sig_figs = 18
  , iter_sampling = 0
)

get_sampling_inits_from_warmup = function(chain_id){
	warmup_draws = warmup$draws(inc_warmup=T)
	final_warmup_value = warmup_draws[iter_warmup,chain_id,]
	init_list = as.list(final_warmup_value)
	names(init_list) = dimnames(final_warmup_value)[[3]]
	init_list = init_list[names(init_list)!='lp__']
	return(init_list)
}

samples = mod$sample(
  data = data_file
  , chains = parallel_chains
  , parallel_chains = parallel_chains
  , refresh = 0
  , show_messages = F
  , seed = seed
  , iter_warmup = 0
  , adapt_engaged = FALSE
  , inv_metric = warmup$inv_metric(matrix=F)
  , step_size = warmup$metadata()$step_size_adaptation
  , iter_sampling = iter_sample
  , init = get_sampling_inits_from_warmup
)

And now we get divergences from the sampling run. I have more rigorous testing of this detailed here, but the gist is that doing the warmup-then-sample approach is somehow generating more divergences than the combined warmup-and-sample approach despite to all examination all the inputs being as expected. I don't believe this is an issue with cmdstanr's passing of the pertinent info to cmdstan because I've checked the init jsons and output csvs and they have all the expected content.

[mike-lawrence]

Actually, here's some data and a model that seems to more reliably show the difference:
divergence_demo.zip

And here's a script to loop over lots of seeds, showing that with nearly every one the number of divergences in the composed approach is far higher than those in the standard approach:
composed_divergences.zip

[mike-lawrence]

And here's the functions I used to confirm that the inputs are as expected in the csv files (implying that this is not a cmdstanr problem):

check_inputs = function(chain){
	sampling_step_size = as.numeric(strsplit(readLines(samples$output_files()[chain])[42],'=')[[1]][2])
	warmup_step_size = warmup$metadata()$step_size_adaptation[chain]
	step_size_ok = dplyr::near(
		sampling_step_size
		, warmup_step_size
	)
	
	temp = readLines(samples$output_files()[chain])[44]
	temp = str_replace(temp,'#','')
	sampling_inv_metric = as.numeric(unlist(strsplit(temp,',',)))
	warmup_inv_metric = unlist(warmup$inv_metric(matrix=F)[chain])
	inv_metric_ok = all(dplyr::near(
		sampling_inv_metric
		, warmup_inv_metric
	))
	
	samples_inits = 
		(
			tibble::tibble(
				v = readLines(samples$metadata()$init[chain])
			)
			%>% dplyr::slice(2:(n()-1))
			%>% tidyr::separate(
				v
				, into = c('par','value')
				, sep = ':'
			)
			%>% dplyr::mutate(
				value = as.numeric(stringr::str_replace(value,',',''))
			)
			%>% dplyr::pull(value)
		)
	warmup_inits = as.numeric(get_sampling_inits_from_warmup(chain))
	inits_ok = all(dplyr::near(samples_inits,warmup_inits))
	return(tibble::tibble(chain,step_size_ok,inv_metric_ok,inits_ok))
}

purrr::map_dfr(1:parallel_chains,check_inputs)

[mitzimorris]

yes, this happens. it's not a bug - it's an artifact of the sampler trying to take the specified initial parameters values and meet model constraints. also, even when specifying the seed, the RNG is in a different state.

[mitzimorris]

reopening - @bbbales2 et al can investigate.

[mike-lawrence]

Here's a gist containing the stan file, the code to generate data and the code to loop over seeds, generating new data and attempting to sample in the two ways for each seed.

Data and model are a standard hierarchical model where there are 10 subjects each observed with common gaussian measurement noise 10 times in each of 2 conditions, and the subjects' intercept and condition effect are modelled as multivariate normal.

[mike-lawrence]

Huh, running that just crashed my rstudio, but only once it was about halfway done. I'll try again and see if it continues to occur...

[mike-lawrence]

Ran it a few times now and no more crashes, so possibly the crash was unrelated (though I didn't have much else happening on my computer at the time of the crash; plenty of free RAM and CPU).

[mike-lawrence]

Note: confirmed to occur with cmdstanpy as well as cmdstanr, reinforcing that this is core to cmdstan itself. I also confirmed it persists when using the development version of cmdstan.

[ahartikainen]

Should we extract the cmd used for calling cmdstan?

In cmdstanpy

fit.runset.cmd

[mitzimorris]

there's a slight complication because the command is composed on the fly w/ the chain id - cmdstan_call(chain_id)
or something like that

[mike-lawrence]

@mitzimorris @jgabry @rok-cesnovar Tagging you as core interface devs. @bbbales2 Investigated this, using my hierarchical model where I observe the most reliable increased divergences rates. He confirmed that the interfaces are creating an init json with the proper info, but found that if he tweaked the structure of these jsons the increased-divergences behaviour goes away.

Here's what both interfaces produce for the inits json:

{
  "chol_corr[1,1]": 1,
  "chol_corr[2,1]": -0.469573042026972,
  "chol_corr[1,2]": 0,
  "chol_corr[2,2]": 0.882893627908559,
  "mean_coef_[1]": -0.0305049133852781,
  "mean_coef_[2]": 0.164671727959854,
  "sd_coef_[1]": 0.00665194669853791,
  "sd_coef_[2]": 0.148205452770937,
  "multi_normal_helper[1,1]": -0.0216987316133846,
  "multi_normal_helper[2,1]": 0.724235080770058,
  "multi_normal_helper[1,2]": 0.43239986307413,
  "multi_normal_helper[2,2]": 1.75514903811538,
  "multi_normal_helper[1,3]": -0.691867572248652,
  "multi_normal_helper[2,3]": -2.66409204621652,
  "multi_normal_helper[1,4]": -0.395599105436074,
  "multi_normal_helper[2,4]": -0.231732460746632,
  "multi_normal_helper[1,5]": 0.439963477396995,
  "multi_normal_helper[2,5]": -0.460156658379562,
  "multi_normal_helper[1,6]": 1.46417143002824,
  "multi_normal_helper[2,6]": -1.2165899374329,
  "multi_normal_helper[1,7]": 1.54988112300058,
  "multi_normal_helper[2,7]": -0.90213563970599,
  "multi_normal_helper[1,8]": 0.740397448247888,
  "multi_normal_helper[2,8]": 1.32996651536118,
  "multi_normal_helper[1,9]": 0.456738742343979,
  "multi_normal_helper[2,9]": -0.306703562089387,
  "multi_normal_helper[1,10]": -0.530873450669662,
  "multi_normal_helper[2,10]": 1.49202297709515,
  "noise_": 0.981719532981492
}

And here's the alternative structure that he observes as solving the issue:

{
    "chol_corr": [[1, 0], [-0.469573042026972, 0.882893627908559]],
    "mean_coef_": [-0.0305049133852781, 0.164671727959854],
    "sd_coef_": [0.00665194669853791, 0.148205452770937],
    "multi_normal_helper": [[-0.0216987316133846, 0.43239986307413, -0.691867572248652, -0.395599105436074, 0.439963477396995, 1.46417143002824, 1.54988112300058, 0.740397448247888, 0.456738742343979, -0.530873450669662],
			    [0.724235080770058, 1.75514903811538, -2.66409204621652, -0.231732460746632, -0.460156658379562, -1.2165899374329, -0.90213563970599, 1.32996651536118, -0.306703562089387, 1.49202297709515]],
    "noise_": 0.981719532981492
}

So is this something that needs to be addressed in the interfaces, or in cmdstan?

[rok-cesnovar]

In that its a bug in the interfaces. The below example is how it should be.

[ahartikainen]

At least CmdStanPy doesn't create multidimensional objects yet.
Is there possibility to fix this in CmdStan side?