Repository: Ouranosinc/xclim
Issue category: Performance Issue

sdba - improve the speed of DetrendedQuantileMapping
### Description
Doing a few preliminary tests **on a single pixel and without dask**, I obtained the following times : 

DetrendedQuantileMapping, groups='month', 150 years of daily data : ~15 seconds
DetrendedQuantileMapping, groups='day', window=31, 150 years of daily data : ~24 seconds

EmpiricalQuantileMapping, groups='month', 150 years of daily data : ~0.35 seconds
EmpiricalQuantileMapping, groups='day', window=31, 150 years of daily data : ~2.5 seconds

There seems to be a bottleneck in the detrending version. Discussions with @aulemahal point towards the normalization process and `sdba.utils.interp_on_quantiles`. I will continue investigating and post my findings here.


Functions to edit:
xclim/sdba/adjustment.py:DetrendedQuantileMapping._train
xclim/sdba/adjustment.py:DetrendedQuantileMapping._adjust
xclim/sdba/adjustment.py:DetrendedQuantileMapping
xclim/sdba/base.py:Grouper.get_index
xclim/sdba/detrending.py:BaseDetrend.fit
xclim/sdba/detrending.py:BaseDetrend.get_trend
xclim/sdba/detrending.py:BaseDetrend._set_fitds
xclim/sdba/processing.py:jitter_under_thresh
xclim/sdba/processing.py:normalize
xclim/sdba/utils.py:interp_on_quantiles
Functions added: xclim/sdba/adjustment.py:DetrendedQuantileMapping.__init__
xclim/sdba/detrending.py:BaseDetrend._set_ds

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 1d6a6ebc5..0bd3055e1 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -4,6 +4,7 @@ History
 
 0.18.x
 ------
+* Optimization options for `xclim.sdba` : different grouping for the normalization steps of DQM and save training or fitting datasets to temporary files.
 * `xclim.sdba.detrending` objects can now act on groups.
 * Replaced `dask[complete]` with `dask[array]` in basic installation and added `distributed` to `docs` build dependencies.
 * `xclim.core.locales` now supported in Windows build environments.
diff --git a/xclim/sdba/adjustment.py b/xclim/sdba/adjustment.py
index 81c6f446f..9a4a8d8ff 100644
--- a/xclim/sdba/adjustment.py
+++ b/xclim/sdba/adjustment.py
@@ -224,9 +224,10 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
     The algorithm follows these steps, 1-3 being the 'train' and 4-6, the 'adjust' steps.
 
     1. A scaling factor that would make the mean of `hist` match the mean of `ref` is computed.
-    2. `ref` and `hist` are normalized by removing the group-wise mean.
+    2. `ref` and `hist` are normalized by removing the "dayofyear" mean.
     3. Adjustment factors are computed between the quantiles of the normalized `ref` and `hist`.
-    4. `sim` is corrected by the scaling factor, and detrended group-wise using a linear fit.
+    4. `sim` is corrected by the scaling factor, and either normalized by "dayofyear" and  detrended group-wise
+       or directly detrended per "dayofyear", using a linear fit (modifiable).
     5. Values of detrended `sim` are matched to the corresponding quantiles of normalized `hist` and corrected accordingly.
     6. The trend is put back on the result.
 
@@ -246,7 +247,10 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
     kind : {'+', '*'}
       The adjustment kind, either additive or multiplicative.
     group : Union[str, Grouper]
-      The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
+      The grouping information used in the quantile mapping process. See :py:class:`xclim.sdba.base.Grouper` for details.
+      the normalization step is always performed on each day of the year.
+    norm_window : 1
+      The window size used in the normalization grouping. Defaults to 1.
 
     In adjustment:
 
@@ -256,26 +260,54 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
       The method to use when detrending. If an int is passed, it is understood as a PolyDetrend (polynomial detrending) degree. Defaults to 1 (linear detrending)
     extrapolation : {'constant', 'nan'}
       The type of extrapolation to use. See :py:func:`xclim.sdba.utils.extrapolate_qm` for details. Defaults to "constant".
