Select the set of indicators for available variables

HISTORY.rst

@@ -12,6 +12,7 @@ Announcements
 
 New features and enhancements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+* New function ``xscen.indicators.select_inds_for_avail_vars`` to filter the indicators that can be calculated with the variables available in a ``xarray.Dataset``. (:pull:`xyz`).
 * Replaced aggregation function ``climatological_mean()`` with ``climatological_op()`` offering more types of operations to aggregate over climate periods. (:pull:`290`)
 * Added the ability to search for simulations that reach a given warming level. (:pull:`251`).
 * ``xs.spatial_mean`` now accepts the ``region="global"`` keyword to perform a global average (:issue:`94`, :pull:`260`).

tests/test_indicators.py

@@ -150,3 +150,59 @@ def test_as_jul(self, restrict_years):
             assert len(out.time) == 4
             assert out.time[0].dt.strftime("%Y-%m-%d").item() == "2000-07-01"
             assert out.time[-1].dt.strftime("%Y-%m-%d").item() == "2003-07-01"
+
+    def test_select_inds_for_avail_vars(self):
+        # Test that select_inds_for_avail_vars filters a list of indicators to only those
+        # that can be computed with the variables available in a dataset.
+        ds = xs.testing.datablock_3d(
+            np.ones((365, 3, 3)),
+            variable="tas",
+            x="lon",
+            x_start=-75,
+            y="lat",
+            y_start=45,
+            x_step=1,
+            y_step=1,
+            start="2001-01-01",
+            freq="D",
+            units="K",
+            as_dataset=True,
+        )
+        indicators = [
+            xclim.core.indicator.Indicator.from_dict(
+                data={"base": "tg_min", "parameters": {"freq": "QS-DEC"}},
+                identifier="tg_min_qs",
+                module="tests",
+            ),
+            xclim.core.indicator.Indicator.from_dict(
+                data={"base": "days_over_precip_thresh", "parameters": {"freq": "MS"}},
+                identifier="precip_average_ms",
+                module="tests",
+            ),
+        ]
+        # indicators as different types
+        inds_for_avail_vars_from_list = xs.indicators.select_inds_for_avail_vars(
+            ds=ds, indicators=indicators
+        )
+        module = xclim.core.indicator.build_indicator_module(
+            "indicators", {i.base: i for i in indicators}, reload=True
+        )
+        inds_for_avail_vars_from_module = xs.indicators.select_inds_for_avail_vars(
+            ds=ds, indicators=module
+        )
+        inds_for_avail_vars_from_list_of_tuples = (
+            xs.indicators.select_inds_for_avail_vars(
+                ds=ds, indicators=[(n, i) for n, i in module.iter_indicators()]
+            )
+        )
+
+        for inds_for_avail_vars in [
+            inds_for_avail_vars_from_list,
+            inds_for_avail_vars_from_module,
+            inds_for_avail_vars_from_list_of_tuples,
+        ]:
+            assert len(list(inds_for_avail_vars.iter_indicators())) == 1
+            assert [n for n, _ in inds_for_avail_vars.iter_indicators()] == ["tg_min"]
+            assert [i for _, i in inds_for_avail_vars.iter_indicators()] == [
+                indicators[0]
+            ]

xscen/indicators.py

@@ -20,7 +20,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 __all__ = ["compute_indicators", "load_xclim_module"]
 
 
@@ -291,3 +290,55 @@ def func(ds, *, ind, nout):
 
     func.__name__ = ind.identifier
     return partial(func, ind=ind, nout=nout)
+
+
+def select_inds_for_avail_vars(
+    ds: xr.Dataset,
+    indicators: Union[
+        str, Path, Sequence[Indicator], Sequence[tuple[str, Indicator]], ModuleType
+    ],
+) -> Sequence[Indicator]:
+    """Filter the indicators for which the necessary variables are available.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset to use for the indicators.
+    indicators : Union[str, Path, Sequence[Indicator], Sequence[Tuple[str, Indicator]]]
+        Path to a YAML file that instructs on how to calculate missing variables.
+        Can also be only the "stem", if translations and custom indices are implemented.
+        Can be the indicator module directly, or a sequence of indicators or a sequence of
+        tuples (indicator name, indicator) as returned by `iter_indicators()`.
+
+    Returns
+    -------
+    ModuleType – An indicator module.
+
+    See Also
+    --------
+    xclim.indicators, xclim.core.indicator.build_indicator_module_from_yaml
+    """
+    # transform indicator input into a list of tuples (name, indicator)
+    is_list_of_tuples = isinstance(indicators, list) and all(
+        isinstance(i, tuple) for i in indicators
+    )
+    if isinstance(indicators, (str, Path)):
+        logger.debug("Loading indicator module.")
+        indicators = load_xclim_module(indicators, reload=True)
+    if hasattr(indicators, "iter_indicators"):
+        indicators = [(name, ind) for name, ind in indicators.iter_indicators()]
+    elif isinstance(indicators, (list, tuple)) and not is_list_of_tuples:
+        indicators = [(ind.base, ind) for ind in indicators]
+
+    available_vars = {
+        var for var in ds.data_vars if var in xc.core.utils.VARIABLES.keys()
+    }
+    available_inds = [
+        (name, ind)
+        for var in available_vars
+        for name, ind in indicators
+        if var in ind.parameters.keys()
+    ]
+    return xc.core.indicator.build_indicator_module(
+        "inds_for_avail_vars", available_inds
+    )