Block-mapped resample with the help of flox

CHANGELOG.rst

@@ -19,6 +19,7 @@ New features and enhancements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 * New generic ``xclim.indices.generic.spell_mask`` that returns a mask of which days are part of a spell. Supports multivariate conditions and weights. Used in new generic index ``xclim.indices.generic.bivariate_spell_length_statistics`` that extends ``spell_length_statistics`` to two variables. (:pull:`1885`).
 * Indicator parameters can now be assigned a new name, different from the argument name in the compute function. (:pull:`1885`).
+* New global option ``resample_map_blocks`` to wrap all ``resample().map()`` code inside a ``xr.map_blocks`` to lower the number of dask tasks. Uses utility ``xclim.indices.helpers.resample_map`` and requires ``flox`` to ensure the chunking allows such block-mapping. Defaults to False. (:pull:`1848`).
 
 Bug fixes
 ^^^^^^^^^

environment.yml

@@ -11,6 +11,7 @@ dependencies:
  - click >=8.1
  - dask >=2.6.0
  - filelock >=3.14.0
+ - flox >= 0.9
  - jsonpickle >=3.1.0
  - numba >=0.54.1
  - numpy >=1.23.0

pyproject.toml

@@ -111,7 +111,7 @@ docs = [
  "sphinxcontrib-bibtex",
  "sphinxcontrib-svg2pdfconverter[Cairosvg]"
 ]
-extras = ["fastnanquantile >=0.0.2", "POT >=0.9.4"]
+extras = ["fastnanquantile >=0.0.2", "flox >=0.9", "POT >=0.9.4"]
 all = ["xclim[dev]", "xclim[docs]", "xclim[extras]"]
 
 [project.scripts]

tests/test_helpers.py

@@ -5,8 +5,11 @@
 import pytest
 import xarray as xr
 
+from xclim.core.options import set_options
 from xclim.core.units import convert_units_to
+from xclim.core.utils import uses_dask
 from xclim.indices import helpers
+from xclim.testing.helpers import assert_lazy
 
 
 @pytest.mark.parametrize("method,rtol", [("spencer", 5e3), ("simple", 1e2)])
@@ -132,3 +135,45 @@ def test_cosine_of_solar_zenith_angle():
  ]
  )
  np.testing.assert_allclose(cza[:4, :], exp_cza, rtol=1e-3)
+
+
+def _test_function(da, op, dim):
+ return getattr(da, op)(dim)
+
+
+@pytest.mark.parametrize(
+ ["in_chunks", "exp_chunks"], [(60, 6 * (2,)), (30, 12 * (1,)), (-1, (12,))]
+)
+def test_resample_map(tas_series, in_chunks, exp_chunks):
+ pytest.importorskip("flox")
+ tas = tas_series(365 * [1]).chunk(time=in_chunks)
+ with assert_lazy:
+ out = helpers.resample_map(
+ tas, "time", "MS", lambda da: da.mean("time"), map_blocks=True
+ )
+ assert out.chunks[0] == exp_chunks
+ out.load() # Trigger compute to see if it actually works
+
+
+def test_resample_map_dataset(tas_series, pr_series):
+ pytest.importorskip("flox")
+ tas = tas_series(3 * 365 * [1], start="2000-01-01").chunk(time=365)
+ pr = pr_series(3 * 365 * [1], start="2000-01-01").chunk(time=365)
+ ds = xr.Dataset({"pr": pr, "tas": tas})
+ with set_options(resample_map_blocks=True):
+ with assert_lazy:
+ out = helpers.resample_map(
+ ds,
+ "time",
+ "YS",
+ lambda da: da.mean("time"),
+ )
+ assert out.chunks["time"] == (1, 1, 1)
+ out.load()
+
+
+def test_resample_map_passthrough(tas_series):
+ tas = tas_series(365 * [1])
+ with assert_lazy:
+ out = helpers.resample_map(tas, "time", "MS", lambda da: da.mean("time"))
+ assert not uses_dask(out)

tests/test_indices.py

@@ -1428,13 +1428,24 @@ def test_1d(self, tasmax_series, thresh, window, op, expected):
  def test_resampling_order(self, tasmax_series, resample_before_rl, expected):
  a = np.zeros(365)
  a[5:35] = 31
