Arviz: InferenceData extensions [Discussion]

[OriolAbril]

InferenceData objects are central to ArviZ, and even though a common subset of tasks using InferenceData can be done directly with ArviZ plotting and stats functions, any task that deviates from this becomes more and more convoluted and long.

The aim of this issue is to start a discussion about new capabilities to add to InferenceData and generate a proposal (which will be added to xarray_examples for discussion with xarray team).

I also think there are several groups of functions, if it may help start brainstroming or generating different proposals per group. Ideas on all levels are welcome!

Straightforward extensions to xr.Dataset methods

.sel is a good example of this. I think several methods could fit in this category and very roughly follow a similar pattern:

def idata_extension(self, groups, ... , **kwargs):
    for group in groups:
        if group not in self._groups:
            raise Error
        # some kind of check to make method as convenient as possible
        # an example is sel using only the dimensions present in current group to index
        dataset = getattr(self, group)
        setattr(self, group, datasel.method(**kwargs)

In addition to groups we should think about other ArviZ specific args, common in most functions and not passed to xarray. Maybe inplace and/or copy?

Also, groups could accept groups and some _metagroups_ so that one keyword represents several proper groups. We could go as far as adding the metagroups dict in rcParams. One metagroup example could be "posteriors" -> ("posterior", "sample_stats", "log_likelihood", "posterior_predictive")

Some ideas of functions that could fit in this category are:

• .isel
• stack and unstack
• rename, rename_dims and rename_vars
• .load and other dask related methods like chunk would be interesting after ArviZ starts becoming Dask friendly.

Many dataset methods make sense to extend, so I think we should focus on the ones that solve more issues on our side. For example, if we make an extension to apply_ufunc compatible with inference data or extend the map method, the mean, median, max... are not really necessary, only convenient, whereas other methods may have no alternative.

Commenting the ones you expect to use the most seems like a good start to choose where to begin with.

Specific inference data methods

This category requires a much more detailed and custom implementation. Some examples that would fall here are:

• InferenceData html repr
• InferenceData method to print all dims, coords and variables in all groups at once. Maybe add option to show values? In general no values to make it readable. In jupyter may not have much sense because html could already cover this, but for terminal-like environments it would.
• Extension to xr.where to select from one group with a condition on another, ideally similar to pandas query function
• InferenceData compatible apply_ufunc to apply the same transform to several groups, e.g. shift and rescale all values in prior_predictive, posterior_predictive, predictions and observed_data (could also be an extension to map but apply_ufunc should be more versatile)

[percygautam]

Regarding specific inference data method : InferenceData html repr

I have been working on this for some time. The possible implementation is https://dfm.io/nbview/?url=https%3A%2F%2Fraw.githubusercontent.com%2Fpercygautam%2Farviz-examples%2Fmaster%2FHTML%2520repr.ipynb . I used xarray's implementation as reference.

[ahartikainen]

Looks great.

Do you think we should Add any extra information there? Like number of variables in each group? (not sure if this is a good idea or not)

[OriolAbril]

This got me thinking, should inference data objects have a name attribute?

[percygautam]

Looks great.

Do you think we should Add any extra information there? Like number of variables in each group? (not sure if this is a good idea or not)

Maybe, we could add dimensions for each group (when not checked). But personally, I like the minimalistic repr better.

[OriolAbril]

more ideas:

• a context so that all idata.fun inside it take the same groups as input (i.e. if groups=None, check context, if ouside context use all groups, otherwise use groups in context)? should not be too hard to implement, I think modifying _group_names should be enough
• maybe some functions will need both groups and var_names input? i.e. if we were to extend astype, not all variables in a groups need to have the same dtype.