reorder with window

CHANGES.rst

@@ -21,6 +21,7 @@ New features and enhancements
 * Support ``indexer`` keyword in YAML indicator description. (:issue:`1522`, :pull:`1561`).
 * New ``xclim.core.calendar.stack_periods`` and ``unstack_periods`` for performing ``rolling(time=...).construct(..., stride=...)`` but with non-uniform temporal periods like years or months. They replace ``xclim.sdba.processing.construct_moving_yearly_window`` and ``unpack_moving_yearly_window`` which are deprecated and will be removed in a future release.
 * New ``as_dataset`` options for ``xclim.set_options``. When True, indicators will output Datasets instead of DataArrays. (:issue:`1257`, :pull:`1625`).
+* Add the possibility to use a group with a window in ``xc.sdba.processing.reordering``. (:pull:`1566`).
 
 Breaking changes
 ^^^^^^^^^^^^^^^^

tests/test_sdba/test_processing.py

@@ -5,6 +5,7 @@
 import pytest
 import xarray as xr
 
+from xclim.core.calendar import date_range
 from xclim.core.units import units
 from xclim.sdba.adjustment import EmpiricalQuantileMapping
 from xclim.sdba.base import Grouper
@@ -185,6 +186,31 @@ def test_reordering():
     assert out.attrs == y.attrs
 
 
+def test_reordering_with_window():
+    time = list(
+        date_range("2000-01-01", "2000-01-04", freq="D", calendar="noleap")
+    ) + list(date_range("2001-01-01", "2001-01-04", freq="D", calendar="noleap"))
+
+    x = xr.DataArray(
+        np.arange(1, 9, 1),
+        dims=("time"),
+        coords={"time": time},
+    )
+
+    y = xr.DataArray(
+        np.arange(8, 0, -1),
+        dims=("time"),
+        coords={"time": time},
+    )
+
+    group = Grouper(group="time.dayofyear", window=3)
+    out = reordering(x, y, group=group)
+
+    np.testing.assert_array_equal(out, [3.0, 3.0, 2.0, 2.0, 7.0, 7.0, 6.0, 6.0])
+    out.attrs.pop("history")
+    assert out.attrs == y.attrs
+
+
 def test_to_additive(pr_series, hurs_series):
     # log
     pr = pr_series(np.array([0, 1e-5, 1, np.e**10]))

xclim/sdba/_processing.py

@@ -142,7 +142,7 @@ def _normalize(
     )
 
 
-@map_groups(reordered=[Grouper.DIM], main_only=True)
+@map_groups(reordered=[Grouper.DIM], main_only=False)
 def _reordering(ds, *, dim):
     """Group-wise reordering.
 
@@ -159,17 +159,47 @@ def _reordering(ds, *, dim):
     def _reordering_1d(data, ordr):
         return np.sort(data)[np.argsort(np.argsort(ordr))]
 
-    return (
-        xr.apply_ufunc(
-            _reordering_1d,
-            ds.sim,
-            ds.ref,
-            input_core_dims=[[dim], [dim]],
-            output_core_dims=[[dim]],
-            vectorize=True,
-            dask="parallelized",
-            output_dtypes=[ds.sim.dtype],
+    def _reordering_2d(data, ordr):
+        data_r = data.ravel()
+        ordr_r = ordr.ravel()
+        reorder = np.sort(data_r)[np.argsort(np.argsort(ordr_r))]
+        return reorder.reshape(data.shape)[
+            :, int(data.shape[1] / 2)
+        ]  # pick the middle of the window
+
+    if "window" in dim:
+        return (
+            xr.apply_ufunc(
+                _reordering_2d,
+                ds.sim,
+                ds.ref,
+                input_core_dims=[["time", "window"], ["time", "window"]],
+                output_core_dims=[["time"]],
+                vectorize=True,
+                dask="parallelized",
+                output_dtypes=[ds.sim.dtype],
+            )
+            .rename("reordered")
+            .to_dataset()
+        )
+    elif len(dim) == 1:
+        return (
+            xr.apply_ufunc(
+                _reordering_1d,
+                ds.sim,
+                ds.ref,
+                input_core_dims=[dim, dim],
+                output_core_dims=[dim],
+                vectorize=True,
+                dask="parallelized",
+                output_dtypes=[ds.sim.dtype],
+            )
+            .rename("reordered")
+            .to_dataset()
+        )
+    else:
+        raise ValueError(
+            f"Reordering can only be done along one dimension."
+            f" If there is more than one, one of the dimension should be `window`."
+            f" The dimensions are {dim}."
         )
-        .rename("reordered")
-        .to_dataset()
-    )