Select the set of indicators for available variables

CHANGES.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:`291`).
 * Replaced aggregation function ``climatological_mean()`` with ``climatological_op()`` offering more types of operations to aggregate over climatological 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,61 @@ 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"
+
+    @pytest.mark.parametrize("indicator_iter", ["list", "tuples", "module"])
+    def test_select_inds_for_avail_vars(self, indicator_iter):
+        # 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
+        module = xclim.core.indicator.build_indicator_module(
+            "indicators", {i.base: i for i in indicators}, reload=True
+        )
+        if indicator_iter == "list":
+            inds_for_avail_vars = xs.indicators.select_inds_for_avail_vars(
+                ds=ds, indicators=indicators
+            )
+        elif indicator_iter == "tuples":
+            inds_for_avail_vars = xs.indicators.select_inds_for_avail_vars(
+                ds=ds, indicators=[(n, i) for n, i in module.iter_indicators()]
+            )
+        elif indicator_iter == "module":
+            inds_for_avail_vars = xs.indicators.select_inds_for_avail_vars(
+                ds=ds, indicators=module
+            )
+
+        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]]
+        # no indicators found
+        inds_for_avail_vars = xs.indicators.select_inds_for_avail_vars(
+            ds=ds, indicators=indicators[1:]
+        )
+        assert len(list(inds_for_avail_vars.iter_indicators())) == 0
+        assert [(n, i) for n, i in inds_for_avail_vars.iter_indicators()] == []

xscen/indicators.py

@@ -20,7 +20,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 __all__ = ["compute_indicators", "load_xclim_module"]
 
 
@@ -293,3 +292,59 @@ 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,
+        os.PathLike,
+        Sequence[Indicator],
+        Sequence[tuple[str, Indicator]],
+        ModuleType,
+    ],
+) -> ModuleType:
+    """Filter the indicators for which the necessary variables are available.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset to use for the indicators.
+    indicators : Union[str, os.PathLike, Sequence[Indicator], Sequence[Tuple[str, Indicator]]]
+        Path to a YAML file that instructs on how to calculate indicators.
+        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 of 'length' ∈ [0, n].
+
+    See Also
+    --------
+    xclim.indicators, xclim.core.indicator.build_indicator_module_from_yaml
+    """
+    # Transform the 'indicators' 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, os.PathLike)):
+        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, reload=True
+    )