warmup samples not being properly split between posterior and warmup posterior

Describe the bug when using save warmup in cmdstanpy and then passing to arviz InferenceData object, the entire chain(s) are passed to the warmup posterior instead of being split between the posterior and the warmup posterior.

Though, the posterior shows the dimension of the entire chain… but having an empty array.

To Reproduce

versions:

Attributes:
    created_at:                 2021-04-19T23:47:52.469962
    arviz_version:              0.11.2
    inference_library:          cmdstanpy
    inference_library_version:  0.9.68

sample with cmdstanpy and set:

iter_warmup=1000 iter_sampling=400 save_warmup=True

full call is:

fit = m.model.sample(data=data,
                     parallel_chains = n_chains,
                     chains=n_chains,
                     threads_per_chain=n_threads,
                     seed=234568,
                     iter_warmup=1000,
                     iter_sampling=400,
                     max_treedepth=11,
                     show_progress=True,
                     save_warmup=True
                   )

s = av.from_cmdstanpy(fit, save_warmup=True)



av.to_netcdf(s, "fit.h5")

now I reload the fit:

fit = arviz.from_netcdf("fit.h5")

fit.posterior.dim

Out[16]: Frozen(SortedKeysDict({'chain': 0, 'draw': 1400, 'index_raw_dim_0': 36, 'log_K_raw_dim_0': 36, 'index_dim_0': 36, 'log_K_dim_0': 36, 'K_dim_0': 36}))

and the values of one of the parameters

In [17]: fit.posterior.index_mu.values
Out[17]: array([], shape=(0, 1400), dtype=float64)

the warmup should be 1000 draws and the posterior 400.

If I look at the warmup posterior

In [18]: x.warmup_posterior.dims
Out[18]: Frozen(SortedKeysDict({'chain': 2, 'draw': 1400, 'index_raw_dim_0': 36, 'log_K_raw_dim_0': 36, 'index_dim_0': 36, 'log_K_dim_0': 36, 'K_dim_0': 36}))

In [19]: x.warmup_posterior.index_mu.values
Out[19]:
array([[ 0.81948,  0.81948,  0.81948, ..., -1.97329, -2.02469, -1.99088],
       [-1.44893, -1.44893, -1.44893, ..., -2.13466, -2.1432 , -2.17423]])

Expected behavior I would expect the warmup posterior and posterior to be split according to the number of samples specified in the sampling call, but perhaps I miss something?

[ahartikainen]

This looks like a bug in our end. Probably some typo in the splitting procedure. Let’s see, I try to replicate the results.

[grburgess]

Just confirming that this is fixed with a pip install from the master.