Feat/boundary dataloader

README.md

@@ -108,7 +108,9 @@ Once `neural-lam` is installed you will be able to train/evaluate models. For th
    interface that provides the data in a data-structure that can be used within
    neural-lam. A datastore is used to create a `pytorch.Dataset`-derived
    class that samples the data in time to create individual samples for
-   training, validation and testing.
+   training, validation and testing. A secondary datastore can be provided
+   for the boundary data. Currently, boundary datastore must be of type `mdp`
+   and only contain forcing features. This can easily be expanded in the future.
 
 2. **The graph structure** is used to define message-passing GNN layers,
    that are trained to emulate fluid flow in the atmosphere over time. The
@@ -121,7 +123,7 @@ different aspects about the training and evaluation of the model.
 
 The path you provide to the neural-lam config (`config.yaml`) also sets the
 root directory relative to which all other paths are resolved, as in the parent
-directory of the config becomes the root directory. Both the datastore and
+directory of the config becomes the root directory. Both the datastores and
 graphs you generate are then stored in subdirectories of this root directory.
 Exactly how and where a specific datastore expects its source data to be stored
 and where it stores its derived data is up to the implementation of the
@@ -134,6 +136,7 @@ assume you placed `config.yaml` in a folder called `data`):
 data/
 ├── config.yaml           - Configuration file for neural-lam
 ├── danra.datastore.yaml  - Configuration file for the datastore, referred to from config.yaml
+├── era5.datastore.zarr/  - Optional configuration file for the  boundary datastore, referred to from config.yaml
 └── graphs/               - Directory containing graphs for training
 ```
 
@@ -142,18 +145,20 @@ And the content of `config.yaml` could in this case look like:
 datastore:
   kind: mdp
   config_path: danra.datastore.yaml
+datastore_boundary:
+  kind: mdp
+  config_path: era5.datastore.yaml
 training:
   state_feature_weighting:
     __config_class__: ManualStateFeatureWeighting
-    values:
+    weights:
       u100m: 1.0
       v100m: 1.0
 ```
 
-For now the neural-lam config only defines two things: 1) the kind of data
-store and the path to its config, and 2) the weighting of different features in
-the loss function. If you don't define the state feature weighting it will default
-to weighting all features equally.
+For now the neural-lam config only defines two things:
+1) the kind of datastores and the path to their config
+2) the weighting of different features in the loss function. If you don't define the state feature weighting it will default to weighting all features equally.
 
 (This example is taken from the `tests/datastore_examples/mdp` directory.)
 
@@ -525,5 +530,4 @@ Furthermore, all tests in the ```tests``` directory will be run upon pushing cha
 
 # Contact
 If you are interested in machine learning models for LAM, have questions about the implementation or ideas for extending it, feel free to get in touch.