+    normalize_sim : bool
+      If True, scaled sim is normalized by its "dayofyear" mean and then detrended using `group`.
+        The norm is broadcasted and added back on scen using `interp='nearest'`, ignoring the passed `interp`.
+      If False, scaled sim is detrended per "dayofyear".
+      This is useful on large datasets using dask, in which case "dayofyear" is a very small division,
+        because normalization is a more efficient operation than detrending for similarly sized groups.
 
     References
     ----------
     .. [Cannon2015] 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
     """
 
+    @parse_group
+    def __init__(
+        self,
+        *,
+        nquantiles: int = 20,
+        kind: str = ADDITIVE,
+        group: Union[str, Grouper] = "time",
+        norm_window: int = 1,
+    ):
+        super().__init__(
+            nquantiles=nquantiles, kind=kind, group=group,
+        )
+        self["norm_group"] = Grouper("time.dayofyear", window=norm_window)
+
     def _train(self, ref, hist):
+        refn = normalize(ref, group=self.norm_group, kind=self.kind)
+        histn = normalize(hist, group=self.norm_group, kind=self.kind)
+        super()._train(refn, histn)
+
         mu_ref = self.group.apply("mean", ref)
         mu_hist = self.group.apply("mean", hist)
-        ref = normalize(ref, group=self.group, kind=self.kind)
-        hist = normalize(hist, group=self.group, kind=self.kind)
-        super()._train(ref, hist)
-
         self.ds["scaling"] = get_correction(mu_hist, mu_ref, kind=self.kind)
         self.ds.scaling.attrs.update(
             standard_name="Scaling factor",
             description="Scaling factor making the mean of hist match the one of hist.",
         )
 
-    def _adjust(self, sim, interp="nearest", extrapolation="constant", detrend=1):
+    def _adjust(
+        self,
+        sim,
+        interp="nearest",
+        extrapolation="constant",
+        detrend=1,
+        normalize_sim=False,
+    ):
+
         # Apply preliminary scaling from obs to hist
         sim = apply_correction(
             sim,
@@ -283,9 +315,17 @@ def _adjust(self, sim, interp="nearest", extrapolation="constant", detrend=1):
             self.kind,
         )
 
+        if normalize_sim:
+            sim_norm = self.norm_group.apply("mean", sim)
+            sim = normalize(sim, group=self.norm_group, kind=self.kind, norm=sim_norm)
+
         # Find trend on sim
         if isinstance(detrend, int):
-            detrend = PolyDetrend(degree=detrend, kind=self.kind, group=self.group)
+            detrend = PolyDetrend(
+                degree=detrend,
+                kind=self.kind,
+                group=self.group if normalize_sim else self.norm_group,
+            )
 
         sim_fit = detrend.fit(sim)
         sim_detrended = sim_fit.detrend(sim)
@@ -295,7 +335,15 @@ def _adjust(self, sim, interp="nearest", extrapolation="constant", detrend=1):
             sim_detrended, extrapolation=extrapolation, interp=interp
         )
         # Retrend
-        return sim_fit.retrend(scen_detrended)
+        scen = sim_fit.retrend(scen_detrended)
+
+        if normalize_sim:
+            return apply_correction(
+                scen,
+                broadcast(sim_norm, scen, group=self.norm_group, interp="nearest"),
+                self.kind,
+            )
+        return scen
 
 
 class QuantileDeltaMapping(EmpiricalQuantileMapping):
diff --git a/xclim/sdba/base.py b/xclim/sdba/base.py
index 7eaef29eb..a0222136d 100644
--- a/xclim/sdba/base.py
+++ b/xclim/sdba/base.py
@@ -170,9 +170,9 @@ def get_index(
         ind = da.indexes[self.dim]
         i = getattr(ind, self.prop)
 
-        if i.dtype != np.int:
+        if not np.issubdtype(i.dtype, np.integer):
             raise ValueError(
-                f"Index {self.name} is not of type int (rather {ind.dtype}), but {self.__class__.__name__} requires integer indexes."
+                f"Index {self.name} is not of type int (rather {i.dtype}), but {self.__class__.__name__} requires integer indexes."
             )
 
         interp = (
diff --git a/xclim/sdba/detrending.py b/xclim/sdba/detrending.py
index cce24308a..20cc22790 100644
--- a/xclim/sdba/detrending.py
+++ b/xclim/sdba/detrending.py
@@ -52,16 +52,23 @@ def fit(self, da: xr.DataArray):
         Returns a new object storing the fit data that can be used for detrending and retrending.
         """
         new = self.copy()
-        new._set_fitds(new.group.apply(new._fit, da, main_only=True))
+        new._set_ds(new.group.apply(new._fit, da, main_only=True))
         new.__fitted = True
         return new
 
     def get_trend(self, da: xr.DataArray):
