feat: add optional boundary datastore support

AGENTS.md

@@ -9,15 +9,15 @@ Mandatory rules for AI coding agents. Violations will result in rejected PRs.
 Neural-LAM: graph-based neural weather prediction for Limited Area Modeling. Models: `GraphLAM`,
 `HiLAM`, `HiLAMParallel`.
 
-**Data flow:** Raw zarr/numpy → `Datastore` → `WeatherDataset` → `WeatherDataModule` → Model →
-Predictions
+**Data flow:** Raw zarr/numpy → `Datastore` (+ optional boundary `Datastore`) → `WeatherDataset` →
+`WeatherDataModule` → Model → Predictions
 
 **Key modules:**
 - `datastore/` — `BaseDatastore` (abstract), `MDPDatastore` (zarr via mllam-data-prep)
-- `models/` — `ARModel` (autoregressive base, Lightning) → `BaseGraphModel` (encode-process-decode)
-  → `GraphLAM` / `HiLAM` / `HiLAMParallel`
-- `weather_dataset.py` — `WeatherDataset` + `WeatherDataModule`
-- `config.py` — YAML config via dataclass-wizard
+- `models/` — `ForecasterModule` (Lightning) → `ARForecaster` (Forecaster) →
+  `GraphLAM` / `HiLAM` / `HiLAMParallel` (StepPredictor)
+- `weather_dataset.py` — `WeatherDataset` + `WeatherDataModule` (supports optional boundary datastore)
+- `config.py` — YAML config via dataclass-wizard (`NeuralLAMConfig` with optional `datastore_boundary`)
 - `create_graph.py` — builds mesh graphs (must run before training)
 - `interaction_net.py` — `InteractionNet` GNN layer (PyG `MessagePassing`)
 - `utils.py` — `make_mlp`, normalization helpers

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add optional boundary datastore support: `NeuralLAMConfig` now takes a named `datastores` dict. `WeatherDataset` loads boundary forcing from such a datastore and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,
@@ -25,6 +27,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:` [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov
+
 - Move data normalization from CPU (`WeatherDataset`) to GPU
   (`ForecasterModule.on_after_batch_transfer`) for improved performance and
   multi-GPU compatibility. `WeatherDataset` / `WeatherDataModule` no longer

README.md

@@ -154,9 +154,13 @@ data/
 
 And the content of `config.yaml` could in this case look like:
 ```yaml
-datastore:
-  kind: mdp
-  config_path: danra.datastore.yaml
+datastores:
+  danra:
+    kind: mdp
+    config_path: danra.datastore.yaml
+  era5_boundary:  # optional; no `state` data, so used for input (boundary) only
+    kind: mdp
+    config_path: era5_boundary.datastore.yaml
 training:
   state_feature_weighting:
     __config_class__: ManualStateFeatureWeighting
@@ -175,7 +179,16 @@ training:
 
 For now the neural-lam config only defines few things:
 
-1. The kind of datastore and the path to its config
+1. A named mapping of datastores (`datastores`), each giving the kind of
+   datastore and the path to its config. The role of each datastore is implied
+   by the categories of data it provides: a datastore that contains `state`
+   data is used for both model input and output (the interior domain), while a
+   datastore without `state` data is used for input only (e.g. boundary forcing
+   from a separate domain such as ERA5 for a LAM domain). Exactly one datastore
+   must provide `state` data, and at most one datastore may omit it, so there is
+   a single interior and a single (optional) boundary datastore. When a boundary
+   datastore is present its forcing is windowed and included as an additional
+   tensor in each training sample.
 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.
@@ -215,6 +228,20 @@ the input-data representation is split into two parts:
    `WeatherDataset` class is also responsible for normalising the values and
    returning `torch.Tensor`-objects.
 
