Repository: arviz-devs/arviz

Issue:
Add helper function to easily stack `chain` and `draws` dimensions
Stacking chains and draws is often useful when one doesn't care about which chain a draw is coming from. This is currently possible by doing `idata.posterior.stack(sample=("chain", "draw"))`, but very few people seem to know that.

Adding a helper function to easily stack the `chain` and `draws` dimensions into a `sample` dimension would go a long way towards making users' life easier, as well as making better use of `xarray`'s capabilities. Feel free to comment if you want to take on this issue 🖖 


Discussion:
Is there a reason we can not add this when we create InferenceData? Would anything change or break? I think it would only add one more dimension, but everything else is same?
I don't think anything would break, as you can still access `chain` and `draw` after stacking, but I'm not sure I'd add it by default -- if another easy solution to add this feature exists I mean.
The reason is just that it adds yet another dimension, so it makes the ID heavier and, most importantly, can be confusing for people who don't need this. Generating it on-demand would be the best case scenario IMO 
Proposal for this, a `get_dataset` or `extract_dataset` function, something like:

```python
def get_dataset(idata, group="posterior", combined=False, var_names=None, filter_vars=None):
    """Extracts an inference data group or subset of it as xarray dataset

    Parameters
    ----------
    idata : InferenceData
        InferenceData from which to extract the data. 
        <not sure if it should be idata or anything that can be converted to idata>
    group : str, default "posterior"
    combined : bool, default False?
    var_names : str or list of str, optional
        Like with plotting, sometimes it's easier to subset saying what to exclude instead of what to include
    filter_vars : like with plotting

    Returns
    -------
    xarray.Dataset (or xarray.DataArray?)
        I am not sure whether we should return a dataarray iff `var_names` is a string 
        and a dataset otherwise or always a dataset.
    """
```

