Issues with plotting when combining chains and draws of InferenceData

Questions v5
1

Hi all,

I think this a bit of a basic question regarding working with InferenceData returned from pm.sample(), but here goes:
I’m combining chains and draws using the “extract” function in ArviZ (as described here) in order to simplify some manipulations I want to perform on the posterior variables.

For example, let’s say the variables in my posterior originally have dimensionality (4, 1000). Using “extract”, the dimensionality (in the normal Numpy sense, not xarray sense) of the returned variables is what I expect (i.e., (4000,). I perform some computations with the lower dimensional data and everything is good so far. However, rather than being treated as 4000 “draws” from a single chain, which would work nicely with ArviZ plotting functions like “plot_posterior”, they are now 4000 “samples”, which plot_posterior doesn’t like (it tries to create a figure panel for each sample). The simplest workaround I’ve found is to use the “unstack” function to return the manipulated variables to their original dimensionality (4, 1000) and then input them into plot_posterior. While this works fine, it seems pretty clumsy and like there must be a simpler way to plot the combined (stacked) version of the posterior variables.

Appreciate any advice you may have.

2

Does plot_posterior(..., combine_dims={"sample"}) work?

3

Yup, that did it. Thank you. I’m not sure I understand why. I thought combining chains and draws would be equivalent to concatenating chains, but it seemed to increase dimensionality (i.e., creating “sample” axis). However, that’s for me to dig deeper.

4

It reduces the number of dimensions, but it creates a new dimension with a different name that didn’t exist before (and does so keeping the info on the original dims that were flattened to create it). In xarray and arviz only the dimension name matters. For now arviz plotting functions default to reducing dimensions called chain and/or draw and facetting over the rest. To reduce dimensions with a different name you need to use combine_dims argument

5

Gotcha. Thanks for the helpful explanation.