feat(apply mask): apply a boolean mask when opening a dataset (#496)

## Description
This PR introduces `mask` as an argument to `open_dataset`. This
argument allows to apply an arbitrary spatial boolean mask from a .npy
file to the dataset.

## What problem does this change solve?
It is a new feature for `open_dataset` that allows to select arbitrary
spatial points of a dataset, it could be used e.g. to thin data on an
irregular grid or select a region of interest.
It was tested to reduce the density of grid points of the triangular
grid of the ICON model. The plots show a zoom over Switzerland with the
dense grid at approx. 1km resolution and the grid after applying a mask
selecting only every second triangle.

Original grid:
<img width="500" alt="image"
src="https://github.com/user-attachments/assets/9a8e1036-c77a-4f35-ab1f-1ef67014e363"
/>

Half the points masked using external mask:
<img width="500" alt="image"
src="https://github.com/user-attachments/assets/d3cfcb36-3f5f-43a6-aee0-c29b1d9f66d7"
/>


## What issue or task does this change relate to?

##  Additional notes ##

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)


<!-- readthedocs-preview anemoi-datasets start -->
----
📚 Documentation preview 📚:
https://anemoi-datasets--496.org.readthedocs.build/en/496/

<!-- readthedocs-preview anemoi-datasets end -->

---------

Co-authored-by: Francesco Zanetta <62377868+frazane@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Gert Mertes <13658335+gmertes@users.noreply.github.com>

Author: 4 people

docs/using/grids.rst

@@ -37,6 +37,25 @@ dataset:
    :width: 75%
    :align: center
 
+*********
+ masking
+*********
+
+You can apply an arbitrary spatial mask to a dataset by specifying the
+``mask`` parameter in the ``open_dataset`` function. The mask must be a
+NumPy .npy file containing a boolean array, where True indicates points
+to be kept and False indicates points to be removed.
+
+.. code:: python
+
+   ds = open_dataset(dataset, mask="path/to/mask.npy")
+
+The mask array must have the same total number of grid points and
+dimension as the dataset. Please note that this masking will not be
+automatically applied to the input data provided during inference.
+Users must ensure they are applying the mask explicitly, for instance
+via the use of pre-processors.
+
 ******
  area
 ******

src/anemoi/datasets/data/dataset.py

@@ -245,6 +245,12 @@ def __subset(self, **kwargs: Any) -> "Dataset":
 
             return Statistics(self, open_dataset(statistics))._subset(**kwargs).mutate()
 
+        if "mask" in kwargs:
+            from .masked import Masking
+
+            mask_file = kwargs.pop("mask")
+            return Masking(self, mask_file)._subset(**kwargs).mutate()
+
         # Note: trim_edge should go before thinning
         if "trim_edge" in kwargs:
             from .masked import TrimEdge

src/anemoi/datasets/data/masked.py

@@ -10,6 +10,7 @@
 
 import logging
 from functools import cached_property
+from pathlib import Path
 from typing import Any
 
 import numpy as np
@@ -207,6 +208,60 @@ def field_shape(self) -> Shape:
         return x, y
 
 
+class Masking(Masked):
+    """A class that applies a precomputed boolean mask from a .npy file."""
+
+    def __init__(self, forward: Dataset, mask_file: str) -> None:
+        """Initialize the Masking class.
+
+        Parameters
+        ----------
+        forward : Dataset
+            The dataset to be masked.
+        mask_file : str
+            Path to a .npy file containing a boolean mask of same shape as fields.
+        """
+        self.mask_file = mask_file
+
+        # Check path
+        if not Path(self.mask_file).exists():
+            raise FileNotFoundError(f"Mask file not found: {self.mask_file}")
+        # Load mask
+        try:
+            mask = np.load(self.mask_file)
+        except Exception as e:
+            raise ValueError(f"Could not load data from {mask_file}: {e}")
+
+        if mask.dtype != bool:
+            raise ValueError(f"Mask file {mask_file} does not contain boolean values.")
+        if mask.shape != forward.field_shape:
+            raise ValueError(f"Mask length {mask.shape} does not match field size {forward.field_shape}.")
+        if sum(mask) == 0:
+            LOG.warning(f"Mask in {mask_file} eliminates all points in field.")
+
+        super().__init__(forward, mask)
+
+    def tree(self) -> Node:
+        """Get the tree representation of the dataset.
+
+        Returns
+        -------
+        Node
+            The tree representation of the dataset.
+        """
+        return Node(self, [self.forward.tree()], mask_file=self.mask_file)
+
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        """Get the metadata specific to the Masking subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the Masking subclass.
+        """
+        return dict(mask_file=self.mask_file)
+
+
 class Cropping(Masked):
     """A class to represent a cropped dataset."""
 