- tx = tasmax_series(a + K2C)
+ tx = tasmax_series(a + K2C).chunk()
 
  hsf = xci.hot_spell_frequency(
  tx, resample_before_rl=resample_before_rl, freq="MS"
- )
+ ).load()
  assert hsf[1] == expected
 
+ @pytest.importorskip("flox")
+ @pytest.mark.parametrize("resample_map", [True, False])
+ def test_resampling_map(self, tasmax_series, resample_map):
+ a = np.zeros(365)
+ a[5:35] = 31
+ tx = tasmax_series(a + K2C).chunk()
+
+ with set_options(resample_map_blocks=resample_map):
+ hsf = xci.hot_spell_frequency(tx, resample_before_rl=True, freq="MS").load()
+ assert hsf[1] == 1
+
 
 class TestHotSpellMaxLength:
  @pytest.mark.parametrize(
@@ -1708,10 +1719,10 @@ def test_run_start_at_0(self, pr_series):
  def test_resampling_order(self, pr_series, resample_before_rl, expected):
  a = np.zeros(365) + 10
  a[5:35] = 0
- pr = pr_series(a)
+ pr = pr_series(a).chunk()
  out = xci.maximum_consecutive_dry_days(
  pr, freq="ME", resample_before_rl=resample_before_rl
- )
+ ).load()
  assert out[0] == expected
 
 

xclim/core/indicator.py

@@ -165,6 +165,7 @@
  infer_kind_from_parameter,
  is_percentile_dataarray,
  load_module,
+ split_auxiliary_coordinates,
 )
 
 # Indicators registry
@@ -1446,13 +1447,12 @@ def _postprocess(self, outs, das, params):
  # Reduce by or and broadcast to ensure the same length in time
  # When indexing is used and there are no valid points in the last period, mask will not include it
  mask = reduce(np.logical_or, miss)
- if (
- isinstance(mask, DataArray)
- and "time" in mask.dims
- and mask.time.size < outs[0].time.size
- ):
- mask = mask.reindex(time=outs[0].time, fill_value=True)
- outs = [out.where(np.logical_not(mask)) for out in outs]
+ if isinstance(mask, DataArray): # mask might be a bool in some cases
+ if "time" in mask.dims and mask.time.size < outs[0].time.size:
+ mask = mask.reindex(time=outs[0].time, fill_value=True)
+ # Remove any aux coord to avoid any unwanted dask computation in the alignment within "where"
+ mask, _ = split_auxiliary_coordinates(mask)
+ outs = [out.where(~mask) for out in outs]
 
  return outs
 

xclim/core/options.py

@@ -25,6 +25,7 @@
 SDBA_ENCODE_CF = "sdba_encode_cf"
 KEEP_ATTRS = "keep_attrs"
 AS_DATASET = "as_dataset"
+MAP_BLOCKS = "resample_map_blocks"
 
 MISSING_METHODS: dict[str, Callable] = {}
 
@@ -39,6 +40,7 @@
  SDBA_ENCODE_CF: False,
  KEEP_ATTRS: "xarray",
  AS_DATASET: False,
+ MAP_BLOCKS: False,
 }
 
 _LOUDNESS_OPTIONS = frozenset(["log", "warn", "raise"])
@@ -71,6 +73,7 @@ def _valid_missing_options(mopts):
  SDBA_ENCODE_CF: lambda opt: isinstance(opt, bool),
  KEEP_ATTRS: _KEEP_ATTRS_OPTIONS.__contains__,
  AS_DATASET: lambda opt: isinstance(opt, bool),
+ MAP_BLOCKS: lambda opt: isinstance(opt, bool),
 }
 
 
@@ -185,6 +188,9 @@ class set_options:
  Note that xarray's "default" is equivalent to False. Default: ``"xarray"``.
  as_dataset : bool
  If True, indicators output datasets. If False, they output DataArrays. Default :``False``.
+ resample_map_blocks: bool
+ If True, some indicators will wrap their resampling operations with `xr.map_blocks`, using :py:func:`xclim.indices.helpers.resample_map`.
+ This requires `flox` to be installed in order to ensure the chunking is appropriate.git
 
  Examples
  --------

xclim/core/utils.py

