Can't initialize chains with different inits via list of dicts. #362

import cmdstanpy
import numpy as np
import re

model = cmdstanpy.CmdStanModel(stan_file='mwe_inits.stan')
no_states = 2
ts = np.linspace(0,1,4)[1:]
no_timesteps = len(ts)
data = dict(
    no_states=no_states,
    no_timesteps=no_timesteps,
    ts=ts,
    observed_states=np.zeros((no_timesteps, no_states))
)

prior_fit = model.sample(
    dict(data, likelihood=0),
)
predicted_states = prior_fit.stan_variable('predicted_states')
idxs = np.random.choice(len(predicted_states), size=5, replace=False)

try:
    print('This does not work. Wants either a list of strings or a dict.')
    inits = [
        dict(log_initial_states=prior_fit.stan_variable('log_initial_states')[idx])
        for idx in idxs[1:]
    ]
    post_fit = model.sample(
        dict(data, likelihood=1, observed_states=predicted_states[idxs[0]]),
        inits=inits
    )
except Exception as ex:
    print(ex)

print('This does the wrong thing, initializing all chains with the same values.')
inits = dict(log_initial_states=prior_fit.stan_variable('log_initial_states')[idxs[1:]])
post_fit = model.sample(
    dict(data, likelihood=1, observed_states=predicted_states[idxs[0]]),
    inits=inits,
    iter_warmup=1, iter_sampling=0
)

for i, path in enumerate(post_fit.runset.stdout_files):
    with open(path) as fd:
        out = fd.read()
    print(f'chain {i}: ' + re.findall('initial.+', out)[0])
print(f'correct inits: ', inits['log_initial_states'])

functions {
  vector dydt(real t, vector y){
    return -y;
  }
}
data {
  int no_states;
  int no_timesteps;

  real ts[no_timesteps];
  vector[no_states] observed_states[no_timesteps];

  int likelihood;
}
parameters {
  vector[no_states] log_initial_states;
}
transformed parameters {
  print("initial state: ", log_initial_states);
  vector[no_states] initial_states = exp(log_initial_states);
  vector[no_states] true_states[no_timesteps] = ode_rk45(
    dydt, initial_states, 0, ts
  );
}
model {
  log_initial_states ~ normal(0,1);
  if (likelihood) {
    for (i in 1:no_timesteps) {
      observed_states[i] ~ lognormal(log(true_states[i]), .1);
    }
  }
}
generated quantities {
  vector[no_states] predicted_states[no_timesteps];
  for (i in 1:no_timesteps) {
    predicted_states[i] = to_vector(lognormal_rng(log(true_states[i]), .1));
  }
}

This does not work. Wants either a list of strings or a dict.
List element 0 must be a filename string, found {'log_initial_states': array([-2.10363 ,  0.931937])}
List element 1 must be a filename string, found {'log_initial_states': array([-1.54445, -1.14433])}
List element 2 must be a filename string, found {'log_initial_states': array([-0.712222, -0.865938])}
List element 3 must be a filename string, found {'log_initial_states': array([-0.958756 ,  0.0936526])}
This does the wrong thing, initializing all chains with the same values.
chain 0: initial state: [-2.10363,-1.54445]
chain 1: initial state: [-2.10363,-1.54445]
chain 2: initial state: [-2.10363,-1.54445]
chain 3: initial state: [-2.10363,-1.54445]
correct inits:  [[-2.10363    0.931937 ]
 [-1.54445   -1.14433  ]
 [-0.712222  -0.865938 ]
 [-0.958756   0.0936526]]