Option to output datasets

Ouranosinc · Jan 25, 2024 · a1e3017 · a1e3017
1 parent 53267be
commit a1e3017
Show file tree

Hide file tree

Showing 7 changed files with 67 additions and 24 deletions.
diff --git a/.github/workflows/label-on-approval.yml b/.github/workflows/label-on-approval.yml
@@ -74,7 +74,7 @@ jobs:
           (steps.fc.outputs.comment-id == '') &&
           (!contains(github.event.pull_request.labels.*.name, 'approved')) &&
           (github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name)
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}
@@ -85,7 +85,7 @@ jobs:
       - name: Update comment
         if: |
           contains(github.event.pull_request.labels.*.name, 'approved')
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}

diff --git a/.github/workflows/testdata-version.yml b/.github/workflows/testdata-version.yml
@@ -60,7 +60,7 @@ jobs:
             core.setFailed('Configured `xclim-testdata` tag is not `latest`.')
       - name: Update Failure Comment
         if: ${{ failure() }}
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}
@@ -76,7 +76,7 @@ jobs:
           edit-mode: replace
       - name: Update Success Comment
         if: ${{ success() }}
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}

diff --git a/.github/workflows/workflow-warning.yml b/.github/workflows/workflow-warning.yml
@@ -45,7 +45,7 @@ jobs:
           (steps.fc.outputs.comment-id == '') &&
           (!contains(github.event.pull_request.labels.*.name, 'approved')) &&
           (github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name)
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}
@@ -57,7 +57,7 @@ jobs:
       - name: Update comment
         if: |
           contains(github.event.pull_request.labels.*.name, 'approved')
-        uses: peter-evans/create-or-update-comment@23ff15729ef2fc348714a3bb66d2f655ca9066f2 # v3.1.0
+        uses: peter-evans/create-or-update-comment@71345be0265236311c031f5c7866368bd1eff043 # v4.0.0
         with:
           comment-id: ${{ steps.fc.outputs.comment-id }}
           issue-number: ${{ github.event.pull_request.number }}