@@ -758,3 +758,44 @@ def _chunk_like(*inputs: xr.DataArray | xr.Dataset, chunks: dict[str, int] | Non
  da.chunk(**{d: c for d, c in chunks.items() if d in da.dims})
  )
  return tuple(outputs)
+
+
+def split_auxiliary_coordinates(
+ obj: xr.DataArray | xr.Dataset,
+) -> tuple[xr.DataArray | xr.Dataset, xr.Dataset]:
+ """Split auxiliary coords from the dataset.
+
+ An auxiliary coordinate is a coordinate variable that does not define a dimension and thus is not necessarily needed for dataset alignment.
+ Any coordinate that has a name different than its dimension(s) is flagged as auxiliary. All scalar coordinates are flagged as auxiliary.
+
+ Parameters
+ ----------
+ obj : DataArray or Dataset
+ Xarray object
+
+ Returns
+ -------
+ clean_obj : DataArray or Dataset
+ Same as `obj` but without any auxiliary coordinate.
+ aux_coords : Dataset
+ The auxiliary coordinates as a dataset. Might be empty.
+
+ Note
+ ----
+ This is useful to circumvent xarray's alignment checks that will sometimes look the auxiliary coordinate's data, which can trigger
+ unwanted dask computations.
+
+ The auxiliary coordinates can be merged back with the dataset with
+ :py:meth:`xarray.Dataset.assign_coords` or :py:meth:`xarray.DataArray.assign_coords`.
+
+ >>> # xdoctest: +SKIP
+ >>> clean, aux = split_auxiliary_coordinates(ds)
+ >>> merged = clean.assign_coords(da.coords)
+ >>> merged.identical(ds) # True
+ """
+ aux_crd_names = [
+ nm for nm, crd in obj.coords.items() if len(crd.dims) != 1 or crd.dims[0] != nm
+ ]
+ aux_crd_ds = obj.coords.to_dataset()[aux_crd_names]
+ clean_obj = obj.drop_vars(aux_crd_names)
+ return clean_obj, aux_crd_ds

xclim/indices/_threshold.py

@@ -29,6 +29,7 @@
  spell_length_statistics,
  threshold_count,
 )
+from xclim.indices.helpers import resample_map
 
 # Frequencies : YS: year start, QS-DEC: seasons starting in december, MS: month start
 # See http://pandas.pydata.org/pandas-docs/stable/timeseries.html#offset-aliases
@@ -1491,12 +1492,17 @@ def last_spring_frost(
  thresh = convert_units_to(thresh, tasmin)
  cond = compare(tasmin, op, thresh, constrain=("<", "<="))
 
- out = cond.resample(time=freq).map(
+ out = resample_map(
+ cond,
+ "time",
+ freq,
  rl.last_run_before_date,
- window=window,
- date=before_date,
- dim="time",
- coord="dayofyear",
+ map_kwargs=dict(
+ window=window,
+ date=before_date,
+ dim="time",
+ coord="dayofyear",
+ ),
  )
  out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(tasmin))
  return out
@@ -1662,11 +1668,12 @@ def first_snowfall(
  thresh = convert_units_to(thresh, prsn, context="hydro")
  cond = prsn >= thresh
 
- out = cond.resample(time=freq).map(
+ out = resample_map(
+ cond,
+ "time",
+ freq,
  rl.first_run,
- window=1,
- dim="time",
- coord="dayofyear",
+ map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
  )
  out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(prsn))
  return out
@@ -1717,11 +1724,12 @@ def last_snowfall(
  thresh = convert_units_to(thresh, prsn, context="hydro")
  cond = prsn >= thresh
 
- out = cond.resample(time=freq).map(
+ out = resample_map(
+ cond,
+ "time",
+ freq,
  rl.last_run,
- window=1,
- dim="time",
- coord="dayofyear",
+ map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
  )
  out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(prsn))
  return out
@@ -3097,7 +3105,7 @@ def _exceedance_date(grp):
  never_reached_val = never_reached
  return xarray.where((cumsum <= sum_thresh).all("time"), never_reached_val, out)
 
- dded = c.clip(0).resample(time=freq).map(_exceedance_date)
+ dded = resample_map(c.clip(0), "time", freq, _exceedance_date)
  dded = dded.assign_attrs(
  units="", is_dayofyear=np.int32(1), calendar=get_calendar(tas)
  )