Repository: Ouranosinc/xclim

Issue:
Ensemble percentiles as dimension instead of variables
<!--Describe what you were trying to get done.
Tell us what happened, what went wrong, and what you expected to happen.-->

Currently, `xclim.ensembles.ensemble_percentiles` takes in a dataset and returns a new one with one variable per percentile and data variable. A more "xarray-like" behaviour would be to replace the `realization` dimension by the `percentile` dimension. In fact, `ensemble_percentiles` is more or less a xr-wrapper and workaround because `xr.quantiles` doesn't work on dask arrays. 

Moreover, it would be nice if `ensemble_percentiles` accepted DataArrays directly.

Discussion points:
- Should we keep the old behaviour (creating new variables)?
- If yes, should it be the default?

Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 5c6413eb..47eda3f1 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,7 @@ History
 =======
 0.18.x
 ------
+* `ensembles.ensemble_percentiles` modified to compute along a `percentiles` dimension by default, instead of creating different variables.
 * Added indicator `first_day_below` and run length helper `first_run_after_date`
 * Added ANUCLIM model climate indices mappings.
 
diff --git a/docs/notebooks/ensembles.ipynb b/docs/notebooks/ensembles.ipynb
index ce96f320..c7b3d07d 100644
--- a/docs/notebooks/ensembles.ipynb
+++ b/docs/notebooks/ensembles.ipynb
@@ -136,7 +136,7 @@
     "### Ensemble percentiles\n",
     "\n",
     "Here we use xclim's `ensemble_percentiles()` to calculate percentile values across the 10 realizations. \n",
-    "Output variables are created for combining the original variable name `tas` with the addtional ending `_p{x}` where x are the input percentile values `default = [10, 50, 90]`. Compared to numpy's `percentile()` and xarray's `quantile()`, this method handles more efficiently dataset with invalid values and the chunking along the realization dimension (which is automatic when dask arrays are used)."
+    "The output has now a `percentiles` dimension instead of `realization`. Split variables can be created instead, by specifying `split=True` (the variable name `tas` will be appended with `_p{x}`). Compared to numpy's `percentile()` and xarray's `quantile()`, this method handles more efficiently dataset with invalid values and the chunking along the realization dimension (which is automatic when dask arrays are used)."
    ]
   },
   {
@@ -157,13 +157,20 @@
    "source": [
     "fig, ax = plt.subplots()\n",
     "ax.fill_between(ens_stats.time.values, ens_stats.tas_min, ens_stats.tas_max, alpha=0.3, label='Min-Max')\n",
-    "ax.fill_between(ens_perc.time.values, ens_perc.tas_p15, ens_perc.tas_p85, alpha=0.5, label='Perc. 15-85')\n",
+    "ax.fill_between(ens_perc.time.values, ens_perc.tas.sel(percentiles=15), ens_perc.tas.sel(percentiles=85), alpha=0.5, label='Perc. 15-85')\n",
     "ax._get_lines.get_next_color()  # Hack to get different line\n",
     "ax._get_lines.get_next_color()\n",
     "ax.plot(ens_stats.time.values, ens_stats.tas_mean, linewidth=2, label='Mean')\n",
-    "ax.plot(ens_perc.time.values, ens_perc.tas_p50, linewidth=2, label='Median')\n",
+    "ax.plot(ens_perc.time.values, ens_perc.tas.sel(percentiles=50), linewidth=2, label='Median')\n",
     "ax.legend()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -182,7 +189,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.0"
+   "version": "3.8.1"
   }
  },
  "nbformat": 4,
diff --git a/xclim/ensembles.py b/xclim/ensembles.py
index d9aea683..5d11e0db 100644
--- a/xclim/ensembles.py
+++ b/xclim/ensembles.py
@@ -14,6 +14,7 @@ from sklearn.cluster import KMeans
 
 from xclim.core.calendar import convert_calendar
 from xclim.core.calendar import get_calendar
+from xclim.core.formatting import merge_attributes
 from xclim.core.formatting import update_history
 
 # Avoid having to include matplotlib in xclim requirements