+Each variable in a datastore is assigned to one of three data *categories*,
+which fix whether it is fed to the model as input, predicted as output, or both:
+
+| Category  | Model input | Model output | Description |
+|-----------|:-----------:|:------------:|-------------|
+| `state`   | ✓           | ✓            | Prognostic variables the model both reads and predicts (autoregressed forward in time). |
+| `forcing` | ✓           |              | Time-varying inputs known in advance (e.g. solar radiation, boundary forcing). |
+| `static`  | ✓           |              | Time-invariant inputs (e.g. orography, land-sea mask). |
+
+These categories are
+also what determine a datastore's role: a datastore that provides `state` data
+is the interior domain (model input and output), while one without `state` data
+is used for input only, e.g. boundary forcing from a separate domain.
+
 There are currently two different datastores implemented in the codebase:
 
 1. `neural_lam.datastore.MDPDatastore` which represents loading of
@@ -380,9 +407,10 @@ Which you can then use in a neural-lam configuration file like this:
 
 ```yaml
 # config.yaml
-datastore:
-  kind: npyfilesmeps
-  config_path: meps.datastore.yaml
+datastores:
+  meps:
+    kind: npyfilesmeps
+    config_path: meps.datastore.yaml
 training:
   state_feature_weighting:
     __config_class__: ManualStateFeatureWeighting
@@ -584,13 +612,18 @@ Canonical dimension names used in tensor shape annotations throughout the codeba
 
 - `B` - batch size
 - `pred_steps` - number of autoregressive prediction steps
+- `num_times` - number of time steps along the time axis of a raw or batched timeseries (a trailing `'`, e.g. `num_times'`, denotes a pre-subsampling / pre-differencing variant)
 - `num_grid_nodes` - number of nodes in the flattened spatial grid
+- `num_boundary_grid_nodes` - number of nodes in the flattened boundary spatial grid
 - `num_mesh_nodes` - number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l`
 - `num_state_vars` - number of atmospheric state variables
 - `num_forcing_vars` - number of forcing input variables
+- `num_windowed_forcing_vars` - forcing variables stacked over the past/future forcing window
+- `num_windowed_boundary_vars` - boundary forcing variables stacked over the past/future boundary window
 - `num_variables` - generic variable dimension used in metric functions
 - `hidden_dim` - internal hidden representation size in GNN layers and MLPs
 - `input_dim` - input feature dimensionality to a layer before transformation
+- `d_mesh_static` - number of static features per mesh node
 - `num_edges` - number of edges in a graph (g2m, m2g, same-level, up, down)
 - `num_send` - number of sender nodes in a message-passing step
 - `num_rec` - number of receiver nodes in a message-passing step

neural_lam/config.py

@@ -7,6 +7,7 @@
 
 # Third-party
 import dataclass_wizard
