feat: adopt #652 multi-datastore schema + ForecastBatch return shape

Pre-emptively land the public API shape proposed in #652
on top of the boundary-datastore work in #635, so the model-side
adapter (Joel's follow-up) doesn't have to break the schema again
later.

Config schema (neural_lam/config.py):
- Replace `datastore` + `datastore_boundary` top-level keys with a
  single `datastores: Dict[str, DatastoreSelection]` mapping. The dict
  key becomes the canonical source name used throughout the pipeline
  (ForecastBatch field keys, weight / clamping disambiguation, etc).
- DatastoreSelection grows optional `inputs:` and `outputs:` per-
  category variable include-lists, with `None` per-category values
  meaning "all variables in that category". `outputs:` declares which
  datastore is the interior (prognostic source); omitted means input-
  only.
- Validate at config load: raise InvalidConfigError when two
  datastores declare the same variable as an output, pointing the
  user at mdp's `dim_mapping.name_format` (new builds) or
  `xr.Dataset.assign_coords` on the existing zarr's small
  `{category}_feature` coord (milliseconds regardless of zarr size).
- `load_config_and_datastore` returns `(config, Dict[str, BaseDatastore])`
  - the full multi-source mapping - rather than the legacy
  `(config, interior, boundary)` triple. No transitional adapter.

Dataset return shape (neural_lam/weather_dataset.py):
- New `ForecastBatch` NamedTuple with per-source dict fields
  (`init_states`, `target_states`, `forcing`) plus a global
  `target_times` tensor. PyTorch's default_collate recurses through
  the NamedTuple and the dicts, stacking per-source tensors along a
  new batch axis - no custom collate_fn required.
- `WeatherDataset.__init__` takes `(datastores, selections, ...)`
  dicts directly. Internally it still resolves the single interior +
  optional boundary pair for its slicing/windowing logic - that
  internal multi-source rewrite is part of #652's model-side
  follow-up. The PUBLIC return type is already shaped for the
  multi-source case so the future change is additive on the producer
  side.
- `WeatherDataModule` signature updated to match.
- `create_dataarray_from_tensor` refactored to expose a
  `build_dataarray_from_tensor` staticmethod for callers (e.g.
  `ForecasterModule._create_dataarray_from_tensor`) that have a
  datastore but not a full dataset.

Production call site (neural_lam/train_model.py):
- Unpacks the new `(config, datastores)` return shape.
- Uses `_resolve_datastore_roles` to pick out the interior +
  boundary for the legacy ForecasterModule constructor (which keeps
  its single-datastore + boundary shape pending Joel's adapter).
- WeatherDataModule receives the full multi-source dicts.

Example YAMLs (tests/datastore_examples/):
- Single-source danra: top-level `datastores: {danra: ...}` wrapping.
- danra + era5 boundary: `datastores: {interior: ..., boundary: ...}`
  with explicit `outputs: {state: }` on interior so the resolver
  knows which one is the prognostic source.

Intentionally NOT in this PR (per #652 follow-up scope):
- Model-side adapter. `ForecasterModule.training_step` and friends
  still unpack the legacy 5-tuple `(init_states, target_states,
  forcing, boundary, target_times)`, which will fail at runtime when
  Lightning hands them a `ForecastBatch`. Joel's follow-up replaces
  those positional unpacks with `batch.init_states["interior"]`-style
  per-source dict access.
- Test updates. Every test that builds a `WeatherDataset` or
  `NeuralLAMConfig` with the old keyword signature will fail. They
  need mechanical updates to the new dict shape (and the model-
  exercising ones should stay skipped until the model adapter
  lands).
- Variable include-lists honoured at runtime. The schema parses
  `inputs:` / `outputs:` but `WeatherDataset` still fetches all
  variables per category from each datastore. Filtering by the
  declared subsets is a small follow-up that depends on whether
  Joel's adapter wants to do it at dataset level or model level.
- Diagnostic outputs. The schema accepts `outputs.diagnostic: [...]`
  but nothing concatenates them into the target tensor yet; same
  story (small follow-up after the model adapter decides on the
  prognostic/diagnostic split semantic).

Refs #635, refs #652.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: sadamov

neural_lam/config.py

@@ -1,7 +1,7 @@
 # Standard library
 import dataclasses
 from pathlib import Path
-from typing import Dict, Union
+from typing import Dict, List, Optional, Union
 
 # Third-party
 import dataclass_wizard
@@ -18,26 +18,47 @@
 @dataclasses.dataclass
 class DatastoreSelection:
     """
-    Configuration for selecting a datastore to use with neural-lam.
+    Configuration for selecting a datastore and declaring how its variables
+    are consumed by the model.
 
     Attributes
     ----------
     kind : str
         The kind of datastore to use, currently `mdp` or `npyfilesmeps` are
         implemented.
     config_path : str
-        The path to the configuration file for the selected datastore, this is
-        assumed to be relative to the configuration file for neural-lam.
+        The path to the configuration file for the selected datastore, this
+        is assumed to be relative to the configuration file for neural-lam.
+    inputs : Dict[str, List[str] or None] or None, optional
+        Per-category lists of variable names this datastore contributes as
+        model inputs. Categories are typically ``state``, ``forcing``,
+        ``static``. If the whole field is ``None`` (the default), every
+        variable in every category that the datastore exposes is treated
+        as an input. If a category key is present with a ``null`` / ``None``
+        value, all variables in that category are used. An explicit empty
+        list excludes the category.
+    outputs : Dict[str, List[str] or None] or None, optional
+        Per-category lists of variable names this datastore contributes as
+        model outputs (prediction targets). Categories may include ``state``
+        for prognostic outputs (those that are fed back as input in the
+        next autoregressive step) and ``diagnostic`` for predict-only
+        outputs. Prognostic outputs are the intersection of
+        ``inputs["state"]`` and ``outputs["state"]``; everything in
+        ``outputs`` that is not also in ``inputs["state"]`` is diagnostic.
+        If ``None`` (the default), this datastore is treated as input-only
+        (no contribution to predictions). ``null`` per-category values
+        follow the same "all available" convention as ``inputs``.
     """
 
     kind: str
+    config_path: str
+    inputs: Optional[Dict[str, Optional[List[str]]]] = None
+    outputs: Optional[Dict[str, Optional[List[str]]]] = None
 
     def __post_init__(self):
         if self.kind not in DATASTORES:
             raise ValueError(f"Datastore kind {self.kind} is not implemented")
 
-    config_path: str
-
 
 @dataclasses.dataclass
 class ManualStateFeatureWeighting:
@@ -89,10 +110,10 @@ class TrainingConfig:
     Attributes
     ----------
     state_feature_weighting : Union[ManualStateFeatureWeighting,
-                                    UnformFeatureWeighting]
+                                    UniformFeatureWeighting]
         The method to use for weighting the state features in the loss
-        function. Defaults to uniform weighting (`UnformFeatureWeighting`, i.e.
-        all features are weighted equally).
+        function. Defaults to uniform weighting (`UniformFeatureWeighting`,
+        i.e. all features are weighted equally).
     """
 
     state_feature_weighting: Union[
@@ -112,31 +133,35 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
 
     Attributes
     ----------
-    datastore : DatastoreSelection
-        The configuration for the datastore to use.
+    datastores : Dict[str, DatastoreSelection]
+        Mapping from user-chosen datastore name to its selection and role
+        declaration. The dict key becomes the canonical source name used
+        throughout the pipeline (in ``ForecastBatch`` field dict keys, in
+        weight / clamping config keys when collisions need to be
+        disambiguated, etc).
     training : TrainingConfig
         The configuration for training the model.
     """
 
-    datastore: DatastoreSelection
-    datastore_boundary: Union[DatastoreSelection, None] = None
+    datastores: Dict[str, DatastoreSelection]
     training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)
 
     class _(dataclass_wizard.JSONWizard.Meta):
         """
         Define the configuration class as a JSON wizard class.
 
-        Together `tag_key` and `auto_assign_tags` enable that when a `Union` of
-        types are used for an attribute, the specific type to deserialize to
-        can be specified in the serialised data using the `tag_key` value. In
-        our case we call the tag key `__config_class__` to indicate to the
-        user that they should pick a dataclass describing configuration in
-        neural-lam. This Union-based selection allows us to support different
-        configuration attributes for different choices of methods for example
-        and is used when picking between different feature weighting methods in
-        the `TrainingConfig` class. `auto_assign_tags` is set to True to
-        automatically set that tag key (i.e. `__config_class__` in the config
-        file) should just be the class name of the dataclass to deserialize to.
+        Together `tag_key` and `auto_assign_tags` enable that when a `Union`
+        of types are used for an attribute, the specific type to deserialize
+        to can be specified in the serialised data using the `tag_key`
+        value. In our case we call the tag key `__config_class__` to
+        indicate to the user that they should pick a dataclass describing
+        configuration in neural-lam. This Union-based selection allows us
+        to support different configuration attributes for different choices
+        of methods for example and is used when picking between different
+        feature weighting methods in the `TrainingConfig` class.
+        `auto_assign_tags` is set to True to automatically set that tag key
+        (i.e. `__config_class__` in the config file) should just be the
+        class name of the dataclass to deserialize to.
         """
 
         tag_key = "__config_class__"
@@ -145,25 +170,62 @@ class _(dataclass_wizard.JSONWizard.Meta):
         # dataclasses used
         # TODO: this should be enabled once
         # https://github.com/rnag/dataclass-wizard/issues/137 is fixed, but
-        # currently cannot be used together with `auto_assign_tags` due to a
-        # bug it seems
+        # currently cannot be used together with `auto_assign_tags` due to
+        # a bug it seems
         # raise_on_unknown_json_key = True
 
 
 class InvalidConfigError(Exception):
     pass
 
 
+def _validate_output_name_collisions(
+    datastores: Dict[str, Union[MDPDatastore, NpyFilesDatastoreMEPS]],
+    selections: Dict[str, DatastoreSelection],
+) -> None:
+    """Raise :class:`InvalidConfigError` if two datastores would contribute
+    a variable with the same name to the model's output set, since
+    downstream sites (weight dicts, metric keys, saved zarr coords) cannot
+    disambiguate.
+
+    Fix is to give the colliding variable a unique name in one of the
+    contributing zarrs (mdp's ``dim_mapping.name_format`` for new builds,
+    or ``xr.Dataset.assign_coords`` on the small ``{category}_feature``
+    coord array of an existing zarr - a milliseconds operation regardless
+    of zarr size).
+    """
+    seen: Dict[str, str] = {}
+    for ds_name, sel in selections.items():
+        if sel.outputs is None:
+            continue
+        for category, var_list in sel.outputs.items():
+            if var_list is None:
+                var_list = datastores[ds_name].get_vars_names(category)
+            for var in var_list:
+                if var in seen:
+                    raise InvalidConfigError(
+                        f"Variable '{var}' is declared as an output in "
+                        f"both datastores '{seen[var]}' and '{ds_name}'. "
+                        "Rename the variable in one of the source zarrs "
+                        "(via mdp's `dim_mapping.name_format` for new "
+                        "builds, or via `xr.Dataset.assign_coords` on the "
+                        "existing zarr's `{category}_feature` coord - a "
+                        "milliseconds operation regardless of zarr size). "
+                        "See mllam/neural-lam#652."
+                    )
+                seen[var] = ds_name
+
+
 def load_config_and_datastore(
     config_path: str,
 ) -> tuple[
     NeuralLAMConfig,
-    Union[MDPDatastore, NpyFilesDatastoreMEPS],
-    Union[MDPDatastore, NpyFilesDatastoreMEPS, None],
+    Dict[str, Union[MDPDatastore, NpyFilesDatastoreMEPS]],
 ]:
-    """
-    Load the neural-lam configuration and the datastores specified in the
-    configuration.
+    """Load the neural-lam configuration and instantiate each datastore.
+
+    The configuration uses the multi-datastore schema introduced for #652:
+    a top-level ``datastores`` mapping with one entry per source.
 
     Parameters
     ----------
@@ -172,9 +234,11 @@ def load_config_and_datastore(
 
     Returns
     -------
-    tuple[NeuralLAMConfig, datastore, datastore_boundary]
-        The Neural-LAM configuration, the loaded (interior) datastore,
-        and the boundary datastore (or None if not configured).
+    config : NeuralLAMConfig
+        The parsed configuration.
+    datastores : Dict[str, BaseDatastore]
+        Mapping from each user-chosen datastore name to the loaded
+        datastore object, in the same order as declared in the config.
     """
     try:
         config = NeuralLAMConfig.from_yaml_file(config_path)
@@ -183,22 +247,21 @@ def load_config_and_datastore(
             "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
-    )
 
-    datastore_boundary = None
-    if config.datastore_boundary is not None:
-        datastore_boundary_config_path = (
-            Path(config_path).parent / config.datastore_boundary.config_path
+    if not config.datastores:
+        raise InvalidConfigError(
+            f"Configuration at {config_path} declares no datastores. "
+            "Add at least one entry under the top-level `datastores:` key."
         )
-        datastore_boundary = init_datastore(
-            datastore_kind=config.datastore_boundary.kind,
-            config_path=datastore_boundary_config_path,
+
+    config_dir = Path(config_path).parent
+    loaded: Dict[str, Union[MDPDatastore, NpyFilesDatastoreMEPS]] = {}
+    for name, selection in config.datastores.items():
+        datastore_config_path = config_dir / selection.config_path
+        loaded[name] = init_datastore(
+            datastore_kind=selection.kind,
+            config_path=datastore_config_path,
         )
 
-    return config, datastore, datastore_boundary
+    _validate_output_name_collisions(loaded, config.datastores)
+    return config, loaded

neural_lam/models/module.py

@@ -243,10 +243,17 @@ def _create_dataarray_from_tensor(
         split: str,
         category: str,
     ) -> xr.DataArray:
-        weather_dataset = WeatherDataset(datastore=self.datastore, split=split)
+        # Use the staticmethod variant so we don't instantiate a full
+        # WeatherDataset (which now requires the multi-datastore dict from
+        # #652) just to build a single DataArray. The reference dataarray
+        # only needs the per-grid coords from the datastore.
+        reference = self.datastore.get_dataarray(category=category, split=split)
         time = np.array(time.cpu(), dtype="datetime64[ns]")
-        da = weather_dataset.create_dataarray_from_tensor(
-            tensor=tensor, time=time, category=category
+        da = WeatherDataset.build_dataarray_from_tensor(
+            reference_dataarray=reference,
+            tensor=tensor,
+            time=time,
+            category=category,
         )
         return da
 

neural_lam/train_model.py

@@ -391,8 +391,18 @@ def main(input_args=None):
     seed.seed_everything(args.seed)
 
     # Load neural-lam configuration and datastores to use
-    config, datastore, datastore_boundary = load_config_and_datastore(
-        config_path=args.config_path
+    config, datastores = load_config_and_datastore(config_path=args.config_path)
+
+    # Resolve interior + boundary roles for the legacy single-source
+    # model side (ForecasterModule, predictor). Multi-source consumption
+    # on the model side is tracked in #652.
+    # Local
+    from .weather_dataset import _resolve_datastore_roles
+
+    interior_name, boundary_name = _resolve_datastore_roles(config.datastores)
+    datastore = datastores[interior_name]
+    datastore_boundary = (
+        datastores[boundary_name] if boundary_name is not None else None
     )
 
     # Check --var_leads_metrics_watch variable indices against the datastore
@@ -408,16 +418,16 @@ def main(input_args=None):
                 f"{len(state_var_names)} state variables)."
             )
 
-    # Create datamodule
+    # Create datamodule - takes the full multi-source dicts
     data_module = WeatherDataModule(
-        datastore=datastore,
+        datastores=datastores,
+        selections=config.datastores,
         ar_steps_train=args.ar_steps_train,
         ar_steps_eval=args.ar_steps_eval,
         num_past_forcing_steps=args.num_past_forcing_steps,
         num_future_forcing_steps=args.num_future_forcing_steps,
         num_past_boundary_steps=args.num_past_boundary_steps,
         num_future_boundary_steps=args.num_future_boundary_steps,
-        datastore_boundary=datastore_boundary,
         load_single_member=args.load_single_member,
         batch_size=args.batch_size,
         num_workers=args.num_workers,