Arviz: More intuitive chain extraction methods

[grburgess]

Short Description

I'm happy to work on this myself, but before going down that path, is there a simple way to extract samples combined across multiple chains? I know how to do it programmatically, but it would be have this upfront for the user

Code Example or link

something like

param_div_free, param_div = fit.posterior.parameter.extract(combined=True, split_divergence=True)

Also include the ArviZ version and version of any other relevant packages.

Relevant documentation or public examples

I base this off of some of the functionality I enjoyed in the stan utility here:

https://github.com/betanalpha/jupyter_case_studies/issues/4

and things I added here:

https://github.com/grburgess/stan_utility

which I would prefer to abandon and add effort to this project.

[OriolAbril]

We were planning on adding an example about this somewhere, maybe xarray_examples to begin with.

Currently the way to go is one of:

idata.posterior.stack(sample=("chain", "draw")).parameter.values

will generate a xarray dataset object, adding .parameter will get a dataarray which is generally interchangeable with numpy arrays which could be directly retrieved adding .values too. Depends on the use at hand. The shape of the retrieved array will be (*shape, sample) instead of the (chain, draw, *shape) used by ArviZ as convention.

Another option is to do

idata.posterior.parameter.values.flatten()

Which does not do the same unless parameter is a 1D variable but can still be useful in some cases.

There is also the option of adding a new method to inferencedata:

idata.extract(var_name, group, coords, combined)

This would return a dataarray/numpy array given a variable name (required argument), a group that could be set to default posterior, and optionally, subsetting with coords dict like in plots and a combined argument defaulting to False

I don't know what is split_divergence argument supposed to do though, so for now I can't help with that.

[ahartikainen]

I think split (or groupby) could take same specific array / variable and split idata to multiple idata.

So idata.groupby('sample_stats/divergence') would group by divergence. There is extra logic needed for coords.

[OriolAbril]

group by divergence as in generating a groupby with 2 groups, one with all samples with diverging=False and the other with True?

[ahartikainen]

For example in boolean case

[OriolAbril]

@percygautam can you try something like the code below on #1254?

def diverging_groupby(ds, diverging):
    return ds.groupby(diverging)

idata_aux = idata.stack(sample=("chain", "draw"))
idata_aux.map(diverging_groupby, groups="posterior", diverging=idata_aux.sample_stats.diverging)

given that patterns like this are probably common, would it make sense to add a idata_args to idata.map? then the code below could be simplified to:

idata.stack(sample=("chain", "draw")).map(diverging_groupby, groups="posterior", idata_args="sample_stats/diverging")

idata.map would then first retrieve and store idata_args and assume that the call signature of the function is func(ds, *(*args, *idata_args), **kwargs). idata_args would be a string or list of strings following the pattern group/var_name (thus following netCDF convention).

I personally think that we should not include coords/sel options here, if necessary, the user can always call .sel or .isel before like with stack.

[percygautam]

I tried the code. It's working fine and returns a DatasetGroupBy object for idata.posterior.

[canyon289]

General note. Going to start tagging this with xarray label for things that fall under "Would be nice if we could do this". Ive got a couple from my bambi work as well

[grburgess]

Thanks for the update. I'm not at all familiar with xarray and will have to get used to it. I will play around with some fits and perhaps make some notebooks/docs for a PR from my own experience.

[grburgess]

Just to update, the above specified instructions work very well and give me the access I wanted.

[canyon289]

Closing as this seems to have been addressed

[OriolAbril]

It is probably worth it to tag #1469 to add even better ways to extract the draws combining chains, as simple chains... and to #641 to create a "Working with InferenceData" notebook that should probably show this