Arviz: Unexpected behavior when using `pm.hpd` with multidimensional (ndim>2)posterior arrays

[narendramukherjee]

Describe the bug
I am using hpd on multi-dimensional posteriors with larger than 2 dimensions. One example of where this can come up is if I use a 2D array to represent the coefficients for the interaction between levels of 2 different predictors. The output trace of this 2D array of coefficients from PyMC3 will now be 3D, with the 0th dimension representing the MCMC samples.

I find that with a 3D trace/array, hpd treats dimension 1 as the MCMC dimension (instead of dimension 0, the actual MCMC dimension)

To Reproduce

#12000 MCMC samples from a 10x3 dimensional distribution
a = np.random.normal(size=(12000, 10, 3))
pm.hpd(a).shape

Output: (3, 12000, 2)
Should instead be: (10, 3, 2)

This happens because the assumption in hpd is that ndim==2 if ndim>1: https://github.com/arviz-devs/arviz/blob/master/arviz/stats/stats.py#L346

for row in ary.T works properly for arrays that have 2 dimensions, but for anything larger than that, it infers the wrong MCMC dimension.

Suggestion: Will it make sense to have an axis parameter to specify the axis of the MCMC samples, much like say np.mean etc?

[ahartikainen]

Possible solution:

First dimension = chains
Second dimension = draws
3rd --> = parameter dimensions.

You need to add one dimension to you data.

Or use pymc3 trace object with vars info?

[AlexAndorra]

Ah yes I had to do that the other day. My solution was to slice the array you give to az.hpd: az.hpd(a[:, :, 0]).shape should give you (10, 2) -- I'm using az instead of pm here as PyMC is based on ArviZ for diagnostics.

So, if the 3in the shape corresponds to categories for instance, you do az.hpd(a[:, :, i]) for each of them and you should get a (10, 3, 2) by concatenation and transposition 🎉

That being said, implementing an axis paremeter could indeed be interesting -- if easily doable -- but I'm no expert here.

[narendramukherjee]

@AlexAndorra Yeah, I used the concatenation and transposition idea that you outlined for my use-case, but wondering if its good to have some sort of general structure in arviz that clears up the confusion that might be caused by multiple array dimensions.

[AlexAndorra]

Glad it worked out!
As for adding a general structure to ArviZ, I find it a good idea but I'll defer to the core devs -- in particular, I don't know how widespread the use case is, and how painful it would be to add this functionnality to ArviZ.

[OriolAbril]

I think the most coherent approach here would be to use make_ufunc and wrap_xarray_ufunc, similarly to what is done with ess or rhat

[percygautam]

I'll like to work on this.