diff --git a/CHANGES.rst b/CHANGES.rst
@@ -18,6 +18,7 @@ New features and enhancements
 * Validate YAML indicators description before trying to build module. (:issue:`1523`, :issue:`1595`, :pull:`1560`, :pull:`1596`, :pull:`1600`).
 * 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`).
 
 Breaking changes
 ^^^^^^^^^^^^^^^^

diff --git a/tests/test_indicators.py b/tests/test_indicators.py
@@ -198,6 +198,28 @@ def test_keep_attrs(tasmin_series, tasmax_series, xcopt, xropt, exp):
     assert "bing" not in tg.attrs
 
 
+def test_as_dataset(tasmax_series, tasmin_series):
+    tx = tasmax_series(np.arange(360.0))
+    tn = tasmin_series(np.arange(360.0))
+    tx.attrs.update(something="blabla", bing="bang", foo="bar")
+    tn.attrs.update(something="blabla", bing="bong")
+    dsin = xr.Dataset({"tasmax": tx, "tasmin": tn}, attrs={"fou": "barre"})
+    with xclim.set_options(keep_attrs=True, as_dataset=True):
+        dsout = multiOptVar(ds=dsin)
+    assert isinstance(dsout, xr.Dataset)
+    assert dsout.attrs["fou"] == "barre"
+    assert dsout.multiopt.attrs.get("something") == "blabla"
+
+
+def test_as_dataset_multi(tas_series):
+    tg = tas_series(np.arange(360.0))
+    with xclim.set_options(as_dataset=True):
+        dsout = multiTemp(tas=tg, freq="YS")
+    assert isinstance(dsout, xr.Dataset)
+    assert "tmin" in dsout.data_vars
+    assert "tmax" in dsout.data_vars
+
+
 def test_opt_vars(tasmin_series, tasmax_series):
     tn = tasmin_series(np.zeros(365))
     tx = tasmax_series(np.zeros(365))

diff --git a/xclim/core/indicator.py b/xclim/core/indicator.py
@@ -143,6 +143,7 @@
     read_locale_file,
 )
 from .options import (
+    AS_DATASET,
     CHECK_MISSING,
     KEEP_ATTRS,
     METADATA_LOCALES,
@@ -809,7 +810,7 @@ def __call__(self, *args, **kwds):
         if self._version_deprecated:
             self._show_deprecation_warning()  # noqa
 
-        das, params = self._parse_variables_from_call(args, kwds)
+        das, params, dsattrs = self._parse_variables_from_call(args, kwds)
 
         if OPTIONS[KEEP_ATTRS] is True or (
             OPTIONS[KEEP_ATTRS] == "xarray"
@@ -881,6 +882,20 @@ def __call__(self, *args, **kwds):
             out.attrs.update(attrs)
             out.name = var_name
 
+        if OPTIONS[AS_DATASET]:
+            out = Dataset({o.name: o for o in outs})
+            if OPTIONS[KEEP_ATTRS] is True or (
+                OPTIONS[KEEP_ATTRS] == "xarray"
+                and xarray.core.options._get_keep_attrs(False)
+            ):
+                out.attrs.update(dsattrs)
+            out.attrs["history"] = update_history(
+                self._history_string(das, params),
+                out,
+                new_name=self.identifier,
+            )
+            return out
+
         # Return a single DataArray in case of single output, otherwise a tuple
         if self.n_outs == 1:
             return outs[0]
@@ -912,7 +927,9 @@ def _parse_variables_from_call(self, args, kwds) -> tuple[OrderedDict, dict]:
             else:
                 params[name] = param.value
 
-        return das, params
+        ds = ba.arguments.get("ds")
+        dsattrs = ds.attrs if ds is not None else {}
+        return das, params, dsattrs
 
     def _assign_named_args(self, ba):
         """Assign inputs passed as strings from ds."""
@@ -1065,28 +1082,24 @@ def _update_attrs(
             if "cell_methods" in out:
                 attrs["cell_methods"] += " " + out.pop("cell_methods")
 
-        # Use of OrderedDict to ensure inputs (das) get listed before parameters (args).
-        # In the history attr, call signature will be all keywords and might be in a
-        # different order than the real function (but order doesn't really matter with keywords).
-        kwargs = OrderedDict(**das)
-        for k, v in args.items():
-            if self._all_parameters[k].injected:
-                continue
-            if self._all_parameters[k].kind == InputKind.KWARGS:
-                kwargs.update(**v)
-            elif self._all_parameters[k].kind != InputKind.DATASET:
-                kwargs[k] = v
-
         attrs["history"] = update_history(
-            self._history_string(**kwargs),
+            self._history_string(das, args),
             new_name=out.get("var_name"),
             **das,
         )
 
         attrs.update(out)
         return attrs
 
-    def _history_string(self, **kwargs):
+    def _history_string(self, das, params):
+        kwargs = dict(**das)
+        for k, v in params.items():
+            if self._all_parameters[k].injected:
+                continue
+            if self._all_parameters[k].kind == InputKind.KWARGS:
+                kwargs.update(**v)
+            elif self._all_parameters[k].kind != InputKind.DATASET:
+                kwargs[k] = v
         return gen_call_string(self._registry_id, **kwargs)
 
     @staticmethod
@@ -1396,7 +1409,7 @@ def __init__(self, **kwds):
 
         super().__init__(**kwds)
 
-    def _history_string(self, **kwargs):
+    def _history_string(self, das, params):
         if self.missing == "from_context":
             missing = OPTIONS[CHECK_MISSING]
         else:
@@ -1408,7 +1421,7 @@ def _history_string(self, **kwargs):
             if mopts:
                 opt_str += f", missing_options={mopts}"
 
-        return super()._history_string(**kwargs) + opt_str
+        return super()._history_string(das, params) + opt_str
 
     def _get_missing_freq(self, params):
         """Return the resampling frequency to be used in the missing values check."""

diff --git a/xclim/core/options.py b/xclim/core/options.py
@@ -23,6 +23,7 @@
 SDBA_EXTRA_OUTPUT = "sdba_extra_output"
 SDBA_ENCODE_CF = "sdba_encode_cf"
 KEEP_ATTRS = "keep_attrs"
+AS_DATASET = "as_dataset"
 
 MISSING_METHODS: dict[str, Callable] = {}
 
@@ -36,6 +37,7 @@
     SDBA_EXTRA_OUTPUT: False,
     SDBA_ENCODE_CF: False,
     KEEP_ATTRS: "xarray",
+    AS_DATASET: False,
 }
 
 _LOUDNESS_OPTIONS = frozenset(["log", "warn", "raise"])
@@ -67,6 +69,7 @@ def _valid_missing_options(mopts):
     SDBA_EXTRA_OUTPUT: lambda opt: isinstance(opt, bool),
     SDBA_ENCODE_CF: lambda opt: isinstance(opt, bool),
     KEEP_ATTRS: _KEEP_ATTRS_OPTIONS.__contains__,
+    AS_DATASET: lambda opt: isinstance(opt, bool),
 }
 
 
@@ -175,8 +178,12 @@ class set_options:
     keep_attrs : bool or str
         Controls attributes handling in indicators. If True, attributes from all inputs are merged
         using the `drop_conflicts` strategy and then updated with xclim-provided attributes.
+        If ``as_dataset`` is also True and a dataset was passed to the ``ds`` argument of the Indicator,
+        the dataset's attributes are copied to the indicator's output.
         If False, attributes from the inputs are ignored. If "xarray", xclim will use xarray's `keep_attrs` option.
         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``.
 
     Examples
     --------