Fix map_blocks for DataArray with cftime and sdba_encode_cf

CHANGES.rst

@@ -10,10 +10,15 @@ Announcements
 ^^^^^^^^^^^^^
 * `xclim` has migrated its development branch name from `master` to `main`. (:issue:`1667`, :pull:`1669`).
 
+Bug fixes
+^^^^^^^^^
+* Fixed an bug in sdba's ``map_groups`` that prevented passing DataArrays with cftime coordinates if the ``sdba_encode_cf`` option was True. (:issue:`1673`, :pull:`1674`).
+
 Internal changes
 ^^^^^^^^^^^^^^^^
 * Added "doymin" and "doymax" to the possible operations of ``generic.stats``. Fixed a warning issue when ``op`` was "integral". (:pull:`1672`).
 
+
 v0.48.2 (2024-02-26)
 --------------------
 Contributors to this version: Juliette Lavoie (:user:`juliettelavoie`).

tests/test_sdba/test_base.py

@@ -6,6 +6,7 @@
 import pytest
 import xarray as xr
 
+from xclim import set_options
 from xclim.sdba.base import Grouper, Parametrizable, map_blocks, map_groups
 
 
@@ -234,3 +235,17 @@ def func(ds, *, group, lon=None):
 
         with pytest.raises(ValueError, match="cannot be chunked"):
             func(xr.Dataset(dict(tas=tas)), group="time")
+
+    @pytest.mark.parametrize("use_dask", [True, False])
+    def test_dataarray_cfencode(self, use_dask, open_dataset):
+        ds = open_dataset("sdba/CanESM2_1950-2100.nc")
+        if use_dask:
+            ds = ds.chunk()
+
+        @map_blocks(reduces=["location"], data=[])
+        def func(ds, *, group):
+            d = ds.mean("location")
+            return d.rename("data").to_dataset()
+
+        with set_options(sdba_encode_cf=True):
+            func(ds.convert_calendar("noleap").tasmax, group=Grouper("time"))

xclim/sdba/base.py

@@ -591,7 +591,6 @@ def _map_blocks(ds, **kwargs):  # noqa: C901
                         f"Dimension {dim} is meant to be added by the "
                         "computation but it is already on one of the inputs."
                     )
-
             if uses_dask(ds):
                 # Use dask if any of the input is dask-backed.
                 chunks = (
@@ -673,9 +672,12 @@ def _map_blocks(ds, **kwargs):  # noqa: C901
             if OPTIONS[SDBA_ENCODE_CF]:
                 ds = ds.copy()
                 # Optimization to circumvent the slow pickle.dumps(cftime_array)
-                for name, crd in ds.coords.items():
-                    if xr.core.common._contains_cftime_datetimes(crd.variable):  # noqa
-                        ds[name] = xr.conventions.encode_cf_variable(crd.variable)
+                # List of the keys to avoid changing the coords dict while iterating over it.
+                for crd in list(ds.coords.keys()):
+                    if xr.core.common._contains_cftime_datetimes(
+                        ds[crd].variable
+                    ):  # noqa
+                        ds[crd] = xr.conventions.encode_cf_variable(ds[crd].variable)
 
             def _call_and_transpose_on_exit(dsblock, **f_kwargs):
                 """Call the decorated func and transpose to ensure the same dim order as on the template."""