More flexible produce_horizon (#483)

### Pull Request Checklist: - [x] This PR addresses an already opened issue (for bug fixes / features) - This PR fixes #xyz - [x] (If applicable) Documentation has been added / updated (for bug fixes / features). - [x] (If applicable) Tests have been added. - [x] This PR does not seem to break the templates. - [x] CHANGELOG.rst has been updated (with summary of main changes). - [x] Link to issue (:issue:`number`) and pull request (:pull:`number`) has been added. ### What kind of change does this PR introduce? * You can now skip the computation of indicators, if you've already done that step and just want the subsetting/climatological_op utilities of `produce_horizon`. * You can now control the `op` in `produce_horizon`. * Fixed a bug in `compute_indicators` where the `cat:variable` was determined too early and only gave the first indicator even when multiple were computed. ### Does this PR introduce a breaking change? - `indicators` was moved to a `kwarg`, since it is now optional. ### Other information:
Ouranosinc · Nov 4, 2024 · 6abe87b · 6abe87b
2 parents 7ab1cf9 + b40ce24
commit 6abe87b
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 21 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -12,17 +12,19 @@ New features and enhancements
 * New function ``xs.get_warming_level_from_period`` to get the warming level associated with a given time horizon. (:pull:`474`).
 * Added ability to skip whole folders to ``xs.parse_directory`` with argument ``skip_dirs``. (:pull:`478`, :pull:`479`).
 * `diagnostics.measures_improvement` now accepts `dim`, which specifies `dimension(s)` on which the proportion of improved pixels are computed. (:pull:`416`)
-
+* The argument `indicators` in ``xs.produce_horizon`` is now optional. Added an argument `op` to control the climatological operation. (:pull:`483`).
 
 Breaking changes
 ^^^^^^^^^^^^^^^^
 * ``xs.get_warming_level`` has been renamed to ``xs.get_period_from_warming_level``. Its argument `return_horizon` was reversed and renamed `return_central_year` (:pull:`474`).
 * Removed support for the deprecated `xclim` function `change_significance` in `ensemble_stats`. (:pull:`482`).
+* The argument `indicators` in ``xs.produce_horizon`` is no longer positional. (:pull:`483`).
 
 Bug fixes
 ^^^^^^^^^
 * ``xs.io.save_to_table`` now correctly handles the case where the input is a `DataArray` or a `Dataset` with a single variable. (:pull:`473`).
 * Fixed a bug in ``xs.utils.change_units`` where the original dataset was also getting modified. (:pull:`482`).
+* Fixed a bug in ``xs.compute_indicators`` where the `cat:variable` attribute was not correctly set. (:pull:`483`).
 
 Internal changes
 ^^^^^^^^^^^^^^^^

diff --git a/src/xscen/aggregate.py b/src/xscen/aggregate.py
@@ -933,16 +933,18 @@ def spatial_mean(  # noqa: C901
 @parse_config
 def produce_horizon(  # noqa: C901
     ds: xr.Dataset,
+    *,
     indicators: (
         str
         | os.PathLike
         | Sequence[Indicator]
         | Sequence[tuple[str, Indicator]]
         | ModuleType
-    ),
-    *,
+        | None
+    ) = None,
     periods: list[str] | list[list[str]] | None = None,
     warminglevels: dict | None = None,
+    op: str | dict = "mean",
     to_level: str | None = "horizons",
 ) -> xr.Dataset:
     """
@@ -957,7 +959,8 @@ def produce_horizon(  # noqa: C901
     ----------
     ds : xr.Dataset
         Input dataset with a time dimension.
-    indicators :  str | os.PathLike | Sequence[Indicator] | Sequence[Tuple[str, Indicator]] | ModuleType
+        If 'indicators' is None, the dataset should contain the precomputed indicators.
+    indicators :  str | os.PathLike | Sequence[Indicator] | Sequence[Tuple[str, Indicator]] | ModuleType, optional
         Indicators to compute. It will be passed to the `indicators` argument of `xs.compute_indicators`.
     periods : list of str or list of lists of str, optional
         Either [start, end] or list of [start_year, end_year] for the period(s) to be evaluated.
@@ -966,6 +969,8 @@ def produce_horizon(  # noqa: C901
         Dictionary of arguments to pass to `py:func:xscen.subset_warming_level`.
         If 'wl' is a list, the function will be called for each value and produce multiple horizons.
         If both periods and warminglevels are None, the full time series will be used.
+    op : str or dict
+        Operation to perform over the time dimension. See `py:func:xscen.climatological_op` for details. Default is 'mean'.
     to_level : str, optional
         The processing level to assign to the output.
         If there is only one horizon, you can use "{wl}", "{period0}" and "{period1}" in the string to dynamically
@@ -1019,23 +1024,26 @@ def produce_horizon(  # noqa: C901
             ds_sub = subset_warming_level(ds, **period)
 
         if ds_sub is not None:
-            # compute indicators
-            ind_dict = compute_indicators(
-                ds=(
-                    ds_sub.squeeze(dim="warminglevel")
-                    if "warminglevel" in ds_sub.dims
-                    else ds_sub
-                ),
-                indicators=indicators,
-            )
+            if indicators is not None:
+                # compute indicators
+                ind_dict = compute_indicators(
+                    ds=(
+                        ds_sub.squeeze(dim="warminglevel")
+                        if "warminglevel" in ds_sub.dims
+                        else ds_sub
+                    ),
+                    indicators=indicators,
+                )
+            else:
+                ind_dict = {"skipped": ds_sub}
 
             # Compute the window-year mean
             ds_merge = xr.Dataset()
             for freq, ds_ind in ind_dict.items():
                 if freq != "fx":
                     ds_mean = climatological_op(
                         ds_ind,
-                        op="mean",  # ToDo: make op an argument of produce_horizon
+                        op=op,
                         rename_variables=False,
                         horizons_as_dim=True,
                     ).drop_vars("time")
@@ -1052,6 +1060,8 @@ def produce_horizon(  # noqa: C901
                 if "warminglevel" in ds_mean.coords:
                     wl = np.array([ds_mean["warminglevel"].item()])
                     wl_attrs = ds_mean["warminglevel"].attrs
+                    if "warminglevel" in ds_mean.dims:
+                        ds_mean = ds_mean.squeeze("warminglevel")
                     ds_mean = ds_mean.drop_vars("warminglevel")
                     ds_mean["horizon"] = wl
                     ds_mean["horizon"].attrs.update(wl_attrs)
@@ -1070,7 +1080,7 @@ def produce_horizon(  # noqa: C901
                 [
                     all(
                         [
-                            out[0][v].attrs[attr] == out[i][v].attrs[attr]
+                            out[0][v].attrs.get(attr) == out[i][v].attrs.get(attr)
                             for i in range(1, len(out))
                         ]
                     )

diff --git a/src/xscen/indicators.py b/src/xscen/indicators.py
@@ -136,8 +136,7 @@ def compute_indicators(  # noqa: C901
         If True, cut the time axis to be within the same years as the input.
         This is mostly useful for frequencies that do not start in January, such as QS-DEC.
         In that instance, `xclim` would start on previous_year-12-01 (DJF), with a NaN.
-        `restrict_years` will cut that first timestep.
-        This should have no effect on YS and MS indicators.
+        `restrict_years` will cut that first timestep. This should have no effect on YS and MS indicators.
     to_level : str, optional
         The processing level to assign to the output.
         If None, the processing level of the inputs is preserved.
@@ -262,11 +261,7 @@ def compute_indicators(  # noqa: C901
         key = freq
         if key not in out_dict:
             out_dict[key] = out
-            # TODO: Double-check History, units, attrs, add missing variables (grid_mapping), etc.
             out_dict[key].attrs = ds.attrs
-            out_dict[key].attrs["cat:variable"] = parse_from_ds(
-                out_dict[key], ["variable"]
-            )["variable"]
             out_dict[key].attrs["cat:xrfreq"] = freq
             out_dict[key].attrs["cat:frequency"] = CV.xrfreq_to_frequency(freq, None)
             if to_level is not None:
@@ -275,6 +270,10 @@ def compute_indicators(  # noqa: C901
         else:
             for v in out.data_vars:
                 out_dict[key][v] = out[v]
+    for key in out_dict:
+        out_dict[key].attrs["cat:variable"] = parse_from_ds(
+            out_dict[key], ["variable"]
+        )["variable"]
 
     return out_dict
 

diff --git a/tests/test_aggregate.py b/tests/test_aggregate.py
@@ -340,6 +340,39 @@ def test_single(self):
         assert len(out.horizon) == 1
         np.testing.assert_array_equal(out.horizon, ["1982-1988"])
 
+    def test_op(self):
+        ds = self.ds.copy()
+        ds.tas.loc["1995-01-01":"1995-12-31"] = 2
+        out = xs.produce_horizon(
+            ds, indicators=self.yaml_file, periods=[["1995", "2007"]]
+        )
+        np.testing.assert_array_almost_equal(out["tg_min"], [14 / 13])
+        out2 = xs.produce_horizon(
+            ds, indicators=self.yaml_file, periods=[["1995", "2007"]], op="max"
+        )
+        np.testing.assert_array_almost_equal(out2["tg_min"], [2])
+
+    def test_precomputed(self):
+        ds = timeseries(
+            np.arange(30),
+            variable="tas",
+            start="1981-01-01",
+            freq="YS-JAN",
+            as_dataset=True,
+        )
+        ds = ds.rename({"tas": "tg_min"})
+        ds.attrs["cat:source"] = "CanESM5"
+        ds.attrs["cat:experiment"] = "ssp585"
+        ds.attrs["cat:member"] = "r1i1p1f1"
+        ds.attrs["cat:mip_era"] = "CMIP6"
+
+        out = xs.produce_horizon(ds, periods=[["1981", "1985"], ["1991", "1995"]])
+        assert len(out.horizon) == 2
+        assert "time" not in out.dims
+        np.testing.assert_array_almost_equal(
+            out["tg_min"], [np.mean(np.arange(5)), np.mean(np.arange(10, 15))]
+        )
+
     def test_warminglevel_in_ds(self):
         ds = self.ds.copy().expand_dims({"warminglevel": ["+1Cvs1850-1900"]})
         out = xs.produce_horizon(

diff --git a/tests/test_indicators.py b/tests/test_indicators.py
@@ -77,6 +77,10 @@ def test_output(self, periods):
             xrfreq in ind_dict[xrfreq].attrs["cat:xrfreq"] for xrfreq in ind_dict.keys()
         )
         assert all(v in ind_dict["YS-JAN"] for v in ["tg_min", "growing_degree_days"])
+        assert all(
+            v in ind_dict["YS-JAN"].attrs["cat:variable"]
+            for v in ["tg_min", "growing_degree_days"]
+        )
         if periods is None:
             assert "time" not in ind_dict["fx"].dims
             assert len(ind_dict["YS-JAN"].time) == 2