@@ -150,9 +151,10 @@ def ensemble_mean_std_max_min(ens: xr.Dataset) -> xr.Dataset:
 
 
 def ensemble_percentiles(
-    ens: xr.Dataset,
-    values: Tuple[int, int, int] = (10, 50, 90),
+    ens: Union[xr.Dataset, xr.DataArray],
+    values: Tuple[int] = (10, 50, 90),
     keep_chunk_size: Optional[bool] = None,
+    split: bool = False,
 ) -> xr.Dataset:
     """Calculate ensemble statistics between a results from an ensemble of climate simulations.
 
@@ -160,20 +162,24 @@ def ensemble_percentiles(
 
     Parameters
     ----------
-    ens: xr.Dataset
-      Ensemble dataset (see xclim.ensembles.create_ensemble).
-    values : Tuple[int, int, int]
+    ens: Union[xr.Dataset, xr.DataArray]
+      Ensemble dataset or dataarray (see xclim.ensembles.create_ensemble).
+    values : Tuple[int]
       Percentile values to calculate. Default: (10, 50, 90).
     keep_chunk_size : Optional[bool]
       For ensembles using dask arrays, all chunks along the 'realization' axis are merged.
       If True, the dataset is rechunked along the dimension with the largest chunks, so that the chunks keep the same size (approx)
       If False, no shrinking is performed, resulting in much larger chunks
       If not defined, the function decides which is best
+    split : bool
+      Whether to split each percentile into a new variable of concatenate the ouput along a new
+      "percentiles" dimension.
 
     Returns
     -------
-    xr.Dataset
-      Dataset with containing data variables of requested ensemble statistics
+    Union[xr.Dataset, xr.DataArray]
+      If split is True, same type as ens; dataset otherwise,
+      containing data variable(s) of requested ensemble statistics
 
     Examples
     --------
@@ -190,58 +196,77 @@ def ensemble_percentiles(
     >>> ens_percs = ensemble_percentiles(ens, keep_chunk_size=False)
     """
 