-There is an open [mllam slack channel](https://join.slack.com/t/ml-lam/shared_invite/zt-2t112zvm8-Vt6aBvhX7nYa6Kbj_LkCBQ) that anyone can join (after following the link you have to request to join, this is to avoid spam bots).
-You can also open a github issue on this page, or (if more suitable) send an email to [[email protected]](mailto:[email protected]).
+There is an open [mllam slack channel](https://join.slack.com/t/ml-lam/shared_invite/zt-2t112zvm8-Vt6aBvhX7nYa6Kbj_LkCBQ) that anyone can join. You can also open a github issue on this page, or (if more suitable) send an email to [[email protected]](mailto:[email protected]).

neural_lam/config.py

@@ -97,11 +97,15 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
     ----------
     datastore : DatastoreSelection
         The configuration for the datastore to use.
+    datastore_boundary : Union[DatastoreSelection, None]
+        The configuration for the boundary datastore to use, if any. If None,
+        no boundary datastore is used.
     training : TrainingConfig
         The configuration for training the model.
     """
 
     datastore: DatastoreSelection
+    datastore_boundary: Union[DatastoreSelection, None] = None
     training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)
 
     class _(dataclass_wizard.JSONWizard.Meta):
@@ -136,11 +140,11 @@ class InvalidConfigError(Exception):
     pass
 
 
-def load_config_and_datastore(
+def load_config_and_datastores(
     config_path: str,
 ) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]:
     """
-    Load the neural-lam configuration and the datastore specified in the
+    Load the neural-lam configuration and the datastores specified in the
     configuration.
 
     Parameters
@@ -151,7 +155,7 @@ def load_config_and_datastore(
     Returns
     -------
     tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]
-        The Neural-LAM configuration and the loaded datastore.
+        The Neural-LAM configuration and the loaded datastores.
     """
     try:
         config = NeuralLAMConfig.from_yaml_file(config_path)
@@ -168,4 +172,15 @@ def load_config_and_datastore(
         datastore_kind=config.datastore.kind, config_path=datastore_config_path
     )
 
-    return config, datastore
+    if config.datastore_boundary is not None:
+        datastore_boundary_config_path = (
+            Path(config_path).parent / config.datastore_boundary.config_path
+        )
+        datastore_boundary = init_datastore(
+            datastore_kind=config.datastore_boundary.kind,
+            config_path=datastore_boundary_config_path,
+        )
+    else:
+        datastore_boundary = None
+
+    return config, datastore, datastore_boundary

neural_lam/create_graph.py

@@ -13,7 +13,7 @@
 from torch_geometric.utils.convert import from_networkx
 
 # Local
-from .config import load_config_and_datastore
+from .config import load_config_and_datastores
 from .datastore.base import BaseRegularGridDatastore
 
 
@@ -595,7 +595,7 @@ def cli(input_args=None):
     ), "Specify your config with --config_path"
 
     # Load neural-lam configuration and datastore to use
-    _, datastore = load_config_and_datastore(config_path=args.config_path)
+    _, datastore = load_config_and_datastores(config_path=args.config_path)
 
     create_graph_from_datastore(
         datastore=datastore,

neural_lam/datastore/base.py

@@ -47,7 +47,6 @@ class BaseDatastore(abc.ABC):
     each of the `x` and `y` coordinates.
     """
 
-    is_ensemble: bool = False
     is_forecast: bool = False
 
     @property
@@ -193,11 +192,9 @@ def get_dataarray(
         """
         Return the processed data (as a single `xr.DataArray`) for the given
         category of data and test/train/val-split that covers all the data (in
-        space and time) of a given category (state/forcing/static). A
-        datastore must be able to return for the "state" category, but
-        "forcing" and "static" are optional (in which case the method should
-        return `None`). For the "static" category the `split` is allowed to be
-        `None` because the static data is the same for all splits.
+        space and time) of a given category (state/forcing/static). For the
+        "static" category the `split` is allowed to be `None` because the static
+        data is the same for all splits.
 
         The returned dataarray is expected to at minimum have dimensions of
         `(grid_index, {category}_feature)` so that any spatial dimensions have
@@ -228,23 +225,6 @@ def get_dataarray(
         """
         pass
 
-    @cached_property
-    @abc.abstractmethod
-    def boundary_mask(self) -> xr.DataArray:
-        """
-        Return the boundary mask for the dataset, with spatial dimensions
-        stacked. Where the value is 1, the grid point is a boundary point, and
-        where the value is 0, the grid point is not a boundary point.
-
-        Returns
-        -------
-        xr.DataArray
-            The boundary mask for the dataset, with dimensions
-            `('grid_index',)`.
-
-        """
-        pass
-
     @abc.abstractmethod
     def get_xy(self, category: str) -> np.ndarray:
         """
@@ -317,6 +297,31 @@ def num_grid_points(self) -> int:
         """
         pass
 
+    @property
+    def num_ensemble_members(self) -> int:
+        """Return the number of ensemble members in the dataset.
+
+        Returns
+        -------
+        int
+            The number of ensemble members in the dataset (default is 1 -
+            not an ensemble).
+
+        """
+        return 1
+
+    @property
+    def is_ensemble(self) -> bool:
+        """Return whether the dataset represents ensemble data.
+
+        Returns
+        -------
+        bool
+            True if the dataset represents ensemble data, False otherwise.
+
+        """
+        return self.num_ensemble_members > 1
+
     @cached_property
     @abc.abstractmethod
     def state_feature_weights_values(self) -> List[float]: