Passing a different mass matrix for each chain in cmdstanr

Interfaces
1

Based on the description in Run Stan's MCMC algorithms — model-method-sample • cmdstanr this doesn’t seem possible, but it would be a required feature when saving the warmed up parameters from multiple chains. By contrast, it’s straightforward to pass different step sizes for each chain.

2

This works in CmdStanPy.

3

It accepts a metric per chain, just not documented as having that capability. You just do inv_metric = warmup$inv_metric(matrix=F). Here’s the full warmup-then-sample code:


warmup = mod$sample(
	data = data_for_stan
	, 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_inits = 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_for_stan
	, chains = parallel_chains
	, parallel_chains = parallel_chains
	, refresh = 0
	, show_messages = F
	
	, seed = seed+1
	
	, 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_inits
)
1 Like
4

Noice!

I won’t always be expecting this, but I love how quickly folks have been responding. Maybe we’re all working on the same problem, rn.

5

@jonah and @rok_cesnovar, the documentation for cmdstanr should be adjusted.

1 Like
6

Yeah, been working on the composed warmup+sampling since last week (inc the bug I alerted you to in another thread). Ben Bales just now worked out why that bug was happening.

2 Likes
7

Thanks @charlesm93 and @mike-lawrence . will make an issue to fix doc.

8

FYI the bug was worked out, and I also posted some code that achieves things properly.