-    ds_out = xr.Dataset(attrs=ens.attrs)
-    for v in ens.data_vars:
-        # Percentile calculation forbids any chunks along realization
-        if len(ens.chunks.get("realization", [])) > 1:
-            if keep_chunk_size is None:
-                # Enable smart rechunking is chunksize exceed 2E8 elements after merging along realization
-                keep_chunk_size = (
-                    np.prod(ens[v].isel(realization=0).data.chunksize)
-                    * ens.realization.size
-                    > 2e8
-                )
-            if keep_chunk_size:
-                # Smart rechunk on dimension where chunks are the largest
-                chkDim, chks = max(
-                    ens.chunks.items(),
-                    key=lambda kv: 0 if kv[0] == "realization" else max(kv[1]),
-                )
-                var = ens[v].chunk(
-                    {"realization": -1, chkDim: len(chks) * ens.realization.size}
+    if isinstance(ens, xr.Dataset):
+        out = xr.merge(
+            [
+                ensemble_percentiles(
+                    da, values, keep_chunk_size=keep_chunk_size, split=split
                 )
-            else:
-                var = ens[v].chunk({"realization": -1})
-        else:
-            var = ens[v]
-
-        for p in values:
-            perc = xr.apply_ufunc(
-                _calc_perc,
-                var,
-                input_core_dims=[["realization"]],
-                output_core_dims=[[]],
-                keep_attrs=True,
-                kwargs=dict(p=p),
-                dask="parallelized",
-                output_dtypes=[ens[v].dtype],
+                for da in ens.data_vars.values()
+                if "realization" in da.dims
+            ]
+        )
+        out.attrs.update(ens.attrs)
+        out.attrs["history"] = update_history(
+            f"Computation of the percentiles on {ens.realization.size} ensemble members.",
+            ens,
+        )
+
+        return out
+
+    # Percentile calculation forbids any chunks along realization
+    if ens.chunks and len(ens.chunks[ens.get_axis_num("realization")]) > 1:
+        if keep_chunk_size is None:
+            # Enable smart rechunking is chunksize exceed 2E8 elements after merging along realization
+            keep_chunk_size = (
+                np.prod(ens.isel(realization=0).data.chunksize) * ens.realization.size
+                > 2e8
             )
+        if keep_chunk_size:
+            # Smart rechunk on dimension where chunks are the largest
+            chkDim, chks = max(
+                enumerate(ens.chunks),
+                key=lambda kv: 0
+                if kv[0] == ens.get_axis_num("realization")
+                else max(kv[1]),
+            )
+            ens = ens.chunk(
+                {"realization": -1, ens.dims[chkDim]: len(chks) * ens.realization.size}
+            )
+        else:
+            ens = ens.chunk({"realization": -1})
+
+    out = xr.apply_ufunc(
+        _calc_perc,
+        ens,
+        input_core_dims=[["realization"]],
+        output_core_dims=[["percentiles"]],
+        keep_attrs=True,
+        kwargs=dict(p=values),
+        dask="parallelized",
+        output_dtypes=[ens.dtype],
+        output_sizes={"percentiles": len(values)},
+    )
 
-            perc.name = f"{v}_p{p:02d}"
-            ds_out[perc.name] = perc
+    out = out.assign_coords(
+        percentiles=xr.DataArray(list(values), dims=("percentiles",))
+    )
 
-            if "description" in ds_out[perc.name].attrs:
-                ds_out[perc.name].attrs[
-                    "description"
-                ] = f"{ds_out[perc.name].attrs['description']} : {p}th percentile of ensemble"
-            else:
-                ds_out[perc.name].attrs["description"] = f"{p}th percentile of ensemble"
+    if split:
+        out = out.to_dataset(dim="percentiles")
+        for p, perc in out.data_vars.items():
+            perc.attrs.update(ens.attrs)
+            perc.attrs["description"] = (
+                perc.attrs.get("descrption", "") + " {p}th percentile of ensemble."
+            )
+            out = out.rename(name_dict={p: f"{ens.name}_p{int(p):02d}"})
 
-    ds_out.attrs["history"] = update_history(
+    out.attrs["history"] = update_history(
         f"Computation of the percentiles on {ens.realization.size} ensemble members.",
-        ds_out,
+        ens,
     )
-    return ds_out
+
+    return out
 
 
 def _ens_align_datasets(
@@ -313,7 +338,7 @@ def _ens_align_datasets(
     return ds_all
 
 
-def _calc_perc(arr, p=50):
+def _calc_perc(arr, p=[50]):
     """Ufunc-like computing a percentile over the last axis of the array.
 
     Processes cases with invalid values separately, which makes it more efficent than np.nanpercentile for array with only a few invalid points.
@@ -330,16 +355,15 @@ def _calc_perc(arr, p=50):
     np.array
     """
     nan_count = np.isnan(arr).sum(axis=-1)
-    out = np.percentile(arr, p, axis=-1)
+    out = np.moveaxis(np.percentile(arr, p, axis=-1), 0, -1)
     nans = (nan_count > 0) & (nan_count < arr.shape[-1])
     if np.any(nans):
-        arr1 = arr.reshape(int(arr.size / arr.shape[-1]), arr.shape[-1])
+        out_mask = np.stack([nans] * len(p), axis=-1)
+        # arr1 = arr.reshape(int(arr.size / arr.shape[-1]), arr.shape[-1])
         # only use nanpercentile where we need it (slow performance compared to standard) :
-        nan_index = np.where(nans)
-        t = np.ravel_multi_index(nan_index, nan_count.shape)
-        out[np.unravel_index(t, nan_count.shape)] = np.nanpercentile(
-            arr1[t, :], p, axis=-1
-        )
+        out[out_mask] = np.moveaxis(
+            np.nanpercentile(arr[nans], p, axis=-1), 0, -1
+        ).ravel()
     return out
 
 


Test patch:
diff --git a/tests/test_ensembles.py b/tests/test_ensembles.py
index 5643d527..aaf40037 100644
--- a/tests/test_ensembles.py
+++ b/tests/test_ensembles.py
@@ -122,7 +122,8 @@ class TestEnsembleStats:
         ens = ensembles.create_ensemble(self.nc_files_simple)
         if transpose:
             ens = ens.transpose()
-        out1 = ensembles.ensemble_percentiles(ens)
+
+        out1 = ensembles.ensemble_percentiles(ens, split=True)
         np.testing.assert_array_equal(
             np.percentile(ens["tg_mean"].isel(time=0, lon=5, lat=5), 10),
             out1["tg_mean_p10"].isel(time=0, lon=5, lat=5),
@@ -137,27 +138,31 @@ class TestEnsembleStats:
         )
         assert np.all(out1["tg_mean_p90"] > out1["tg_mean_p50"])
         assert np.all(out1["tg_mean_p50"] > out1["tg_mean_p10"])
-        out1 = ensembles.ensemble_percentiles(ens, values=(25, 75))
-        assert np.all(out1["tg_mean_p75"] > out1["tg_mean_p25"])
+
+        out2 = ensembles.ensemble_percentiles(ens, values=(25, 75), split=True)
+        assert np.all(out2["tg_mean_p75"] > out2["tg_mean_p25"])
         assert "Computation of the percentiles on" in out1.attrs["history"]
 
+        out3 = ensembles.ensemble_percentiles(ens)
+        xr.testing.assert_equal(
+            out1["tg_mean_p10"], out3.tg_mean.sel(percentiles=10, drop=True)
+        )
+
     @pytest.mark.parametrize("keep_chunk_size", [False, True, None])
     def test_calc_perc_dask(self, keep_chunk_size):
         ens = ensembles.create_ensemble(self.nc_files_simple)
         out2 = ensembles.ensemble_percentiles(
-            ens.chunk({"time": 2}), values=(10, 50, 90), keep_chunk_size=keep_chunk_size
+            ens.chunk({"time": 2}), keep_chunk_size=keep_chunk_size
         )
         out1 = ensembles.ensemble_percentiles(ens.load())
-        np.testing.assert_array_equal(out1["tg_mean_p10"], out2["tg_mean_p10"])
-        np.testing.assert_array_equal(out1["tg_mean_p50"], out2["tg_mean_p50"])
-        np.testing.assert_array_equal(out1["tg_mean_p90"], out2["tg_mean_p90"])
+        np.testing.assert_array_equal(out1["tg_mean"], out2["tg_mean"])
 
     def test_calc_perc_nans(self):
         ens = ensembles.create_ensemble(self.nc_files_simple).load()
 
         ens.tg_mean[2, 0, 5, 5] = np.nan
         ens.tg_mean[2, 7, 5, 5] = np.nan
-        out1 = ensembles.ensemble_percentiles(ens)
+        out1 = ensembles.ensemble_percentiles(ens, split=True)
         np.testing.assert_array_equal(
             np.percentile(ens["tg_mean"][:, 0, 5, 5], 10), np.nan
         )