Repository: arviz-devs/arviz

Issue:
az.loo fails when "sample" already exists as a dimension
**Describe the bug**

`az.loo` fails if there is a dimension named `sample` in the `log_likelihood` data variable. It looks like the `log_likelihood.stack(sample=("chain", "draw"))` line is the culprit.

**To Reproduce**

```python
import arviz as az
data = az.load_arviz_data("centered_eight")
data.sample_stats = data.sample_stats.rename({"school": "sample"})
az.loo(data)
```

<details>
  <summary>Error</summary>
  
  ```
Traceback (most recent call last):
  File "tmp/loo_issue.py", line 5, in <module>
    az.loo(data)
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/arviz/stats/stats.py", line 618, in loo
    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/xarray/core/dataarray.py", line 2035, in stack
    ds = self._to_temp_dataset().stack(dimensions, **dimensions_kwargs)
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/xarray/core/dataset.py", line 3578, in stack
    result = result._stack_once(dims, new_dim)
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/xarray/core/dataset.py", line 3524, in _stack_once
    stacked_var = exp_var.stack(**{new_dim: dims})
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/xarray/core/variable.py", line 1554, in stack
    result = result._stack_once(dims, new_dim)
  File "/Users/gibs/miniconda3/envs/birdman/lib/python3.7/site-packages/xarray/core/variable.py", line 1507, in _stack_once
    "cannot create a new dimension with the same "
ValueError: cannot create a new dimension with the same name as an existing dimension

  ```
</details>

**Expected behavior**

