Repository: Ouranosinc/xclim

Issue:
sdba - "verbose" mode?
As suggested in slack by @dgergel and also discussed this morning with Patrick, it could be interesting to output more than just the training dataset / scenario timeseries. I am thinking of some kind of "verbose" or "debug" mode, where sdba trains and adjusts would provide ancillary "diagnostics" variables to help the user understand what happened.

What would be the implementation? A class/method keyword argument? A new option in `xclim.set_options`?

Non exhaustive list of interesting diagnostics:
- EQM/DQM/QDM
    * `quantiles` : Which quantiles of the traning was used by each element of sim
    * `extrapolated` : Boolean flag True for elements where extrapolation was used
- ExtremeValues
   * `nclusters` : The number of clusters (and extreme values) re-adjusted.
- PrincipalComponents
   * The "best" orientation vector. 

Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index d6dafcde..073c2715 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -13,6 +13,7 @@ New features and enhancements
 * Indicator modules built from YAML can now use custom indices. A mapping or module of them can be given to ``build_indicator_module_from_yaml`` with the ``indices`` keyword.
 * Virtual submodules now include an `iter_indicators` function to iterate over the pairs of names and indicator objects in that module.
 * The indicator string formatter now accepts a "r" modifier which passes the raw strings instead of the adjective version.
+* Addition of the `sdba_extra_output` option to adds extra diagnostic variables to the outputs of Adjustment objects. Implementation of `sim_q` in QuantileDeltaMapping and `nclusters` in ExtremeValues.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
diff --git a/docs/notebooks/sdba-advanced.ipynb b/docs/notebooks/sdba-advanced.ipynb
index 9214456c..10307885 100644
--- a/docs/notebooks/sdba-advanced.ipynb
+++ b/docs/notebooks/sdba-advanced.ipynb
@@ -10,7 +10,7 @@
     "\n",
     "## Optimization with dask\n",
     "\n",
-    "Adjustment processes can be very heavy when we need to compute them over large regions and long timeseries. Using small groupings (like `time.dayofyear` ) adds precision and robustness, but also decuplates the load and computing complexity. Fortunately, unlike the heroic pioneers of scientific computing who managed to write parallelized Fortran, we now have [`dask`](https://dask.org/). With only a few parameters, we can magically distribute the computing load to multiple workers and threads.\n",
+    "Adjustment processes can be very heavy when we need to compute them over large regions and long timeseries. Using small groupings (like `time.dayofyear` ) adds precision and robustness, but also decuplates the load and computing complexity. Fortunately, unlike the heroic pioneers of scientific computing who managed to write parallelized Fortran, we now have [dask](https://dask.org/). With only a few parameters, we can magically distribute the computing load to multiple workers and threads.\n",
     "\n",
     "A good first read on the use of dask within xarray are the latter's [Optimization tips](http://xarray.pydata.org/en/stable/dask.html#optimization-tips).\n",
     "\n",
@@ -18,7 +18,7 @@
     "\n",
     "* Most adjustment method will need to perform operation on the whole `time` coordinate, so it is best to optimize chunking along the other dimensions. This is often different from how public data is shared, where more universal 3D chunks are used. \n",
     "\n",
-    "    Chunking of outputs can be controlled in xarray's [`to_netcdf`](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_netcdf.html?highlight=to_netcdf#xarray.Dataset.to_netcdf). We also suggest using [Zarr](https://zarr.readthedocs.io/en/stable/) files. According to [its creators](https://ui.adsabs.harvard.edu/abs/2018AGUFMIN33A..06A/abstract), `zarr` stores should give better performances, especially because of their better ability for parallel I/O. See [`Dataset.to_zarr`](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_zarr.html?highlight=to_zarr#xarray.Dataset.to_zarr) and this useful [rechunking package](https://rechunker.readthedocs.io).\n",
+    "    Chunking of outputs can be controlled in xarray's [to_netcdf](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_netcdf.html?highlight=to_netcdf#xarray.Dataset.to_netcdf). We also suggest using [Zarr](https://zarr.readthedocs.io/en/stable/) files. According to [its creators](https://ui.adsabs.harvard.edu/abs/2018AGUFMIN33A..06A/abstract), `zarr` stores should give better performances, especially because of their better ability for parallel I/O. See [Dataset.to_zarr](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_zarr.html?highlight=to_zarr#xarray.Dataset.to_zarr) and this useful [rechunking package](https://rechunker.readthedocs.io).\n",
     "\n",
     "\n",
     "* One of the main bottleneck for adjustments with small groups is that dask needs to build and optimize an enormous task graph. This issue has been greatly reduced with xclim 0.27 and the use of `map_blocks` in the adjustment methods. However, not all adjustment methods use this optimized syntax.\n",