tests/test_data.py

@@ -30,6 +30,7 @@
 from anemoi.datasets.data.ensemble import Ensemble
 from anemoi.datasets.data.grids import GridsBase
 from anemoi.datasets.data.join import Join
+from anemoi.datasets.data.masked import Masking
 from anemoi.datasets.data.misc import as_first_date
 from anemoi.datasets.data.misc import as_last_date
 from anemoi.datasets.data.padded import Padded
@@ -1400,6 +1401,67 @@ def test_invalid_trim_edge() -> None:
         )
 
 
+@mockup_open_zarr
+def test_masking() -> None:
+    """Test masking a dataset."""
+    test_mask = np.array([True, False, True, True, True, True, False, False, True, False])
+    with (
+        patch("anemoi.datasets.data.masked.np.load", return_value=test_mask),
+        patch("anemoi.datasets.data.masked.Path.exists", return_value=True),
+    ):
+
+        test = DatasetTester("test-2021-2022-6h-o96-abcd", mask="./test_mask.npy")
+        test.run(
+            expected_class=Masking,
+            expected_length=365 * 2 * 4,
+            date_to_row=lambda date: simple_row(date, "abcd")[..., test_mask],
+            start_date=datetime.datetime(2021, 1, 1),
+            time_increment=datetime.timedelta(hours=6),
+            expected_shape=(365 * 2 * 4, 4, 1, sum(test_mask)),
+            expected_variables="abcd",
+            expected_name_to_index="abcd",
+            statistics_reference_dataset="test-2021-2022-6h-o96-abcd",
+            statistics_reference_variables="abcd",
+        )
+    return
+
+
+@mockup_open_zarr
+def test_masking_wrong_mask_dims() -> None:
+    """Test masking a dataset (wrong dims in mask)."""
+    test_mask = np.array([True, False, True, True, True, True, False, False, True])
+    with (
+        patch("anemoi.datasets.data.masked.np.load", return_value=test_mask),
+        patch("anemoi.datasets.data.masked.Path.exists", return_value=True),
+    ):
+        with pytest.raises(ValueError):
+            _ = DatasetTester("test-2021-2022-6h-o96-abcd", mask="./test_mask.npy")
+    return
+
+
+@mockup_open_zarr
+def test_masking_mask_file_not_found() -> None:
+    """Test masking a dataset (mask file not found)."""
+    test_mask = np.array([True, False, True, True, True, True, False, False, True, False])
+    with patch("anemoi.datasets.data.masked.np.load", return_value=test_mask):
+        with pytest.raises(FileNotFoundError):
+            _ = DatasetTester("test-2021-2022-6h-o96-abcd", mask="./test_mask.npy")
+    return
+
+
+@mockup_open_zarr
+def test_masking_wrong_dtype() -> None:
+    """Test masking a dataset (mask file not found)."""
+    test_mask = np.array([1, 0, 1, 1, 1, 1, 0, 0, 1, 0])
+    with (
+        patch("anemoi.datasets.data.masked.np.load", return_value=test_mask),
+        patch("anemoi.datasets.data.masked.Path.exists", return_value=True),
+    ):
+        with pytest.raises(ValueError):
+            _ = DatasetTester("test-2021-2022-6h-o96-abcd", mask="./test_mask.npy")
+    return
+
+
 def test_save_dataset() -> None:
     """Test save datasets."""