I expected that the computation would proceed as in the [documentation page](https://arviz-devs.github.io/arviz/api/generated/arviz.loo.html):

```
Computed from 2000 by 8 log-likelihood matrix

         Estimate       SE
elpd_loo   -30.81     1.43
p_loo        0.95        -
```

**Additional context**

This was with `arviz` commit 5957527 installed through `pip install git+git://github.com/arviz-devs/arviz.git` and `xarray` version `0.17.0`.

---

I realize that I can just change the name of my dimension from `sample` to something else to address this. However, I think it would be nice to not have to rename things, especially as the temporary dataset is not returned at all.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=h1) Report
> Merging [#1647](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=desc) (993ac43) into [main](https://codecov.io/gh/arviz-devs/arviz/commit/8a1fd2ce88ee974828a14de0d20e58fec322233a?el=desc) (8a1fd2c) will **decrease** coverage by `0.08%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/arviz-devs/arviz/pull/1647/graphs/tree.svg?width=650&height=150&src=pr&token=KZTdQCEf7t)](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##             main    #1647      +/-   ##
==========================================
- Coverage   90.92%   90.83%   -0.09%     
==========================================
  Files         108      108              
  Lines       11666    11720      +54     
==========================================
+ Hits        10607    10646      +39     
- Misses       1059     1074      +15     
```


| [Impacted Files](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [arviz/stats/stats.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovc3RhdHMvc3RhdHMucHk=) | `96.60% <100.00%> (-0.07%)` | :arrow_down: |
| [arviz/plots/energyplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovcGxvdHMvZW5lcmd5cGxvdC5weQ==) | `83.33% <0.00%> (-16.67%)` | :arrow_down: |
| [arviz/data/io\_pyro.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovZGF0YS9pb19weXJvLnB5) | `96.24% <0.00%> (-1.12%)` | :arrow_down: |
| [arviz/data/io\_numpyro.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovZGF0YS9pb19udW1weXJvLnB5) | `93.89% <0.00%> (-0.60%)` | :arrow_down: |
| [arviz/data/io\_pymc3.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovZGF0YS9pb19weW1jMy5weQ==) | `90.20% <0.00%> (-0.56%)` | :arrow_down: |
| [arviz/data/inference\_data.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovZGF0YS9pbmZlcmVuY2VfZGF0YS5weQ==) | `83.57% <0.00%> (-0.03%)` | :arrow_down: |
| [arviz/plots/hdiplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovcGxvdHMvaGRpcGxvdC5weQ==) | `92.45% <0.00%> (ø)` | |
| [arviz/plots/distplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovcGxvdHMvZGlzdHBsb3QucHk=) | `90.47% <0.00%> (ø)` | |
| [arviz/stats/stats\_utils.py](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree#diff-YXJ2aXovc3RhdHMvc3RhdHNfdXRpbHMucHk=) | `96.71% <0.00%> (ø)` | |
| ... and [8 more](https://codecov.io/gh/arviz-devs/arviz/pull/1647/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=footer). Last update [8a1fd2c...993ac43](https://codecov.io/gh/arviz-devs/arviz/pull/1647?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

MgeeeeK: > can you also add a note in the docstring of psislw about this convention?

Okay :+1:

Should I also reflect these changes in the examples? Though, I think the note will do the job but still what do you say...
OriolAbril: Yes,  updating the example too, good catch, I had missed that 
MgeeeeK: For some reason I'm getting this error in the build doc pipeline:

```bash
Exception occurred:
  File "/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/docutils/parsers/rst/__init__.py", line 196, in parse
    if len(line) > self.document.settings.line_length_limit:
AttributeError: 'Values' object has no attribute 'line_length_limit'
 ```
ahartikainen: > For some reason I'm getting this error in the build doc pipeline:
> 
> ```shell
> Exception occurred:
>   File "/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/docutils/parsers/rst/__init__.py", line 196, in parse
>     if len(line) > self.document.settings.line_length_limit:
> AttributeError: 'Values' object has no attribute 'line_length_limit'
> ```

Probably docutils has been updated (docutils==0.17, 2021-04-03 https://pypi.org/project/docutils/#history )

Working workflow has docutils==0.16 https://dev.azure.com/ArviZ/ArviZ/_build/results?buildId=4267&view=logs&j=5e60a62b-366f-55f7-ec2d-5209191ccb17&t=3b19a6af-f535-5ba4-1e9a-379fb7d07d81

I will create a quick fix PR and then we can create a correct fix later.
MgeeeeK: Hey @OriolAbril I've update the docstring and rebased the branch, I think its ready now.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 67d44fd..56af478 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@
 * Upload updated mypy.ini in ci if mypy copilot fails ([1624](https://github.com/arviz-devs/arviz/pull/1624))
 * Added type checking to raise an error whenever `InferenceData` object is passed using `io_pymc3`'s `trace` argument ([1629](https://github.com/arviz-devs/arviz/pull/1629))
 * Fix `xlabels` in `plot_elpd` ([1601](https://github.com/arviz-devs/arviz/pull/1601))
+* Renamed `sample` dim to `__sample__` when stacking `chain` and `draw` to avoid dimension collision ([1647](https://github.com/arviz-devs/arviz/pull/1647))
 
 ### Deprecation
 * Deprecated `index_origin` and `order` arguments in `az.summary` ([1201](https://github.com/arviz-devs/arviz/pull/1201))
diff --git a/arviz/stats/stats.py b/arviz/stats/stats.py
index f3fadb3..5467e34 100644
--- a/arviz/stats/stats.py
+++ b/arviz/stats/stats.py
@@ -648,7 +648,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     log_likelihood = _get_log_likelihood(inference_data, var_name=var_name)
     pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise
 
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     shape = log_likelihood.shape
     n_samples = shape[-1]
     n_data_points = np.product(shape[:-1])
@@ -692,7 +692,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
         warn_mg = True
 
     ufunc_kwargs = {"n_dims": 1, "ravel": False}
-    kwargs = {"input_core_dims": [["sample"]]}
+    kwargs = {"input_core_dims": [["__sample__"]]}
     loo_lppd_i = scale_value * _wrap_xarray_ufunc(
         _logsumexp, log_weights, ufunc_kwargs=ufunc_kwargs, **kwargs
     )
@@ -752,6 +752,14 @@ def psislw(log_weights, reff=1.0):
     """
     Pareto smoothed importance sampling (PSIS).
 
+    Notes
+    -----
+    If the ``log_weights`` input is an :class:`~xarray.DataArray` with a dimension
+    named ``__sample__`` (recommended) ``psislw`` will interpret this dimension as samples,
+    and all other dimensions as dimensions of the observed data, looping over them to
+    calculate the psislw of each observation. If no ``__sample__`` dimension is present or
+    the input is a numpy array, the last dimension will be interpreted as ``__sample__``.
+
     Parameters
     ----------
     log_weights: array
@@ -778,13 +786,17 @@ def psislw(log_weights, reff=1.0):
 
         In [1]: import arviz as az
            ...: data = az.load_arviz_data("centered_eight")
-           ...: log_likelihood = data.sample_stats.log_likelihood.stack(sample=("chain", "draw"))
+           ...: log_likelihood = data.sample_stats.log_likelihood.stack(
+           ...:     __sample__=("chain", "draw")
+           ...: )
            ...: az.psislw(-log_likelihood, reff=0.8)
 
     """
-    if hasattr(log_weights, "sample"):
-        n_samples = len(log_weights.sample)
-        shape = [size for size, dim in zip(log_weights.shape, log_weights.dims) if dim != "sample"]
+    if hasattr(log_weights, "__sample__"):
+        n_samples = len(log_weights.__sample__)
+        shape = [
+            size for size, dim in zip(log_weights.shape, log_weights.dims) if dim != "__sample__"
+        ]
     else:
         n_samples = log_weights.shape[-1]
         shape = log_weights.shape[:-1]
@@ -799,7 +811,7 @@ def psislw(log_weights, reff=1.0):
     # define kwargs
     func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "k_min": k_min, "out": out}
     ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
-    kwargs = {"input_core_dims": [["sample"]], "output_core_dims": [["sample"], []]}
+    kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
     log_weights, pareto_shape = _wrap_xarray_ufunc(
         _psislw,
         log_weights,
@@ -808,7 +820,7 @@ def psislw(log_weights, reff=1.0):
         **kwargs,
     )
     if isinstance(log_weights, xr.DataArray):
-        log_weights = log_weights.rename("log_weights").rename(sample="sample")
+        log_weights = log_weights.rename("log_weights")
     if isinstance(pareto_shape, xr.DataArray):
         pareto_shape = pareto_shape.rename("pareto_shape")
     return log_weights, pareto_shape
@@ -1415,13 +1427,13 @@ def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
     else:
         raise TypeError('Valid scale values are "deviance", "log", "negative_log"')
 
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     shape = log_likelihood.shape
     n_samples = shape[-1]
     n_data_points = np.product(shape[:-1])
 
     ufunc_kwargs = {"n_dims": 1, "ravel": False}
-    kwargs = {"input_core_dims": [["sample"]]}
+    kwargs = {"input_core_dims": [["__sample__"]]}
     lppd_i = _wrap_xarray_ufunc(
         _logsumexp,
         log_likelihood,
@@ -1431,7 +1443,7 @@ def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
         **kwargs,
     )
 
-    vars_lpd = log_likelihood.var(dim="sample")
+    vars_lpd = log_likelihood.var(dim="__sample__")
     warn_mg = False
     if np.any(vars_lpd > 0.4):
         warnings.warn(
@@ -1535,7 +1547,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
 
         In [1]: T = data.observed_data.obs - data.posterior.mu.median(dim=("chain", "draw"))
            ...: T_hat = data.posterior_predictive.obs - data.posterior.mu
-           ...: T_hat = T_hat.stack(sample=("chain", "draw"))
+           ...: T_hat = T_hat.stack(__sample__=("chain", "draw"))
            ...: az.loo_pit(idata=data, y=T**2, y_hat=T_hat**2)
 
     """
@@ -1561,7 +1573,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
         elif not isinstance(y, (np.ndarray, xr.DataArray)):
             raise ValueError(f"y must be of types array, DataArray or str, not {type(y)}")
         if isinstance(y_hat, str):
-            y_hat = idata.posterior_predictive[y_hat].stack(sample=("chain", "draw")).values
+            y_hat = idata.posterior_predictive[y_hat].stack(__sample__=("chain", "draw")).values
         elif not isinstance(y_hat, (np.ndarray, xr.DataArray)):
             raise ValueError(f"y_hat must be of types array, DataArray or str, not {type(y_hat)}")
         if log_weights is None:
@@ -1572,10 +1584,10 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
                     log_likelihood = _get_log_likelihood(idata)
             else:
                 log_likelihood = _get_log_likelihood(idata)
-            log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+            log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
             posterior = convert_to_dataset(idata, group="posterior")
             n_chains = len(posterior.chain)
-            n_samples = len(log_likelihood.sample)
+            n_samples = len(log_likelihood.__sample__)
             ess_p = ess(posterior, method="mean")
             # this mean is over all data variables
             reff = (
@@ -1608,7 +1620,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
         )
 
     kwargs = {
-        "input_core_dims": [[], ["sample"], ["sample"]],
+        "input_core_dims": [[], ["__sample__"], ["__sample__"]],
         "output_core_dims": [[]],
         "join": "left",
     }


Test patch:
diff --git a/arviz/tests/base_tests/test_plots_bokeh.py b/arviz/tests/base_tests/test_plots_bokeh.py
index 396c804..41279a0 100644
--- a/arviz/tests/base_tests/test_plots_bokeh.py
+++ b/arviz/tests/base_tests/test_plots_bokeh.py
@@ -718,9 +718,9 @@ def test_plot_loo_pit_label(models, args):
     if args.get("y_hat") == "str":
         y_hat = "y"
     elif args.get("y_hat") == "DataArray":
-        y_hat = models.model_1.posterior_predictive.y.stack(sample=("chain", "draw"))
+        y_hat = models.model_1.posterior_predictive.y.stack(__sample__=("chain", "draw"))
     elif args.get("y_hat") == "ndarray":
-        y_hat = models.model_1.posterior_predictive.y.stack(sample=("chain", "draw")).values
+        y_hat = models.model_1.posterior_predictive.y.stack(__sample__=("chain", "draw")).values
     else:
         y_hat = None
 
diff --git a/arviz/tests/base_tests/test_stats.py b/arviz/tests/base_tests/test_stats.py
index 4f737d5..260fe30 100644
--- a/arviz/tests/base_tests/test_stats.py
+++ b/arviz/tests/base_tests/test_stats.py
@@ -486,7 +486,7 @@ def test_loo_print(centered_eight, scale):
 def test_psislw(centered_eight):
     pareto_k = loo(centered_eight, pointwise=True, reff=0.7)["pareto_k"]
     log_likelihood = get_log_likelihood(centered_eight)
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     assert_allclose(pareto_k, psislw(-log_likelihood, 0.7)[1])
 
 
@@ -544,9 +544,9 @@ def test_loo_pit(centered_eight, args):
     y_hat = args.get("y_hat", None)
     log_weights = args.get("log_weights", None)
     y_arr = centered_eight.observed_data.obs
-    y_hat_arr = centered_eight.posterior_predictive.obs.stack(sample=("chain", "draw"))
-    log_like = get_log_likelihood(centered_eight).stack(sample=("chain", "draw"))
-    n_samples = len(log_like.sample)
+    y_hat_arr = centered_eight.posterior_predictive.obs.stack(__sample__=("chain", "draw"))
+    log_like = get_log_likelihood(centered_eight).stack(__sample__=("chain", "draw"))
+    n_samples = len(log_like.__sample__)
     ess_p = ess(centered_eight.posterior, method="mean")
     reff = np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples
     log_weights_arr = psislw(-log_like, reff=reff)[0]
@@ -584,9 +584,9 @@ def test_loo_pit_multidim(multidim_models, args):
     log_weights = args.get("log_weights", None)
     idata = multidim_models.model_1
     y_arr = idata.observed_data.y
-    y_hat_arr = idata.posterior_predictive.y.stack(sample=("chain", "draw"))
-    log_like = get_log_likelihood(idata).stack(sample=("chain", "draw"))
-    n_samples = len(log_like.sample)
+    y_hat_arr = idata.posterior_predictive.y.stack(__sample__=("chain", "draw"))
+    log_like = get_log_likelihood(idata).stack(__sample__=("chain", "draw"))
+    n_samples = len(log_like.__sample__)
     ess_p = ess(idata.posterior, method="mean")
     reff = np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples
     log_weights_arr = psislw(-log_like, reff=reff)[0]