Repository: Ouranosinc/xclim

Issue:
Optimize/improve  xclim ensemble stats utilities
There are currently functional versions of the following utilities but they could use improvement / polishing:
```python
xclim.utils.utils.create_ensemble()
xclim.utils.ensemble_percentiles()
xclim.utils.calc_percentiles_simple()
xclim.utils.calc_percentiles_blocks()
xclim.utils.calc_perc()
```
**Ideas to explore:**

For ensemble creation :
-  create dimension using IndexVariable  (instead of just a string) and make sure that it has attribute axis='E'. 


**Try to achieve dask enabled percentile calculation (remove 'time_block' method):**

- try DataArray.reduce(np.nanpercentile, dim='realization') - Check performance (nanpercentile seems slow)

- use DataArray.apply(xclim.utils.calc_perc, dim='realization') - 'calc_perc' only applies nanpercentile where needed ... should technically be faster



Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 1461](https://coveralls.io/builds/27991654)

* **25** of **27**   **(92.59%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.2%**) to **90.777%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [xclim/ensembles.py](https://coveralls.io/builds/27991654/source?filename=xclim%2Fensembles.py#L196) | 25 | 27 | 92.59%
<!-- | **Total:** | **25** | **27** | **92.59%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/27991654/badge)](https://coveralls.io/builds/27991654) |
| :-- | --: |
| Change from base [Build 1460](https://coveralls.io/builds/27989588): |  -0.2% |
| Covered Lines: | 1683 |
| Relevant Lines: | 1854 |

---
##### 💛  - [Coveralls](https://coveralls.io)

huard: I suggest we release 0.13 once this is merged. 
tlogan2000: This looks cool but will need a bit of time to review.  @aulemahal I would maybe want to get you to benchmark using some of the climatedata.ca data (lots of gridcells (all of canada), 24 realizations 151 timesteps  ) ...  

Quick question for @aulemahal ... between the two new option in  https://github.com/Ouranosinc/xclim-benchmark/pull/2 (apply and apply rechunk) if I understand apply is the 'faster' option but can use more memory... anyway to combine the two with a rule of thumb chunksize? e.g. if chunksizes > n apply_rechunk else apply...
aulemahal: @tlogan2000  I generated 25 realizations of random data made of 4 chunks of 50x1000x500. I think this resembles  what you were thinking of? 

It gives the following:
![ensemble_percentile_memuse_ManyRealizations_FewLargeChunks_SmallTimeDim](https://user-images.githubusercontent.com/20629530/72107184-d2ec3980-32fe-11ea-92e6-0ad8002432f4.png)
Without rechunking, it behaves as expected, but the default rechunking fails (highest red line). I think this is due to the small size of the time dimensions and its chunks here (4x50). Thus, when rechunking, we get chunks of 2x1000x500, which seems not efficient. 

I added a test so that the rechunking is done on the dimensions with the largest chunks, that the lowest red line. It's incredibly better, dask is so mysterious... Coming up in the PRs.


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 2da52984..50dd6f92 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -10,6 +10,7 @@ History
 * Fixed `ensembles.create_ensemble` errors for builds against xarray master branch.
 * Reformatted code to make better use of Python3.6 conventions (f-strings and object signatures).
 * Fixed randomly failing tests of `checks.missing_any`.
+* Improvement of `ensemble.ensemble_percentile` and `ensemble.create_ensemble`.
 
 0.12.x-beta (2019-11-18)
 ------------------------
diff --git a/xclim/ensembles.py b/xclim/ensembles.py
index ddbccd9a..f333face 100644
--- a/xclim/ensembles.py
+++ b/xclim/ensembles.py
@@ -36,6 +36,7 @@ def create_ensemble(
     a new dimension (name:'realization'). In the case where input files have unequal time dimensions, the output
     ensemble Dataset is created for maximum time-step interval of all input files.  Before concatenation, datasets not
     covering the entire time span have their data padded with NaN values.
+    Dataset and variable attributes of the first dataset are copied to the resulting dataset.
 
     Parameters
     ----------
@@ -48,7 +49,7 @@ def create_ensemble(
       Only applicable when "datasets" is a sequence of file paths.
 
     xr_kwargs :
-      Any keyword arguments to be given to xarray when opening the files.
+      Any keyword arguments to be given to `xr.open_dataset` when opening the files (or to `xr.open_mfdataset` if mf_flag is True)
 
     Returns
     -------
@@ -74,19 +75,16 @@ def create_ensemble(
     >>> datasets.append(glob.glob('/dir2/*.nc'))
     >>> ens = ensembles.create_ensemble(datasets, mf_flag=True)
     """
-
-    dim = "realization"
-
     time_flag, time_all = _ens_checktimes(datasets, mf_flag, **xr_kwargs)
 
-    ds1 = _ens_align_datasets(datasets, mf_flag, time_flag, time_all)
+    ds = _ens_align_datasets(datasets, mf_flag, time_flag, time_all, **xr_kwargs)
 
-    for v in list(ds1[0].data_vars):
-        list1 = [ds[v] for ds in ds1]
-        data = xr.concat(list1, dim=dim)
-        if v == list(ds1[0].data_vars)[0]:
-            ens = xr.Dataset(data_vars=None, coords=data.coords, attrs=ds1[0].attrs)
-        ens[v] = data
+    dim = xr.IndexVariable("realization", np.arange(len(ds)), attrs={"axis": "E"})
+
+    ens = xr.concat(ds, dim)
+    for vname, var in ds[0].variables.items():
+        ens[vname].attrs.update(**var.attrs)
+    ens.attrs.update(**ds[0].attrs)
 
     return ens
 
@@ -142,12 +140,11 @@ 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),
-    time_block: Optional[int] = None,
+    keep_chunk_size: Optional[bool] = None,
 ) -> xr.Dataset:
     """Calculate ensemble statistics between a results from an ensemble of climate simulations.
 
-    Returns a Dataset containing ensemble statistics for input climate simulations.
-    Alternatively calculate ensemble percentiles (default) or ensemble mean and standard deviation.
+    Returns a Dataset containing ensemble percentiles for input climate simulations.
 
     Parameters
     ----------
@@ -155,9 +152,11 @@ def ensemble_percentiles(
       Ensemble dataset (see xclim.ensembles.create_ensemble).
     values : Tuple[int, int, int]
       Percentile values to calculate. Default: (10, 50, 90).
-    time_block : Optional[int]
-      For large ensembles, iteratively calculate percentiles in time-step blocks (n==time_block).
-      If not defined, the function tries to estimate an appropriate value.
+    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
 
     Returns
     -------
@@ -177,34 +176,58 @@ def ensemble_percentiles(
     Calculate non-default percentiles (25th and 75th)
     >>> ens_percs = ensembles.ensemble_percentiles(ens, values=(25, 50, 75))
     >>> print(ens_percs['tas_p25'])
-    Calculate by time blocks (n=10) if ensemble size is too large to load in memory
-    >>> ens_percs = ensembles.ensemble_percentiles(ens, time_block=10)
+    If the original array has many small chunks, it might be more efficient to do:
+    >>> ens_percs = ensembles.ensemble_percentiles(ens, keep_chunk_size=False)
     >>> print(ens_percs['tas_p25'])
     """
 
     ds_out = ens.drop_vars(names=set(ens.data_vars))
-    dims = list(ens.dims)
     for v in ens.data_vars:
-        # Percentile calculation requires load to memory : automate size for large ensemble objects
-        if not time_block:
-            time_block = round(
-                2e8 / (ens[v].size / ens[v].shape[dims.index("time")]), -1
-            )  # 2E8
-
-        if time_block > len(ens[v].time):
-            out = _calc_percentiles_simple(ens, v, values)
-
+        # 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,}
+                )
+            else:
+                var = ens[v].chunk({"realization": -1})
         else:
-            # loop through blocks
-            warnings.warn(
-                f"Large ensemble size detected: statistics will be"
-                f" calculated in blocks of {int(time_block)} time-steps.",
-                UserWarning,
-                stacklevel=2,
+            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],
             )
-            out = _calc_percentiles_blocks(ens, v, values, time_block)
-        for vv in out.data_vars:
-            ds_out[vv] = out[vv]
+
+            perc.name = f"{v}_p{p:02d}"
+            ds_out[perc.name] = perc
+
+            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"
+
     return ds_out
 
 
@@ -341,98 +364,33 @@ def _ens_align_datasets(
     return ds_all
 
 
-def _calc_percentiles_simple(ens, v, values):
-    ds_out = ens.drop_vars(names=set(ens.data_vars))
-    dims = list(ens[v].dims)
-    outdims = [x for x in dims if "realization" not in x]
-
-    # print('loading ensemble data to memory')
-    arr = ens[v].load()  # percentile calc requires loading the array
-    coords = {}
-    for c in outdims:
-        coords[c] = arr[c]
-    for p in values:
-        outvar = v + "_p" + str(p)
-
-        out1 = _calc_perc(arr, p)
-
-        ds_out[outvar] = xr.DataArray(out1, dims=outdims, coords=coords)
-        ds_out[outvar].attrs = ens[v].attrs
-        if "description" in ds_out[outvar].attrs.keys():
-            ds_out[outvar].attrs[
-                "description"
-            ] = f"{ds_out[outvar].attrs['description']} : {p}th percentile of ensemble"
-        else:
-            ds_out[outvar].attrs["description"] = f"{p}th percentile of ensemble"
-    return ds_out
-
-
-def _calc_percentiles_blocks(ens, v, values, time_block):
-    ds_out = ens.drop_vars(names=set(ens.data_vars))
-    dims = list(ens[v].dims)
-    outdims = [x for x in dims if "realization" not in x]
-
-    blocks = list(range(0, len(ens.time) + 1, int(time_block)))
-    if blocks[-1] != len(ens[v].time):
-        blocks.append(len(ens[v].time))
-    arr_p_all = {}
-    for t in range(0, len(blocks) - 1):
-        # print('Calculating block ', t + 1, ' of ', len(blocks) - 1)
-        time_sel = slice(blocks[t], blocks[t + 1])
-        arr = (
-            ens[v].isel(time=time_sel).load()
-        )  # percentile calc requires loading the array
-        coords = {}
-        for c in outdims:
-            coords[c] = arr[c]
-        for p in values:
-
-            out1 = _calc_perc(arr, p)
-
-            if t == 0:
-                arr_p_all[str(p)] = xr.DataArray(out1, dims=outdims, coords=coords)
-            else:
-                arr_p_all[str(p)] = xr.concat(
-                    [
-                        arr_p_all[str(p)],
-                        xr.DataArray(out1, dims=outdims, coords=coords),
-                    ],
-                    dim="time",
-                )
-    for p in values:
-        outvar = v + "_p" + str(p)
-        ds_out[outvar] = arr_p_all[str(p)]
-        ds_out[outvar].attrs = ens[v].attrs
-        if "description" in ds_out[outvar].attrs.keys():
-            ds_out[outvar].attrs[
-                "description"
-            ] = f"{ds_out[outvar].attrs['description']} : {p}th percentile of ensemble"
+def _calc_perc(arr, p=50):
+    """Ufunc-like computing a percentile over the last axis of the array.
 
-        else:
-            ds_out[outvar].attrs["description"] = f"{p}th percentile of ensemble"
+    Processes cases with invalid values separately, which makes it more efficent than np.nanpercentile for array with only a few invalid points.
 
-    return ds_out
-
-
-def _calc_perc(arr, p):
-    dims = arr.dims
-    # make sure realization is the first dimension
-    if dims.index("realization") != 0:
-        arr = arr.transpose(
-            "realization", *[dim for dim in dims if dim != "realization"]
-        )
+    Parameters
+    ----------
+    arr : np.array
+        Percentile is computed over the last axis.
+    p : scalar
+        Percentile to compute, between 0 and 100. (the default is 50)
 
-    nan_count = np.isnan(arr).sum(axis=0)
-    out = np.percentile(arr, p, axis=0)
-    if np.any((nan_count > 0) & (nan_count < arr.shape[0])):
-        arr1 = arr.values.reshape(arr.shape[0], int(arr.size / arr.shape[0]))
-        # only use nanpercentile where we need it (slow performace compared to standard) :
-        nan_index = np.where((nan_count > 0) & (nan_count < arr.shape[0]))
+    Returns
+    -------
+    np.array
+    """
+    nan_count = np.isnan(arr).sum(axis=-1)
+    out = np.percentile(arr, p, axis=-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])
+        # 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=0
+            arr1[t, :], p, axis=-1
         )
-
     return out
 
 


Test patch:
diff --git a/tests/test_ensembles.py b/tests/test_ensembles.py
index 60f94f58..34aa4aa7 100644
--- a/tests/test_ensembles.py
+++ b/tests/test_ensembles.py
@@ -137,10 +137,13 @@ class TestEnsembleStats:
         out1 = ensembles.ensemble_percentiles(ens, values=(25, 75))
         assert np.all(out1["tg_mean_p75"] > out1["tg_mean_p25"])
 
-    def test_calc_perc_blocks(self):
+    @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)
-        out1 = ensembles.ensemble_percentiles(ens)
-        out2 = ensembles.ensemble_percentiles(ens, values=(10, 50, 90), time_block=10)
+        out2 = ensembles.ensemble_percentiles(
+            ens.chunk({"time": 2}), values=(10, 50, 90), 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"])