feat: support area masking of boundary forcing

CHANGELOG.md

@@ -19,6 +19,7 @@ Keep it human-readable, your future self will thank you!
 
 ### Changed
 - ci: small fixes and updates pre-commit, downsteam-ci (#49)
+- feat: New argument 'forcing_area' in the CutOutZarDatasetNodes class. (#52)
 
 ## [0.3.0 Anemoi-graphs, minor release](https://github.com/ecmwf/anemoi-graphs/compare/0.2.1...0.3.0) - 2024-09-03
 

src/anemoi/graphs/nodes/builders/from_file.py

@@ -52,15 +52,52 @@ def get_coordinates(self) -> torch.Tensor:
 
 
 class CutOutZarrDatasetNodes(ZarrDatasetNodes):
-    """Nodes from Zarr dataset."""
+    """Nodes from Zarr dataset.
+
+    Attributes
+    ----------
+    lam_dataset : str
+        The limited area dataset.
+    forcing_dataset : str
+        The forcing dataset.
+    thinning : int, optional
+        The thinning factor. Defaults to 1, which means no thinning.
+    forcing_area : list[float], optional
+        The area of the forcing dataset. Specify the longitude and
+        latitudes boundaries as (north, west, south, east). Defaults
+        to None, which means the forcing dataset is not cropped.
+
+    Methods
+    -------
+    get_coordinates()
+        Get the lat-lon coordinates of the nodes.
+    register_nodes(graph, name)
+        Register the nodes in the graph.
+    register_attributes(graph, name, config)
+        Register the attributes in the nodes of the graph specified.
+    update_graph(graph, name, attr_config)
+        Update the graph with new nodes and attributes.
+    """
 
     def __init__(
-        self, name: str, lam_dataset: str, forcing_dataset: str, thinning: int = 1, adjust: str = "all"
+        self,
+        name: str,
+        lam_dataset: str,
+        forcing_dataset: str,
+        thinning: int = 1,
+        forcing_area: list[float] | None = None,
+        adjust: str = "all",
     ) -> None:
-        dataset_config = {
-            "cutout": [{"dataset": lam_dataset, "thinning": thinning}, {"dataset": forcing_dataset}],
-            "adjust": adjust,
-        }
+        lam_config = {"dataset": lam_dataset, "thinning": thinning}
+        forcing_config = {"dataset": forcing_dataset}
+
+        # Add area argument to crop the boundary forcing
+        if forcing_area is not None:
+            forcing_config["area"] = tuple(forcing_area)
+            assert len(forcing_area) == 4, "The forcing area must be a list of 4 elements (north, west, south, east)."
+            LOGGER.info("Forcing dataset is cropped to area: %s", forcing_area)
+
+        dataset_config = {"cutout": [lam_config, forcing_config], "adjust": adjust}
         super().__init__(dataset_config, name)
         self.n_cutout, self.n_other = self.dataset.grids