+import yaml
 
 # Local
 from .datastore import (
@@ -127,15 +128,22 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
 
     Attributes
     ----------
-    datastore : DatastoreSelection
-        Configuration specifying which datastore backend to use and its
-        associated settings.
+    datastores : Dict[str, DatastoreSelection]
+        Mapping from a user-chosen datastore name to its selection config. The
+        role of each datastore is implied by the categories of data it
+        provides rather than by a dedicated config key: a datastore that
+        contains `state` data is used for both model input and output (the
+        interior domain), while a datastore without `state` data is used for
+        input only (e.g. boundary forcing from a separate domain). Exactly
+        one datastore must provide `state` data, and at most one datastore
+        may omit it (the boundary); both constraints are enforced in
+        :func:`load_config_and_datastore`.
     training : TrainingConfig
         Configuration for training the model, including loss function and
         feature-weighting strategy. Defaults to ``TrainingConfig()``.
     """
 
-    datastore: DatastoreSelection
+    datastores: Dict[str, DatastoreSelection]
     training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)
 
     class _(dataclass_wizard.JSONWizard.Meta):
@@ -174,9 +182,13 @@ class InvalidConfigError(Exception):
 
 def load_config_and_datastore(
     config_path: str,
-) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]:
+) -> tuple[
+    NeuralLAMConfig,
+    Union[MDPDatastore, NpyFilesDatastoreMEPS],
+    Union[MDPDatastore, NpyFilesDatastoreMEPS, None],
+]:
     """
-    Load the neural-lam configuration and the datastore specified in the
+    Load the neural-lam configuration and the datastores specified in the
     configuration.
 
     Parameters
@@ -186,22 +198,74 @@ def load_config_and_datastore(
 
     Returns
     -------
-    tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]
-        The Neural-LAM configuration and the loaded datastore.
+    tuple[NeuralLAMConfig, datastore, datastore_boundary]
+        The Neural-LAM configuration, the loaded interior datastore (the one
+        providing `state` data), and the boundary datastore (the one without
+        `state` data, or None if no such datastore is configured).
+
+    Raises
+    ------
+    InvalidConfigError
+        If not exactly one datastore provides `state` data, or if more than
+        one datastore omits it (only a single boundary datastore is currently
+        supported).
     """
+    with open(config_path, encoding="utf-8") as f:
+        raw_config = yaml.safe_load(f)
+    if isinstance(raw_config, dict) and (
+        "datastore" in raw_config and "datastores" not in raw_config
+    ):
+        raise InvalidConfigError(
+            "The `datastore:` config key has been replaced by a named "
+            "`datastores:` mapping (the role of each datastore is now implied "
+            "by the categories it provides). Move your datastore under a "
+            "named entry, e.g.:\n"
+            "datastores:\n"
+            "  <name>:\n"
+            "    kind: ...\n"
+            "    config_path: ...\n"
+            "See the README and CHANGELOG for details."
+        )
+
     try:
         config = NeuralLAMConfig.from_yaml_file(config_path)
     except dataclass_wizard.errors.UnknownJSONKey as ex:
         raise InvalidConfigError(
             "There was an error loading the configuration file at "
             f"{config_path}. "
         ) from ex
-    # datastore config is assumed to be relative to the config file
-    datastore_config_path = (
-        Path(config_path).parent / config.datastore.config_path
-    )
-    datastore = init_datastore(
-        datastore_kind=config.datastore.kind, config_path=datastore_config_path
-    )
 
-    return config, datastore
+    # datastore configs are assumed to be relative to the config file. The
+    # role of each datastore is implied by the categories of data it provides:
+    # a datastore with `state` data is the interior (input and output), one
+    # without `state` data is used for input only (e.g. boundary forcing).
+    config_dir = Path(config_path).parent
+    interior_datastores = {}
+    boundary_datastores = {}
+    for name, selection in config.datastores.items():
+        datastore = init_datastore(
+            datastore_kind=selection.kind,
+            config_path=config_dir / selection.config_path,
+        )
+        if datastore.get_num_data_vars(category="state") > 0:
+            interior_datastores[name] = datastore
+        else:
+            boundary_datastores[name] = datastore
+
+    if len(interior_datastores) != 1:
+        raise InvalidConfigError(
+            "Exactly one datastore must provide `state` data (the interior "
+            f"domain), but {len(interior_datastores)} were found in "
+            f"{config_path}: {sorted(interior_datastores)}."
+        )
+    if len(boundary_datastores) > 1:
+        raise InvalidConfigError(
+            "At most one boundary datastore (a datastore without `state` "
+            f"data) is currently supported, but {len(boundary_datastores)} "
+            f"were found in {config_path}: {sorted(boundary_datastores)}."
+        )
+
+    (datastore,) = interior_datastores.values()
+    datastore_boundary = next(iter(boundary_datastores.values()), None)
+
+    return config, datastore, datastore_boundary

neural_lam/create_graph.py

@@ -730,7 +730,7 @@ def cli(input_args: Optional[list[str]] = None) -> None:
         raise ValueError("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_datastore(config_path=args.config_path)
 
     create_graph_from_datastore(
         datastore=datastore,