Puzzled about the array ordering of LKJCC corr as xarray object

The anti-diagonal symmetry makes me suspect this is an issue with order="F" or order="C" when unpacking the cholesky or something of the sort. It is indeed very odd.

The missmatch is between the results shown in plot_pair and in the correlations stored in corr, however, due to xarray being involved and the mixture of positional and label based indexing the situation is a bit more puzzling.

I’d recommend you use idata_kwargs to also define the dims in corr. That should look like:

b0_names = ['gc_omega_b0', 'gc_theta_b0', 'ulr_b0']
self.coords.update({'b0_names': b0_names, "b0_names2": b0_names})
...
    idata_kwargs={"dims": {"lkjcc_corr": ["b0_names", "b0_names2"]}},

something very similar is done in the model at https://www.pymc.io/projects/examples/en/latest/case_studies/multilevel_modeling.html that uses LKJ too.

That will have the corr variable have meaningful labels instead of the default assigned dimensions. I would then do two things:

corr_means = idata.posterior["lkjcc_corr"].mean(("chain", "draw"))
corr_means.sel(b0_names="gc_omega_b0", b0_names2="ulr_b0")

does this match the correlation for that plot in plot_pair?

Another thing you can try is instead of looping using the tril_indices to call plot_posterior you can try the approach from Label guide — ArviZ dev documentation, do the labels here and in plot_pair match?

2 Likes