Skip to content

Commit

Permalink
Select the set of indicators for available variables (#291)
Browse files Browse the repository at this point in the history
### Pull Request Checklist:
- [ ] This PR addresses an already opened issue (for bug fixes /
features)
    - This PR fixes #xyz
- [ ] (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] CHANGES.rst has been updated (with summary of main changes).
  - [x] Link to pull request (:pull:`number`) has been added.

### What kind of change does this PR introduce?

Adds the function `xscen.indicators.select_inds_for_avail_vars` to
filter the indicators that can be calculated with the variables
available in a `xarray.Dataset`.

### Does this PR introduce a breaking change?

No.

### Other information:
  • Loading branch information
vindelico authored Dec 5, 2023
2 parents 1ae4981 + 256a333 commit 53f08d5
Show file tree
Hide file tree
Showing 3 changed files with 115 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`).
Expand Down
58 changes: 58 additions & 0 deletions tests/test_indicators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()] == []
57 changes: 56 additions & 1 deletion xscen/indicators.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@

logger = logging.getLogger(__name__)


__all__ = ["compute_indicators", "load_xclim_module"]


Expand Down Expand Up @@ -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
)

0 comments on commit 53f08d5

Please sign in to comment.