-        return self.group.apply(
+        """Get the trend computed from the fit, the fitting dim as found on da.
+
+        If da is a DataArray (and has a "dtype" attribute), the trend is casted to have the same dtype.
+        """
+        out = self.group.apply(
             self._get_trend,
-            {self.group.dim: da[self.group.dim], **self.fit_ds.data_vars},
+            {self.group.dim: da[self.group.dim], **self.ds.data_vars},
             main_only=True,
         )
+        if hasattr(da, "dtype"):
+            out = out.astype(da.dtype)
+        return out
 
     def detrend(self, da: xr.DataArray):
         """Removes the previously fitted trend from a DataArray."""
@@ -77,8 +84,9 @@ def retrend(self, da: xr.DataArray):
         trend = self.get_trend(da)
         return self._retrend(da, trend)
 
-    def _set_fitds(self, ds):
-        self.fit_ds = ds
+    def _set_ds(self, ds):
+        self.ds = ds
+        self.ds.attrs["fit_params"] = str(self)
 
     def _detrend(self, da, trend):
         # Remove trend from series
diff --git a/xclim/sdba/processing.py b/xclim/sdba/processing.py
index e2239eff5..0b25875d7 100644
--- a/xclim/sdba/processing.py
+++ b/xclim/sdba/processing.py
@@ -1,4 +1,5 @@
 """Pre and post processing for bias adjustement"""
+from typing import Optional
 from typing import Union
 
 import dask.array as dsk
@@ -9,6 +10,7 @@
 from .base import parse_group
 from .utils import ADDITIVE
 from .utils import apply_correction
+from .utils import broadcast
 from .utils import ecdf
 from .utils import invert
 
@@ -167,15 +169,24 @@ def jitter_under_thresh(x: xr.DataArray, thresh: float):
     If thresh is high, this will change the mean value of x.
     """
     epsilon = np.finfo(x.dtype).eps
-    jitter = np.random.uniform(low=epsilon, high=thresh, size=x.shape)
+    if isinstance(x.data, dsk.Array):
+        jitter = dsk.random.uniform(
+            low=epsilon, high=thresh, size=x.shape, chunks=x.chunks
+        )
+    else:
+        jitter = np.random.uniform(low=epsilon, high=thresh, size=x.shape)
     return x.where(~((x < thresh) & (x.notnull())), jitter)
 
 
 @parse_group
 def normalize(
-    x: xr.DataArray, *, group: Union[str, Grouper] = "time", kind: str = ADDITIVE
+    x: xr.DataArray,
+    *,
+    group: Union[str, Grouper] = "time",
+    kind: str = ADDITIVE,
+    norm: Optional[xr.DataArray] = None,
 ):
-    """Normalize an array by remove its mean.
+    """Normalize an array by removing its mean.
 
     Normalization if performed group-wise.
 
@@ -187,9 +198,21 @@ def normalize(
       Grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
     kind : {'+', '*'}
       How to apply the adjustment, either additively or multiplicatively.
+    norm : xr.DataArray
+      If the norm was already computed (for example with `group.apply("mean", x)`), skip the
+        computation step. The array should have the same dimensions as `x` except for "time" that should
+        be replaced by `group.prop`.
+
+    Returns
+    -------
+    xr.DataArray or xr.Dataset
+      Group-wise anomaly of x
     """
 
     def _normalize_group(grp, dim=["time"]):
         return apply_correction(grp, invert(grp.mean(dim=dim), kind), kind)
 
-    return group.apply(_normalize_group, x)
+    if norm is None:
+        return group.apply(_normalize_group, x)
+
+    return apply_correction(x, broadcast(norm, x, group=group, interp="nearest"), kind,)
diff --git a/xclim/sdba/utils.py b/xclim/sdba/utils.py
index 24e51e1ee..4911ee8a5 100644
--- a/xclim/sdba/utils.py
+++ b/xclim/sdba/utils.py
@@ -456,7 +456,7 @@ def _interp_quantiles_2D(newx, newg, oldx, oldy, oldg):
         output_core_dims=[[dim]],
         vectorize=True,
         dask="parallelized",
-        output_dtypes=[np.float],
+        output_dtypes=[yq.dtype],
     )
 
 

Test patch: diff --git a/tests/test_sdba/test_adjustment.py b/tests/test_sdba/test_adjustment.py
index c897ea3ca..532f8a67f 100644
--- a/tests/test_sdba/test_adjustment.py
+++ b/tests/test_sdba/test_adjustment.py
@@ -190,7 +190,7 @@ def test_mon_U(self, mon_series, series, mon_triangular, kind, name, spatial_dim
         if spatial_dims:
             mqm = mqm.isel({crd: 0 for crd in spatial_dims.keys()})
         np.testing.assert_array_almost_equal(mqm, int(kind == MULTIPLICATIVE), 1)
-        np.testing.assert_array_almost_equal(p, ref_t, 1)
+        np.testing.assert_allclose(p, ref_t, rtol=0.1, atol=0.5)
 
     def test_cannon(self, cannon_2015_rvs):
         ref, hist, sim = cannon_2015_rvs(15000)
@@ -202,17 +202,6 @@ def test_cannon(self, cannon_2015_rvs):
         np.testing.assert_almost_equal(p.mean(), 41.6, 0)
         np.testing.assert_almost_equal(p.std(), 15.0, 0)
 
-    def test_group_norm(self, tas_series):
-        x = np.arange(365)
-        ref = tas_series(273 - 10 * np.cos(2 * np.pi * x / 365), start="2001-01-01")
-        hist = sim = tas_series(
-            273 - 8 * np.cos(2 * np.pi * x / 365), start="2001-01-01"
-        )
-        DQM = DetrendedQuantileMapping(group="time.month")
-        DQM.train(ref, hist)
-        scen = DQM.adjust(sim, interp="linear")
-        xr.testing.assert_allclose(ref, scen, rtol=1e-3)
-
 
 class TestQDM:
     @pytest.mark.parametrize("kind,name", [(ADDITIVE, "tas"), (MULTIPLICATIVE, "pr")])