feat(config): adopt #652 multi-datastore schema (config + WeatherDataset constructor only)

.gitignore

@@ -93,3 +93,6 @@ build/
 *.egg-info/
 
 tests/test_outputs/
+
+# Local-only drafts for GitHub comments / reviews; not part of the repo.
+.github/draft-*.md

CHANGELOG.md

@@ -23,6 +23,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Adopt the multi-datastore configuration schema from #652: replace the
+  single `datastore:` top-level field with a `datastores:` mapping of
+  named selections, each carrying optional `inputs:` / `outputs:`
+  declarations. `WeatherDataset` and `WeatherDataModule` now take
+  `datastores` and `selections` dicts; `load_config_and_datastore`
+  returns `(config, Dict[str, BaseDatastore])`. A config-time validator
+  rejects name collisions between datastores' declared outputs. Boundary
+  forcing and multi-source consumption land on the model side in a
+  follow-up that does not require changes to the config or
+  dataset shape introduced here.
+  [\#656](https://github.com/mllam/neural-lam/pull/656) @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

neural_lam/config.py

@@ -26,18 +26,17 @@ class DatastoreSelection:
         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.
     """
 
     kind: str
+    config_path: str
 
     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 +88,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,30 +111,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. The dict
+        key becomes the canonical source name used throughout the pipeline.
+        This PR ships only the dict shape; multi-source support (more than
+        one entry) lands together with the per-category `inputs`/`outputs`
+        filtering follow-up - see mllam/neural-lam#652.
     training : TrainingConfig
         The configuration for training the model.
     """
 
-    datastore: DatastoreSelection
+    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__"
@@ -144,8 +148,8 @@ 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
 
 
@@ -155,10 +159,17 @@ class InvalidConfigError(Exception):
 
 def load_config_and_datastore(
     config_path: str,
-) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]:
-    """
-    Load the neural-lam configuration and the datastore specified in the
-    configuration.
+) -> tuple[
+    NeuralLAMConfig,
+    Dict[str, Union[MDPDatastore, NpyFilesDatastoreMEPS]],
+]:
+    """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. This PR
+    accepts the dict shape but enforces exactly one entry; multi-source
+    support lands together with per-category variable filtering in a
+    follow-up.
 
     Parameters
     ----------
@@ -167,8 +178,11 @@ def load_config_and_datastore(
 
     Returns
     -------
-    tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]
-        The Neural-LAM configuration and the loaded datastore.
+    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)
@@ -177,12 +191,27 @@ 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
-    )
 
-    return config, datastore
+    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."
+        )
+    if len(config.datastores) != 1:
+        raise InvalidConfigError(
+            "This release accepts exactly one datastore under "
+            "`datastores:`. Multi-source support lands together with the "
+            "per-category `inputs`/`outputs` filtering follow-up "
+            "(see mllam/neural-lam#652)."
+        )
+
+    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, loaded

neural_lam/create_graph.py

@@ -603,8 +603,11 @@ def cli(input_args: Optional[list[str]] = None) -> None:
     if args.config_path is 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)
+    # Load neural-lam configuration and datastore to use. This PR enforces
+    # a single-entry `datastores` dict; take the only one for graph
+    # building.
+    _, datastores = load_config_and_datastore(config_path=args.config_path)
+    datastore = next(iter(datastores.values()))
 
     create_graph_from_datastore(
         datastore=datastore,

neural_lam/models/module.py

@@ -207,10 +207,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 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/plot_graph.py

@@ -260,9 +260,10 @@ def main() -> None:
     )
 
     args = parser.parse_args()
-    _, datastore = load_config_and_datastore(
+    _, datastores = load_config_and_datastore(
         config_path=args.datastore_config_path
     )
+    datastore = next(iter(datastores.values()))
 
     xy = datastore.get_xy("state", stacked=True)  # (N_grid, 2)
     pos_max = np.max(np.abs(xy))

neural_lam/train_model.py

@@ -377,8 +377,12 @@ def main(input_args=None):
     # Set seed
     seed.seed_everything(args.seed)
 
-    # Load neural-lam configuration and datastore to use
-    config, datastore = load_config_and_datastore(config_path=args.config_path)
+    # Load neural-lam configuration and datastore to use. The schema is
+    # the multi-datastore dict from #652 but this PR enforces exactly one
+    # entry (single-source); take the only datastore as the legacy
+    # single-source view for ForecasterModule and the predictor.
+    config, datastores = load_config_and_datastore(config_path=args.config_path)
+    datastore = next(iter(datastores.values()))
 
     # Check --var_leads_metrics_watch variable indices against the datastore
     # so users get an immediate error instead of an IndexError deep in the
@@ -393,9 +397,13 @@ def main(input_args=None):
                 f"{len(state_var_names)} state variables)."
             )
 
-    # Create datamodule
+    # Create datamodule - takes the full multi-source dicts so the
+    # follow-up that surfaces non-interior datastores in the per-sample
+    # tuple (mllam/neural-lam#652) only touches WeatherDataset / model
+    # internals, not this call site.
     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,