Repository: Ouranosinc/xclim

Issue:
sdba - Grouped Detrending
We need to implement grouped detrending into sdba.

Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 2e0e319a..1d6a6ebc 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -4,6 +4,7 @@ History
 
 0.18.x
 ------
+* `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.
 * `ensembles.ensemble_percentiles` modified to compute along a `percentiles` dimension by default, instead of creating different variables.
diff --git a/setup.cfg b/setup.cfg
index 07c12a04..4b6bea72 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.17.3-beta
+current_version = 0.17.4-beta
 commit = False
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+).(?P<patch>\d+)(\-(?P<release>[a-z]+))?
diff --git a/setup.py b/setup.py
index 218b76cc..2222a0ea 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@ URL = "https://github.com/Ouranosinc/xclim"
 AUTHOR = "Travis Logan"
 AUTHOR_EMAIL = "logan.travis@ouranos.ca"
 REQUIRES_PYTHON = ">=3.6.0"
-VERSION = "0.17.3-beta"
+VERSION = "0.17.4-beta"
 LICENSE = "Apache Software License 2.0"
 
 with open("README.rst") as readme_file:
diff --git a/xclim/__init__.py b/xclim/__init__.py
index ae5f1b3b..4c244691 100644
--- a/xclim/__init__.py
+++ b/xclim/__init__.py
@@ -10,4 +10,4 @@ from xclim.indicators import seaIce
 
 __author__ = """Travis Logan"""
 __email__ = "logan.travis@ouranos.ca"
-__version__ = "0.17.3-beta"
+__version__ = "0.17.4-beta"
diff --git a/xclim/sdba/__init__.py b/xclim/sdba/__init__.py
index 2a197896..85d8a64c 100644
--- a/xclim/sdba/__init__.py
+++ b/xclim/sdba/__init__.py
@@ -37,10 +37,10 @@ A generic bias adjustment process is laid out as follows:
 - preprocessing on `ref`, `hist` and `sim` (using methods in `xclim.sdba.processing` or `xclim.sdba.detrending`)
 - creating the adjustment object `Adj = Adjustment(**kwargs)` (from `xclim.sdba.adjustment`)
 - training `Adj.train(obs, sim)`
-- adjustment with corresponding arguments `scen = Adj.adjust(sim, **kwargs)`
+- adjustment `scen = Adj.adjust(sim, **kwargs)`
 - post-processing on `scen` (for example: re-trending)
 
-The train-adjust approach allows to inspect the trained adjustment object. The trained information is stored in
+The train-adjust approach allows to inspect the trained adjustment object. The training information is stored in
 the underlying `Adj.ds` dataset and always has a `af` variable with the adjustment factors. Its layout and the
 other available variables vary between the different algorithm, refer to :ref:`bias-adjustment-algos`.
 
diff --git a/xclim/sdba/adjustment.py b/xclim/sdba/adjustment.py
index 6aa1ccde..c1fbccfb 100644
--- a/xclim/sdba/adjustment.py
+++ b/xclim/sdba/adjustment.py
@@ -132,7 +132,7 @@ class EmpiricalQuantileMapping(BaseAdjustment):
 
     Parameters
     ----------
-    At init:
+    At instantiation:
 
     nquantiles : int
       The number of quantiles to use. Two endpoints at 1e-6 and 1 - 1e-6 will be added.
@@ -141,7 +141,7 @@ class EmpiricalQuantileMapping(BaseAdjustment):
     group : Union[str, Grouper]
       The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
 
-    In adjust:
+    In adjustment:
 
     interp : {'nearest', 'linear', 'cubic'}
       The interpolation method to use when interpolating the adjustment factors. Defaults to "nearset".
@@ -201,12 +201,9 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
     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.
     3. Adjustment factors are computed between the quantiles of the normalized `ref` and `hist`.
-    4. `sim` is corrected by the scaling factor, normalized by the group-wise mean and then detrended using a linear fit.*
+    4. `sim` is corrected by the scaling factor, and detrended group-wise using a linear fit.
     5. Values of detrended `sim` are matched to the corresponding quantiles of normalized `hist` and corrected accordingly.
-    6. The group-wise mean and trend are put back on the result.*
-
-    * Steps 4 and 6 include a group-wise normalization to overcome the fact that detrending is not made group-wise.
-    A future release of xclim will have grouped detrending and not require this extra step any more.
+    6. The trend is put back on the result.
 
     .. math::
 
@@ -217,7 +214,7 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
 
     Parameters
     ----------
-    At init:
+    At instantiation:
 
     nquantiles : int
       The number of quantiles to use. Two endpoints at 1e-6 and 1 - 1e-6 will be added.
@@ -226,7 +223,7 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
     group : Union[str, Grouper]
       The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
 
-    In adjust:
+    In adjustment:
 
     interp : {'nearest', 'linear', 'cubic'}
       The interpolation method to use when interpolating the adjustment factors. Defaults to "nearest".
@@ -261,39 +258,19 @@ class DetrendedQuantileMapping(EmpiricalQuantileMapping):
             self.kind,
         )
 
-        # Normalize sim group-wise
-        # This group-wise pre normalization + reapplication further down
-        # is to circumvent #442 (detrending is not groupwise)
-        mu_sim = self.group.apply("mean", sim)
-        sim_norm = apply_correction(
-            sim,
-            broadcast(
-                invert(mu_sim, kind=self.kind), sim, group=self.group, interp=interp,
-            ),
-            kind=self.kind,
-        )
-
-        # Find trend on sim (for debugging purpose, the detrending is overridable)
+        # Find trend on sim
         if isinstance(detrend, int):
-            detrend = PolyDetrend(degree=detrend, kind=self.kind)
+            detrend = PolyDetrend(degree=detrend, kind=self.kind, group=self.group)
 
-        sim_fit = detrend.fit(sim_norm)
-        sim_detrended = sim_fit.detrend(sim_norm)
+        sim_fit = detrend.fit(sim)
+        sim_detrended = sim_fit.detrend(sim)
 
         # Adjust using `EmpiricalQuantileMapping.adjust`
         scen_detrended = super()._adjust(
             sim_detrended, extrapolation=extrapolation, interp=interp
         )
         # Retrend
-        scen_norm = sim_fit.retrend(scen_detrended)
-
-        # # Reapply-mean
-        scen = apply_correction(
-            scen_norm,
-            broadcast(mu_sim, sim, group=self.group, interp=interp),
-            kind=self.kind,
-        )
-        return scen
+        return sim_fit.retrend(scen_detrended)
 
 
 class QuantileDeltaMapping(EmpiricalQuantileMapping):
@@ -311,7 +288,7 @@ class QuantileDeltaMapping(EmpiricalQuantileMapping):
 
     Parameters
     ----------
-    At init:
+    At instantiation:
 
     nquantiles : int
       The number of quantiles to use. Two endpoints at 1e-6 and 1 - 1e-6 will be added.
@@ -320,7 +297,7 @@ class QuantileDeltaMapping(EmpiricalQuantileMapping):
     group : Union[str, Grouper]
       The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
 
-    In adjust:
+    In adjustment:
 
     interp : {'nearest', 'linear', 'cubic'}
       The interpolation method to use when interpolating the adjustment factors. Defaults to "nearest".
@@ -368,14 +345,14 @@ class LOCI(BaseAdjustment):
 
     Parameters
     ----------
-    At init:
+    At instantiation:
 
     group : Union[str, Grouper]
       The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
     thresh : float
       The threshold in `ref` above which the values are scaled.
 
-    In adjust:
+    In adjustment:
 
     interp : {'nearest', 'linear', 'cubic'}
       The interpolation method to use then interpolating the adjustment factors. Defaults to "linear".
@@ -423,14 +400,14 @@ class Scaling(BaseAdjustment):
 
     Parameters
     ----------
-    At init:
+    At instantiation:
 
     group : Union[str, Grouper]
       The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
     kind : {'+', '*'}
       The adjustment kind, either additive or multiplicative.
 
-    In adjust:
+    In adjustment:
 
     interp : {'nearest', 'linear', 'cubic'}
       The interpolation method to use then interpolating the adjustment factors. Defaults to "nearest".
diff --git a/xclim/sdba/base.py b/xclim/sdba/base.py
index b161820c..e875c7e0 100644
--- a/xclim/sdba/base.py
+++ b/xclim/sdba/base.py
@@ -6,6 +6,7 @@ from typing import Optional
 from typing import Sequence
 from typing import Union
 
+import numpy as np
 import xarray as xr
 from boltons.funcutils import wraps
 
@@ -91,11 +92,31 @@ class Grouper(Parametrizable):
         More than one array can be combined to a dataset before grouping using the `das`  kwargs.
         A new `window` dimension is added if `self.window` is larger than 1.
         If `Grouper.dim` is 'time', but 'prop' is None, the whole array is grouped together.
+
+        When multiple arrays are passed, some of them can be grouped along the same group as self.
+        They are boadcasted, merged to the grouping dataset and regrouped in the output.
         """
         if das:
+            from .utils import broadcast  # pylint: disable=cyclic-import
+
             if da is not None:
                 das[da.name] = da
-            da = xr.Dataset(data_vars=das)
+
+            da = xr.Dataset(
+                data_vars={
+                    name: das.pop(name)
+                    for name in list(das.keys())
+                    if self.dim in das[name].dims
+                }
+            )
+
+            # "Ungroup" the grouped arrays
+            da = da.assign(
+                {
+                    name: broadcast(var, da[self.dim], group=self, interp="nearest")
+                    for name, var in das.items()
+                }
+            )
 
         if self.window > 1:
             da = da.rolling(center=True, **{self.dim: self.window}).construct(
@@ -143,6 +164,12 @@ class Grouper(Parametrizable):
 
         ind = da.indexes[self.dim]
         i = getattr(ind, self.prop)
+
+        if i.dtype != np.int:
+            raise ValueError(
+                f"Index {self.name} is not of type int (rather {ind.dtype}), but {self.__class__.__name__} requires integer indexes."
+            )
+
         interp = (
             (interp or self.interp)
             if not isinstance(interp, str)
@@ -215,6 +242,7 @@ class Grouper(Parametrizable):
                 map(
                     lambda d: (
                         d.chunks is not None
+                        and self.dim in d.dims
                         and len(d.chunks[d.get_axis_num(self.dim)]) > 1
                     ),
                     da.values(),
@@ -241,10 +269,12 @@ class Grouper(Parametrizable):
         if isinstance(out, xr.Dataset):
             for name, outvar in out.data_vars.items():
                 if "_group_apply_reshape" in outvar.attrs:
-                    if outvar.attrs["_group_apply_reshape"] and self.prop is not None:
+                    if self.prop is not None:
                         out[name] = outvar.groupby(self.name).first(
                             skipna=False, keep_attrs=True
                         )
+                    else:
+                        out[name] = out[name].isel({self.dim: 0})
                     del out[name].attrs["_group_apply_reshape"]
 
         # Save input parameters as attributes of output DataArray.
@@ -266,7 +296,7 @@ class Grouper(Parametrizable):
             self.window > 1 and "window" in out.dims
         ):  # On non reducing ops, drop the constructed window
             out = out.isel(window=self.window // 2, drop=True)
-        if self.prop in out.dims and out.chunks is not None:
+        if self.prop in out.dims and bool(out.chunks):
             # Same as above : downstream methods expect only one chunk along the group
             out = out.chunk({self.prop: -1})
 
@@ -282,7 +312,7 @@ def parse_group(func):
 
     @wraps(func)
     def _parse_group(*args, **kwargs):
-        group = kwargs.get("group", default_group)
+        group = kwargs.setdefault("group", default_group)
         if not isinstance(group, Grouper):
             if not isinstance(group, str):
                 dim, *add_dims = group
diff --git a/xclim/sdba/detrending.py b/xclim/sdba/detrending.py
index aa82e7ec..cce24308 100644
--- a/xclim/sdba/detrending.py
+++ b/xclim/sdba/detrending.py
@@ -1,11 +1,16 @@
 """Detrending objects"""
+from typing import Union
+
 import xarray as xr
 
+from .base import Grouper
 from .base import Parametrizable
+from .base import parse_group
 from .utils import ADDITIVE
 from .utils import apply_correction
 from .utils import invert
-from .utils import loffsets
+
+# from .utils import loffsets
 
 
 class BaseDetrend(Parametrizable):
@@ -13,48 +18,80 @@ class BaseDetrend(Parametrizable):
 
     Defines three methods:
 
-    fit(da)     : Compute trend from da and return a new _fitted_ Detrend object.
-    detrend(da) : Return detrended array.
-    retrend(da) : Puts trend back on da.
+    fit(da)      : Compute trend from da and return a new _fitted_ Detrend object.
+    get_trend(da): Return the fitted trend along da's coordinate.
+    detrend(da)  : Return detrended array.
+    retrend(da)  : Puts trend back on da.
 
-    * Subclasses should implement _fit(), _detrend() and _retrend(), not the methods themselves.
-    Only _fit() should store data. _detrend() and _retrend() are meant to be used on any dataarray with the trend computed in fit.
+    * Subclasses should implement _fit() and _get_trend(). Both will be called in a `group.apply()`.
+    `_fit()` is called with the dataarray and str `dim` that indicates the fitting dimension,
+        it should return a dataset that will be set as `.fitds`.
+    `_get_trend()` is called with .fitds broadcasted on the main dim of the input DataArray.
     """
 
-    def __init__(self, **kwargs):
+    @parse_group
+    def __init__(
+        self, *, group: Union[Grouper, str] = "time", kind: str = "+", **kwargs
+    ):
+        """Initialize Detrending object.
+
+        Parameters
+        ----------
+        group : Union[str, Grouper]
+            The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
+            The fit is performed along the group's main dim.
+        kind : {'*', '+'}
+            The way the trend is removed or added, either additive or multiplicative.
+        """
         self.__fitted = False
-        super().__init__(**kwargs)
+        super().__init__(group=group, kind=kind, **kwargs)
 
-    def fit(self, da: xr.DataArray, dim="time"):
+    def fit(self, da: xr.DataArray):
         """Extract the trend of a DataArray along a specific dimension.
 
         Returns a new object storing the fit data that can be used for detrending and retrending.
         """
         new = self.copy()
-        new._fit(da, dim=dim)
-        new._fitted_dim = dim
+        new._set_fitds(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(
+            self._get_trend,
+            {self.group.dim: da[self.group.dim], **self.fit_ds.data_vars},
+            main_only=True,
+        )
+
     def detrend(self, da: xr.DataArray):
         """Removes the previously fitted trend from a DataArray."""
         if not self.__fitted:
             raise ValueError("You must call fit() before detrending.")
-        return self._detrend(da)
+        trend = self.get_trend(da)
+        return self._detrend(da, trend)
 
     def retrend(self, da: xr.DataArray):
         """Puts back the previsouly fitted trend on a DataArray."""
         if not self.__fitted:
             raise ValueError("You must call fit() before retrending")
-        return self._retrend(da)
+        trend = self.get_trend(da)
+        return self._retrend(da, trend)
 
-    def _fit(self, da):
-        raise NotImplementedError
+    def _set_fitds(self, ds):
+        self.fit_ds = ds
+
+    def _detrend(self, da, trend):
+        # Remove trend from series
+        return apply_correction(da, invert(trend, self.kind), self.kind)
 
-    def _detrend(self, da):
+    def _retrend(self, da, trend):
+        # Add trend to series
+        return apply_correction(da, trend, self.kind)
+
+    def _get_trend(self, grpd, dim="time"):
         raise NotImplementedError
 
-    def _retrend(self, da):
+    def _fit(self, da):
         raise NotImplementedError
 
 
@@ -62,26 +99,25 @@ class NoDetrend(BaseDetrend):
     """Convenience class for polymorphism. Does nothing."""
 
     def _fit(self, da, dim=None):
-        pass
+        return da.isel({dim: 0})
 
-    def _detrend(self, da):
+    def _detrend(self, da, trend):
         return da
 
-    def _retrend(self, da):
+    def _retrend(self, da, trend):
         return da
 
 
 class MeanDetrend(BaseDetrend):
-    """Simple detrending removing only the mean from the data, quite similar to normalizing in additive mode."""
+    """Simple detrending removing only the mean from the data, quite similar to normalizing."""
 
     def _fit(self, da, dim="time"):
-        self._mean = da.mean(dim=dim)
+        mean = da.mean(dim=dim)
+        mean.name = "mean"
+        return mean
 
-    def _detrend(self, da):
-        return da - self._mean
-
-    def _retrend(self, da):
-        return da + self._mean
+    def _get_trend(self, grpd, dim="time"):
+        return grpd.mean
 
 
 class PolyDetrend(BaseDetrend):
@@ -90,47 +126,33 @@ class PolyDetrend(BaseDetrend):
 
     Parameters
     ----------
+    group : Union[str, Grouper]
+        The grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
+        The fit is performed along the group's main dim.
+    kind : {'*', '+'}
+        The way the trend is removed or added, either additive or multiplicative.
     degree : int
-      The order of the polynomial to fit.
-    freq : Optional[str]
-      If given, resamples the data to this frequency before computing the trend.
-    kind : {'+', '*'}
-      The way the trend is removed and put back, either additively or multiplicatively.
-
-    Notes
-    -----
-    If freq is used to resample at a lower frequency, make sure the series includes full periods.
+        The order of the polynomial to fit.
+    preserve_mean : bool
+        Whether to preserve the mean when de/re-trending. If True, the trend has its mean
+        removed before it is used.
     """
 
-    def __init__(self, degree=4, freq=None, kind=ADDITIVE, preserve_mean=False):
+    def __init__(self, group="time", kind=ADDITIVE, degree=4, preserve_mean=False):
         super().__init__(
-            degree=degree, freq=freq, kind=kind, preserve_mean=preserve_mean
+            group=group, kind=kind, degree=degree, preserve_mean=preserve_mean
         )
 
     def _fit(self, da, dim="time"):
-        if self.freq is not None:
-            da = da.resample(
-                time=self.freq, label="left", loffset=loffsets[self.freq]
-            ).mean()
-        self._fitds = da.polyfit(dim=dim, deg=self.degree, full=True)
+        return da.polyfit(dim=dim, deg=self.degree)
 
-    def _get_trend(self, da):
+    def _get_trend(self, grpd, dim="time"):
         # Estimate trend over da
-        trend = xr.polyval(
-            coord=da[self._fitted_dim], coeffs=self._fitds.polyfit_coefficients
-        )
+        trend = xr.polyval(coord=grpd[dim], coeffs=grpd.polyfit_coefficients)
 
         if self.preserve_mean:
             trend = apply_correction(
-                trend, invert(trend.mean(dim=self._fitted_dim), self.kind), self.kind
+                trend, invert(trend.mean(dim=dim), self.kind), self.kind
             )
 
         return trend
-
-    def _detrend(self, da):
-        # Remove trend from series
-        return apply_correction(da, invert(self._get_trend(da), self.kind), self.kind)
-
-    def _retrend(self, da):
-        # Add trend to series
-        return apply_correction(da, self._get_trend(da), self.kind)


Test patch:
diff --git a/tests/test_sdba/test_adjustment.py b/tests/test_sdba/test_adjustment.py
index ac54ae4e..4047e5f5 100644
--- a/tests/test_sdba/test_adjustment.py
+++ b/tests/test_sdba/test_adjustment.py
@@ -209,7 +209,7 @@ class TestDQM:
         DQM = DetrendedQuantileMapping(group="time.month")
         DQM.train(ref, hist)
         scen = DQM.adjust(sim, interp="linear")
-        xr.testing.assert_allclose(ref, scen, rtol=1.5e-4)
+        xr.testing.assert_allclose(ref, scen, rtol=1e-3)
 
 
 class TestQDM:
diff --git a/tests/test_sdba/test_base.py b/tests/test_sdba/test_base.py
index ee14b65e..a09206f6 100644
--- a/tests/test_sdba/test_base.py
+++ b/tests/test_sdba/test_base.py
@@ -6,7 +6,6 @@ sdba = pytest.importorskip("xclim.sdba")  # noqa
 
 from xclim.sdba.base import Grouper
 from xclim.sdba.base import Parametrizable
-from xclim.sdba.processing import normalize
 
 
 def test_param_class():
@@ -52,45 +51,94 @@ def test_grouper_get_index(tas_series, group, interp, val90):
     assert indx[90] == val90
 
 
-def test_grouper_apply(tas_series):
+@pytest.mark.parametrize(
+    "group,n",
+    [("time", 1), ("time.month", 12), ("time.week", 52), ("mytime.dayofyear", 366)],
+)
+@pytest.mark.parametrize("use_dask", [True, False])
+def test_grouper_apply(tas_series, use_dask, group, n):
     tas1 = tas_series(np.arange(366), start="2000-01-01")
     tas0 = tas_series(np.zeros(366), start="2000-01-01")
     tas = xr.concat((tas1, tas0), dim="lat")
 
-    grouper = Grouper("time.month")
-    out = grouper.apply("mean", tas)
-    assert out.isel(month=0, lat=0) == 15.0
-    out = normalize(tas, group=grouper)
-
-    grouper = Grouper("time.month", add_dims=["lat"])
+    grouper = Grouper(group)
+    if not group.startswith("time"):
+        tas = tas.rename(time=grouper.dim)
+        tas1 = tas1.rename(time=grouper.dim)
+        tas0 = tas0.rename(time=grouper.dim)
+
+    if use_dask:
+        tas = tas.chunk({"lat": 1, grouper.dim: -1})
+        tas0 = tas1.chunk({grouper.dim: -1})
+        tas1 = tas0.chunk({grouper.dim: -1})
+
+    # Normal monthly mean
+    out_mean = grouper.apply("mean", tas)
+    if grouper.prop:
+        exp = tas.groupby(group).mean()
+    else:
+        exp = tas.mean(dim=grouper.dim)
+    np.testing.assert_array_equal(out_mean, exp)
+
+    # With additionnal dimension included
+    grouper = Grouper(group, add_dims=["lat"])
     out = grouper.apply("mean", tas)
-    assert out.ndim == 1
-    assert out.isel(month=0,) == 7.5
-    assert out.attrs["group"] == "time.month"
-    assert out.attrs["group_compute_dims"] == ["time", "lat"]
+    assert out.ndim == int(grouper.prop is not None)
+    np.testing.assert_array_equal(out, exp.mean("lat"))
+    assert out.attrs["group"] == group
+    assert out.attrs["group_compute_dims"] == [grouper.dim, "lat"]
     assert out.attrs["group_window"] == 1
 
-    grouper = Grouper("time.month", window=5)
-    out = grouper.apply("mean", tas)
-    np.testing.assert_almost_equal(out.isel(month=0, lat=0), 15.32236842)
+    # Additionnal but main_only
+    out = grouper.apply("mean", tas, main_only=True)
+    np.testing.assert_array_equal(out, out_mean)
 
-    tas = tas.chunk({"lat": 1})
+    # With window
+    grouper = Grouper(group, window=5)
     out = grouper.apply("mean", tas)
-    assert out.chunks == ((1, 1), (12,))
+    rolld = tas.rolling({grouper.dim: 5}, center=True).construct(window_dim="window")
+    if grouper.prop:
+        exp = rolld.groupby(group).mean(dim=[grouper.dim, "window"])
+    else:
+        exp = rolld.mean(dim=[grouper.dim, "window"])
+    np.testing.assert_array_equal(out, exp)
+
+    # With function + nongrouping-grouped
+    grouper = Grouper(group)
 
-    out = normalize(tas, group=grouper)
-    assert out.chunks == ((1, 1), (366,))
+    def normalize(grp, dim):
+        return grp / grp.mean(dim=dim)
 
+    normed = grouper.apply(normalize, tas)
+    assert normed.shape == tas.shape
+    if use_dask:
+        assert normed.chunks == ((1, 1), (366,))
+
+    # Mixed output
     def mixed_reduce(grdds, dim=None):
         tas1 = grdds.tas1.mean(dim=dim)
         tas0 = grdds.tas0 / grdds.tas0.mean(dim=dim)
         tas1.attrs["_group_apply_reshape"] = True
         return xr.Dataset(data_vars={"tas1_mean": tas1, "norm_tas0": tas0})
 
-    tas1 = tas1.chunk({"time": -1})
     out = grouper.apply(mixed_reduce, {"tas1": tas1, "tas0": tas0})
-    assert "month" not in out.norm_tas0.dims
-    assert "month" in out.tas1_mean.dims
-
-    assert out.tas1_mean.chunks == ((12,),)
-    assert out.norm_tas0.chunks == ((366,),)
+    if grouper.prop:
+        assert grouper.prop not in out.norm_tas0.dims
+        assert grouper.prop in out.tas1_mean.dims
+
+    if use_dask:
+        assert out.tas1_mean.chunks == (((n,),) if grouper.prop else tuple())
+        assert out.norm_tas0.chunks == ((366,),)
+
+    # Mixed input
+    if grouper.prop:
+
+        def normalize_from_precomputed(grpds, dim=None):
+            return (grpds.tas / grpds.tas1_mean).mean(dim=dim)
+
+        out = grouper.apply(
+            normalize_from_precomputed, {"tas": tas, "tas1_mean": out.tas1_mean}
+        ).isel(lat=0)
+        exp = normed.groupby(group).mean().isel(lat=0)
+        assert grouper.prop in out.dims
+        np.testing.assert_array_equal(out, exp)