I believe this will handle most practical cases and be quite flexible while still being very little code as everything is reused from other functions/externalized. 
Hi! I recently submitted a PR for updating [GLM poisson regression](https://github.com/pymc-devs/pymc-examples/pull/154) to best practices, and on executing cell 18 in that notebook 

``` az.summary(np.exp(inf_fish_alt.posterior), kind="stats")```

I get this RuntimeWarning:
```
/home/ada/.local/lib/python3.8/site-packages/xarray/core/computation.py:724: RuntimeWarning: overflow encountered in exp 
  result_data = func(*input_data) 
/home/ada/.local/lib/python3.8/site-packages/numpy/core/_methods.py:193: RuntimeWarning: invalid value encountered in subtract 
  x = asanyarray(arr - arrmean) 
```
Here, `inf_fish_alt` is the trace created using model based on `pymc.glm.GLM.from_formula`

@OriolAbril suggested that this might be relevant to the discussion here.
yeah, here we basically want to exponentiate all varibles _except_ mu because `mu` has alredy been exponentiated in the model code and doing so again results in overflow. If the proposed funciton were available, we'd be able to get the right subset quite easily and then exponentiate this and pass it to summary.

Resolving patch:
diff --git a/arviz/data/__init__.py b/arviz/data/__init__.py
--- a/arviz/data/__init__.py
+++ b/arviz/data/__init__.py
@@ -14,6 +14,7 @@
 from .io_pymc3 import from_pymc3, from_pymc3_predictions
 from .io_pyro import from_pyro
 from .io_pystan import from_pystan
+from .utils import extract_dataset
 
 __all__ = [
     "InferenceData",
@@ -22,6 +23,7 @@
     "list_datasets",
     "clear_data_home",
     "numpy_to_data_array",
+    "extract_dataset",
     "dict_to_dataset",
     "convert_to_dataset",
     "convert_to_inference_data",
diff --git a/arviz/data/utils.py b/arviz/data/utils.py
new file mode 100644
--- /dev/null
+++ b/arviz/data/utils.py
@@ -0,0 +1,101 @@
+"""Data specific utilities."""
+import numpy as np
+
+from ..utils import _var_names
+from .converters import convert_to_dataset
+
+
+def extract_dataset(
+    data,
+    group="posterior",
+    combined=True,
+    var_names=None,
+    filter_vars=None,
+    num_samples=None,
+    rng=None,
+):
+    """Extract an InferenceData group or subset of it as a :class:`xarray.Dataset`.
+
+    Parameters
+    ----------
+    idata : InferenceData or InferenceData_like
+        InferenceData from which to extract the data.
+    group : str, optional
+        Which InferenceData data group to extract data from.
+    combined : bool, optional
+        Combine ``chain`` and ``draw`` dimensions into ``sample``. Won't work if
+        a dimension named ``sample`` already exists.
+    var_names : str or list of str, optional
+        Variables to be plotted, two variables are required. Prefix the variables by `~`
+        when you want to exclude them from the plot.
+    filter_vars: {None, "like", "regex"}, optional
+        If `None` (default), interpret var_names as the real variables names. If "like",
+        interpret var_names as substrings of the real variables names. If "regex",
+        interpret var_names as regular expressions on the real variables names. A la
+        `pandas.filter`.
+        Like with plotting, sometimes it's easier to subset saying what to exclude
+        instead of what to include
+    num_samples : int, optional
+        Extract only a subset of the samples. Only valid if ``combined=True``
+    rng : bool, int, numpy.Generator, optional
+        Shuffle the samples, only valid if ``combined=True``. By default,
+        samples are shuffled if ``num_samples`` is not ``None``, and are left
+        in the same order otherwise. This ensures that subsetting the samples doesn't return
+        only samples from a single chain and consecutive draws.
+
+    Returns
+    -------
+    xarray.Dataset
+
+    Examples
+    --------
+    The default behaviour is to return the posterior group after stacking the chain and
+    draw dimensions.
+
+    .. jupyter-execute::
+
+        import arviz as az
+        idata = az.load_arviz_data("centered_eight")
+        az.extract_dataset(idata)
+
+    You can also indicate a subset to be returned, but in variables and in samples:
+
+    .. jupyter-execute::
+
+        az.extract_dataset(idata, var_names="theta", num_samples=100)
+
+    To keep the chain and draw dimensions, use ``combined=False``.
+
+    .. jupyter-execute::
+
+        az.extract_dataset(idata, group="prior", combined=False)
+
+    """
+    if num_samples is not None and not combined:
+        raise ValueError("num_samples is only compatible with combined=True")
+    if rng is None:
+        rng = num_samples is not None
+    if rng is not False and not combined:
+        raise ValueError("rng is only compatible with combined=True")
+    data = convert_to_dataset(data, group=group)
+    var_names = _var_names(var_names, data, filter_vars)
+    if var_names is not None:
+        data = data[var_names]
+    if combined:
+        data = data.stack(sample=("chain", "draw"))
+    # 0 is a valid seed se we need to check for rng being exactly boolean
+    if rng is not False:
+        if rng is True:
+            rng = np.random.default_rng()
+        # default_rng takes ints or sequences of ints
+        try:
+            rng = np.random.default_rng(rng)
+            random_subset = rng.permutation(np.arange(len(data["sample"])))
+        except TypeError as err:
+            raise TypeError("Unable to initializate numpy random Generator from rng") from err
+        except AttributeError as err:
+            raise AttributeError("Unable to use rng to generate a permutation") from err
+        data = data.isel(sample=random_subset)
+    if num_samples is not None:
+        data = data.isel(sample=slice(None, num_samples))
+    return data
diff --git a/doc/source/conf.py b/doc/source/conf.py
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -62,6 +62,7 @@
     "notfound.extension",
     "sphinx_copybutton",
     "sphinx_codeautolink",
+    "jupyter_sphinx",
 ]
 
 # codeautolink


Test patch:
diff --git a/arviz/tests/base_tests/test_data.py b/arviz/tests/base_tests/test_data.py
--- a/arviz/tests/base_tests/test_data.py
+++ b/arviz/tests/base_tests/test_data.py
@@ -26,6 +26,7 @@
     list_datasets,
     load_arviz_data,
     to_netcdf,
+    extract_dataset,
 )
 
 from ...data.base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs
@@ -1397,3 +1398,38 @@ def test_nd_to_inference_data(self):
         assert inference_data.prior.chain.shape == shape[:1]
         assert inference_data.prior.draw.shape == shape[1:2]
         assert inference_data.prior[var_name].shape == shape
+
+
+class TestExtractDataset:
+    def test_default(self):
+        idata = load_arviz_data("centered_eight")
+        post = extract_dataset(idata)
+        assert isinstance(post, xr.Dataset)
+        assert "sample" in post.dims
+        assert post.theta.size == (4 * 500 * 8)
+
+    def test_seed(self):
+        idata = load_arviz_data("centered_eight")
+        post = extract_dataset(idata, rng=7)
+        post_pred = extract_dataset(idata, group="posterior_predictive", rng=7)
+        assert all(post.sample == post_pred.sample)
+
+    def test_no_combine(self):
+        idata = load_arviz_data("centered_eight")
+        post = extract_dataset(idata, combined=False)
+        assert "sample" not in post.dims
+        assert post.dims["chain"] == 4
+        assert post.dims["draw"] == 500
+
+    def test_var_name_group(self):
+        idata = load_arviz_data("centered_eight")
+        prior = extract_dataset(idata, group="prior", var_names="the", filter_vars="like")
+        assert prior.attrs == idata.prior.attrs
+        assert "theta" in prior.data_vars
+        assert "mu" not in prior.data_vars
+
+    def test_subset_samples(self):
+        idata = load_arviz_data("centered_eight")
+        post = extract_dataset(idata, num_samples=10)
+        assert post.dims["sample"] == 10
+        assert post.attrs == idata.posterior.attrs