@@ -278,6 +278,35 @@
    "source": [
     "QDM2.adjust(sim)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Retrieving extra output diagnostics\n",
+    "\n",
+    "To fully understand what is happening during the bias-adjustment process, sdba can output _diagnostic_ variables, giving more visibility to what the adjustment is doing behind the scene. This behaviour, a `verbose` option, is controlled by the `sdba_extra_output` option, set with `xclim.set_options`. When `True`, `train` calls are instructed to include additional variables to the training datasets. In addition, the `adjust` calls will always output a dataset, with `scen` and, depending on the algorithm, other diagnostics variables. See the documentation of each `Adjustment` objects to see what extra variables are available.\n",
+    "\n",
+    "For the moment, this feature is still in construction and only a few `Adjustment` actually provide extra outputs. Please open issues on the Github repo if you have needs or ideas of interesting diagnostic variables.\n",
+    "\n",
+    "For example, `QDM.adjust` adds `sim_q`, which gives the quantile of each element of `sim` within its group."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from xclim import set_options\n",
+    "\n",
+    "with set_options(sdba_extra_output=True):\n",
+    "    QDM = QuantileDeltaMapping(nquantiles=15, kind='+', group='time.dayofyear')\n",
+    "    QDM.train(ref, hist)\n",
+    "    out = QDM.adjust(sim)\n",
+    "\n",
+    "out.sim_q"
+   ]
   }
  ],
  "metadata": {
diff --git a/xclim/core/options.py b/xclim/core/options.py
index 9a657516..28f25b90 100644
--- a/xclim/core/options.py
+++ b/xclim/core/options.py
@@ -20,6 +20,7 @@ CF_COMPLIANCE = "cf_compliance"
 CHECK_MISSING = "check_missing"
 MISSING_OPTIONS = "missing_options"
 RUN_LENGTH_UFUNC = "run_length_ufunc"
+SDBA_EXTRA_OUTPUT = "sdba_extra_output"
 
 MISSING_METHODS: Dict[str, Callable] = dict()
 
@@ -30,6 +31,7 @@ OPTIONS = {
     CHECK_MISSING: "any",
     MISSING_OPTIONS: dict(),
     RUN_LENGTH_UFUNC: "auto",
+    SDBA_EXTRA_OUTPUT: False,
 }
 
 _LOUDNESS_OPTIONS = frozenset(["log", "warn", "raise"])
@@ -57,6 +59,7 @@ _VALIDATORS = {
     CHECK_MISSING: lambda meth: meth != "from_context" and meth in MISSING_METHODS,
     MISSING_OPTIONS: _valid_missing_options,
     RUN_LENGTH_UFUNC: _RUN_LENGTH_UFUNC_OPTIONS.__contains__,
+    SDBA_EXTRA_OUTPUT: lambda opt: isinstance(opt, bool),
 }
 
 
@@ -148,6 +151,12 @@ class set_options:
     - ``run_length_ufunc``:
       Whether to use the 1D ufunc version of run length algorithms or the dask-ready broadcasting version.
       Default is ``'auto'`` which means the latter is used for dask-backed and large arrays.
+    - ``sdba_extra_output``:
+      Whether to add diagnostic variables to outputs of sdba's `train`, `adjust`
+      and `processing` operations.  Details about these additional variables are given in the object's
+      docstring. When activated,  `adjust` will return a Dataset with `scen`  and those extra diagnostics
+      For `processing` functions, see the doc, the output type might change, or not depending on the
+      algorithm. Default: ``False``.
 
     Examples
     --------
diff --git a/xclim/sdba/_adjustment.py b/xclim/sdba/_adjustment.py
index 5de4c81d..e18701fe 100644
--- a/xclim/sdba/_adjustment.py
+++ b/xclim/sdba/_adjustment.py
@@ -94,7 +94,7 @@ def dqm_scale_sim(ds, *, group, interp, kind):
     return sim.rename("sim").to_dataset()
 
 
-@map_blocks(reduces=[Grouper.PROP, "quantiles"], scen=[])
+@map_blocks(reduces=[Grouper.PROP, "quantiles"], scen=[], sim_q=[])
 def qdm_adjust(ds, *, group, interp, extrapolation, kind):
     """QDM: Adjust process on one block.
 
@@ -111,7 +111,7 @@ def qdm_adjust(ds, *, group, interp, extrapolation, kind):
     af = u.broadcast(af, ds.sim, group=group, interp=interp, sel=sel)
 
     scen = u.apply_correction(ds.sim, af, kind)
-    return scen.rename("scen").to_dataset()
+    return xr.Dataset(dict(scen=scen, sim_q=sim_q))
 
 
 @map_blocks(reduces=[Grouper.DIM], af=[Grouper.PROP], hist_thresh=[Grouper.PROP])
diff --git a/xclim/sdba/adjustment.py b/xclim/sdba/adjustment.py
index 6a086231..a5419cb4 100644
--- a/xclim/sdba/adjustment.py
+++ b/xclim/sdba/adjustment.py
@@ -8,6 +8,7 @@ from xarray.core.dataarray import DataArray
 
 from xclim.core.calendar import get_calendar
 from xclim.core.formatting import update_history
+from xclim.core.options import OPTIONS, SDBA_EXTRA_OUTPUT
 from xclim.core.units import convert_units_to
 from xclim.indices import stats
 
@@ -147,13 +148,22 @@ class BaseAdjustment(ParametrizableWithDataset):
                     stacklevel=4,
                 )
 
-        scen = self._adjust(sim, **kwargs)
+        out = self._adjust(sim, **kwargs)
+
+        if isinstance(out, xr.DataArray):
+            out = out.rename("scen").to_dataset()
+
+        scen = out.scen
+
         params = ", ".join([f"{k}={repr(v)}" for k, v in kwargs.items()])
         infostr = f"{str(self)}.adjust(sim, {params})"
         scen.attrs["xclim_history"] = update_history(
             f"Bias-adjusted with {infostr}", sim
         )
         scen.attrs["bias_adjustment"] = infostr
+
+        if OPTIONS[SDBA_EXTRA_OUTPUT]:
+            return out
         return scen
 
     def set_dataset(self, ds: xr.Dataset):
@@ -423,6 +433,12 @@ class QuantileDeltaMapping(EmpiricalQuantileMapping):
         extrapolation : {'constant', 'nan'}
           The type of extrapolation to use. See :py:func:`xclim.sdba.utils.extrapolate_qm` for details. Defaults to "constant".
 
+        Extra diagnostics
+        -----------------
+        In adjustment:
+
+        quantiles : The quantile of each value of `sim`. The adjustment factor is interpolated using this as the "quantile" axis on `ds.af`.
+
         References
         ----------
         Cannon, A. J., Sobie, S. R., & Murdock, T. Q. (2015). Bias correction of GCM precipitation by quantile mapping: How well do methods preserve changes in quantiles and extremes? Journal of Climate, 28(17), 6938–6959. https://doi.org/10.1175/JCLI-D-14-00754.1
@@ -430,14 +446,17 @@ class QuantileDeltaMapping(EmpiricalQuantileMapping):
         super().__init__(**kwargs)
 
     def _adjust(self, sim, interp="nearest", extrapolation="constant"):
-        scen = qdm_adjust(
+        out = qdm_adjust(
             xr.Dataset({"sim": sim, "af": self.ds.af, "hist_q": self.ds.hist_q}),
             group=self.group,
             interp=interp,
             extrapolation=extrapolation,
             kind=self.kind,
-        ).scen
-        return scen
+        )
+        if OPTIONS[SDBA_EXTRA_OUTPUT]:
+            out.sim_q.attrs.update(long_name="Group-wise quantiles of `sim`.")
+            return out
+        return out.scen
 
 
 class ExtremeValues(BaseAdjustment):
@@ -480,6 +499,12 @@ class ExtremeValues(BaseAdjustment):
         power: float
           Shape of the correction strength, see Notes.
 
+        Extra diagnostics
+        -----------------
+        In training:
+
+        nclusters : Number of extreme value clusters found for each gridpoint.
+
         Notes
         -----
         Extreme values are extracted from `ref`, `hist` and `sim` by finding all "clusters",
@@ -576,6 +601,10 @@ class ExtremeValues(BaseAdjustment):
             description=f"Mean of the {self.q_thresh * 100}th percentile of large values (x > {self.cluster_thresh}) of ref and hist.",
         )
         self["hist_calendar"] = get_calendar(hist)
+
+        if OPTIONS[SDBA_EXTRA_OUTPUT] and ref_params is None:
+            ds = ds.assign(nclusters=ref_clusters.nclusters)
+
         self.set_dataset(ds)
 
     def adjust(


Test patch:
diff --git a/xclim/testing/tests/test_sdba/test_adjustment.py b/xclim/testing/tests/test_sdba/test_adjustment.py
index 8521d5aa..5667a4d7 100644
--- a/xclim/testing/tests/test_sdba/test_adjustment.py
+++ b/xclim/testing/tests/test_sdba/test_adjustment.py
@@ -3,6 +3,7 @@ import pytest
 import xarray as xr
 from scipy.stats import genpareto, norm, uniform
 
+from xclim.core.options import set_options
 from xclim.core.units import convert_units_to
 from xclim.sdba.adjustment import (
     LOCI,
@@ -338,13 +339,19 @@ class TestQDM:
         # Test predict
         np.testing.assert_allclose(p, ref.transpose(*p.dims), rtol=0.1, atol=0.2)
 
-    def test_cannon(self, cannon_2015_dist, cannon_2015_rvs):
+    def test_cannon_and_diagnostics(self, cannon_2015_dist, cannon_2015_rvs):
         ref, hist, sim = cannon_2015_rvs(15000, random=False)
 
         # Quantile mapping
-        QDM = QuantileDeltaMapping(kind="*", group="time", nquantiles=50)
-        QDM.train(ref, hist)
-        bc_sim = QDM.adjust(sim)
+        with set_options(sdba_extra_output=True):
+            QDM = QuantileDeltaMapping(kind="*", group="time", nquantiles=50)
+            QDM.train(ref, hist)
+            scends = QDM.adjust(sim)
+
+        assert isinstance(scends, xr.Dataset)
+
+        sim_q_exp = sim.rank(dim="time", pct=True)
+        np.testing.assert_array_equal(sim_q_exp, scends.sim_q)
 
         # Theoretical results
         # ref, hist, sim = cannon_2015_dist
@@ -357,7 +364,7 @@ class TestQDM:
         # mean = np.trapz(pdf * pu, pu)
         # mom2 = np.trapz(pdf * pu ** 2, pu)
         # std = np.sqrt(mom2 - mean ** 2)
-
+        bc_sim = scends.scen
         np.testing.assert_almost_equal(bc_sim.mean(), 41.5, 1)
         np.testing.assert_almost_equal(bc_sim.std(), 16.7, 0)
 
@@ -523,14 +530,14 @@ class TestPrincipalComponents:
 @pytest.mark.slow
 class TestExtremeValues:
     @pytest.mark.parametrize(
-        "c_thresh,q_thresh,frac,power",
+        "c_thresh,q_thresh,frac,power,diags",
         [
-            ["1 mm/d", 0.95, 0.25, 1],
-            ["1 mm/d", 0.90, 1e-6, 1],
-            ["0.007 m/week", 0.95, 0.25, 2],
+            ["1 mm/d", 0.95, 0.25, 1, True],
+            ["1 mm/d", 0.90, 1e-6, 1, False],
+            ["0.007 m/week", 0.95, 0.25, 2, False],
         ],
     )
-    def test_simple(self, c_thresh, q_thresh, frac, power):
+    def test_simple(self, c_thresh, q_thresh, frac, power, diags):
         n = 45 * 365
 
         def gen_testdata(c, s):
@@ -554,10 +561,16 @@ class TestExtremeValues:
 
         EQM = EmpiricalQuantileMapping(group="time.dayofyear", nquantiles=15, kind="*")
         EQM.train(ref, hist)
+
         scen = EQM.adjust(sim)
 
         EX = ExtremeValues(c_thresh, q_thresh=q_thresh)
-        EX.train(ref, hist)
+
+        with set_options(sdba_extra_output=diags):
+            EX.train(ref, hist)
+
+        if diags:
+            assert "nclusters" in EX.ds
 
         qv = (ref.thresh + hist.thresh) / 2
         np.testing.assert_allclose(EX.ds.fit_params, [-0.1, qv, 2], atol=0.5, rtol=0.1)
@@ -574,7 +587,7 @@ class TestExtremeValues:
         # ONLY extreme values have been touched (but some might not have been modified)
         assert (((scen != scen2) | exval) == exval).all()
 
-    def test_dask_julia(self):
+    def test_dask_julia_diags(self):
 
         dsim = open_dataset("sdba/CanESM2_1950-2100.nc").chunk()
         dref = open_dataset("sdba/ahccd_1950-2013.nc").chunk()