From f40dea699ed24c0b4c4219144ea67c9144909979 Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 9 Jun 2026 09:19:05 +0200 Subject: [PATCH 1/4] feat: adopt #652 multi-datastore schema (rebased onto main) Replace the single `datastore:` top-level config field with a `datastores:` mapping keyed by user-chosen names. Each entry is a DatastoreSelection with optional per-category `inputs` / `outputs` declarations; one datastore must declare outputs (the interior / prognostic source) and zero or more may contribute input-only sources that are reserved for the model-side multi-source consumption (the #652 follow-up). WeatherDataset and WeatherDataModule take `datastores` and `selections` dicts; their per-sample return shape and the model unpack are unchanged from current main, so this is a config + data-loader constructor refactor only. Internally the dataset still operates on the interior datastore alone. load_config_and_datastore returns (config, Dict[str, BaseDatastore]). A config-time validator rejects two datastores declaring the same output variable name, with an error message pointing at mdp's `dim_mapping.name_format` and `xr.Dataset.assign_coords` as the two ways to disambiguate. Other callers updated: - train_model.py resolves interior + boundary roles for the legacy single-source model side. - create_graph.py and plot_graph.py resolve the interior datastore via `_resolve_datastore_roles` instead of the old 2-tuple. - module.py refactors `_create_dataarray_from_tensor` to use a new `WeatherDataset.build_dataarray_from_tensor` staticmethod so the model doesn't need to instantiate a full WeatherDataset with the new dict signature. This PR is an alternative to #635: it adopts the public schema proposed in #652 without bringing in #635's internal boundary loading. Boundary forcing, multi-source inputs and diagnostic outputs land via the #652 model-side follow-up. Co-Authored-By: Claude Opus 4.7 --- CHANGELOG.md | 12 ++ neural_lam/config.py | 157 +++++++++++---- neural_lam/create_graph.py | 11 +- neural_lam/models/module.py | 13 +- neural_lam/plot_graph.py | 7 +- neural_lam/train_model.py | 21 +- neural_lam/weather_dataset.py | 184 ++++++++++++++---- tests/conftest.py | 15 ++ .../mdp/danra_100m_winds/config.yaml | 7 +- tests/dummy_datastore.py | 1 + tests/test_checkpoint.py | 10 +- tests/test_clamping.py | 8 +- tests/test_config.py | 26 ++- tests/test_datasets.py | 58 ++++-- tests/test_gnn_layers.py | 10 +- tests/test_gpu_normalization.py | 14 +- tests/test_plotting.py | 43 ++-- tests/test_prediction_model_classes.py | 16 +- tests/test_time_slicing.py | 10 +- tests/test_train_model_warnings.py | 17 +- tests/test_training.py | 12 +- 21 files changed, 487 insertions(+), 165 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dba2ba114..a0e33e06e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/neural_lam/config.py b/neural_lam/config.py index f4195ec36..3e460200c 100644 --- a/neural_lam/config.py +++ b/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,7 +18,8 @@ @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 ---------- @@ -26,18 +27,38 @@ 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. + 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,30 +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 (e.g. as keys in future per-source + ``ForecastBatch`` field dicts, or to disambiguate weight / clamping + config when variable names collide between sources). 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 +170,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 @@ -153,12 +179,53 @@ 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]]: - """ - 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. Parameters ---------- @@ -167,8 +234,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 +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 - ) - 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." + ) + + 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, + ) + + _validate_output_name_collisions(loaded, config.datastores) + return config, loaded diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index c0f47f75d..86f59fffb 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -603,8 +603,15 @@ 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. The graph is + # built from the interior (output-producing) datastore; auxiliary + # input-only sources do not contribute graph nodes today. + # Local + from .weather_dataset import _resolve_datastore_roles + + config, datastores = load_config_and_datastore(config_path=args.config_path) + interior_name, _ = _resolve_datastore_roles(config.datastores) + datastore = datastores[interior_name] create_graph_from_datastore( datastore=datastore, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 215edab86..9977a0f9a 100644 --- a/neural_lam/models/module.py +++ b/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 diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py index f79db4ae8..37b6c58db 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -260,9 +260,14 @@ def main() -> None: ) args = parser.parse_args() - _, datastore = load_config_and_datastore( + # Local + from .weather_dataset import _resolve_datastore_roles + + config, datastores = load_config_and_datastore( config_path=args.datastore_config_path ) + interior_name, _ = _resolve_datastore_roles(config.datastores) + datastore = datastores[interior_name] xy = datastore.get_xy("state", stacked=True) # (N_grid, 2) pos_max = np.max(np.abs(xy)) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index dfb199f2e..ff60b07c3 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -377,8 +377,17 @@ 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 datastores to use + config, datastores = load_config_and_datastore(config_path=args.config_path) + + # Resolve the interior (output-producing) datastore for legacy + # single-source consumers (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] # 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 +402,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, diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index bfc956250..1f0841dc4 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -1,7 +1,7 @@ # Standard library import datetime import warnings -from typing import Iterator, Optional, Union +from typing import Any, Dict, Iterator, Optional, Union # Third-party import numpy as np @@ -10,18 +10,86 @@ import xarray as xr # First-party +from neural_lam.config import DatastoreSelection, InvalidConfigError from neural_lam.datastore.base import BaseDatastore -class WeatherDataset(torch.utils.data.Dataset): - """Dataset class for weather data. +def _resolve_datastore_roles( + selections: Dict[str, DatastoreSelection], +) -> tuple[str, Optional[str]]: + """Identify the unique output-producing (interior) datastore and an + optional input-only (boundary) datastore from the multi-source + selections. + + Multi-source prediction (more than one output-producing datastore) and + multi-source inputs (more than one input-only datastore) are tracked in + `mllam/neural-lam#652 + `_ and are not + supported in this release; this function raises with a clear error + message if the configuration goes beyond the supported single-interior + + optional-single-boundary shape. + + Parameters + ---------- + selections : Dict[str, DatastoreSelection] + The datastore selections from ``NeuralLAMConfig.datastores``, + keyed by user-chosen names. + + Returns + ------- + interior_name : str + The name of the output-producing datastore. + boundary_name : str or None + The name of the input-only datastore, or ``None`` if no boundary + is configured. + """ + output_names = [ + name for name, sel in selections.items() if sel.outputs is not None + ] + if not output_names and len(selections) == 1: + # Single-source convenience: omitted `outputs` implies the lone + # datastore is the interior with all its state vars as outputs. + return next(iter(selections)), None + if len(output_names) != 1: + raise InvalidConfigError( + "Exactly one datastore must declare `outputs` in the current " + "release (the prognostic source). Multi-source prediction is " + f"tracked in #652. Got output-producing datastores: {output_names}." + ) + interior_name = output_names[0] + input_only = [n for n in selections if n != interior_name] + if len(input_only) > 1: + raise InvalidConfigError( + "At most one input-only (boundary) datastore is supported in " + "the current release. Multi-source inputs are tracked in #652. " + f"Got input-only datastores: {input_only}." + ) + return interior_name, (input_only[0] if input_only else None) - This class loads and processes weather data from a given datastore. + +class WeatherDataset(torch.utils.data.Dataset): + """Dataset class for weather data with multi-datastore inputs. + + The dataset takes a dict of loaded datastores and a parallel dict of + :class:`DatastoreSelection` configs declaring how each one is consumed. + Exactly one datastore must produce outputs (the "interior" / + prognostic source); zero or more may contribute inputs only. This + release still operates on the interior datastore only and ignores any + input-only datastores in the per-sample return - boundary forcing and + other auxiliary sources land on the model side via + `mllam/neural-lam#652 + `_. Parameters ---------- - datastore : BaseDatastore - The datastore to load the data from (e.g. mdp). + datastores : Dict[str, BaseDatastore] + The loaded datastores, keyed by their user-chosen names. Typically + the return value of + :func:`neural_lam.config.load_config_and_datastore`. + selections : Dict[str, DatastoreSelection] + The matching :class:`DatastoreSelection` configs, with the same + keys. The ``outputs`` field on each selection determines which + datastore is the interior; the others are input-only. split : str, optional The data split to use ("train", "val" or "test"). Default is "train". ar_steps : int, optional @@ -45,7 +113,8 @@ class WeatherDataset(torch.utils.data.Dataset): def __init__( self, - datastore: BaseDatastore, + datastores: Dict[str, BaseDatastore], + selections: Dict[str, DatastoreSelection], split: str = "train", ar_steps: int = 3, num_past_forcing_steps: int = 1, @@ -54,6 +123,18 @@ def __init__( ) -> None: super().__init__() + self._datastores = datastores + self._selections = selections + self._interior_name, self._boundary_name = _resolve_datastore_roles( + selections + ) + + # The legacy single-source ``self.datastore`` attribute is kept as + # the interior alias so the within-class slicing/windowing code and + # external callers (model side, plotting) keep working unchanged. + # Multi-source consumption is the #652 follow-up. + datastore = datastores[self._interior_name] + self.split = split self.ar_steps = ar_steps self.datastore = datastore @@ -531,33 +612,57 @@ def create_dataarray_from_tensor( time: Union[datetime.datetime, list[datetime.datetime]], category: str, ): + """Instance-method wrapper around :meth:`build_dataarray_from_tensor` + that uses this dataset's already-loaded ``da_{category}`` reference + as the coord source. """ - Construct a xarray.DataArray from a `pytorch.Tensor` with coordinates - for `grid_index`, `time` and `{category}_feature` matching the shape - and number of times provided and add the x/y coordinates from the - datastore. + return self.build_dataarray_from_tensor( + reference_dataarray=getattr(self, f"da_{category}"), + tensor=tensor, + time=time, + category=category, + ) + + @staticmethod + def build_dataarray_from_tensor( + reference_dataarray: xr.DataArray, + tensor: torch.Tensor, + time: Union[datetime.datetime, list[datetime.datetime]], + category: str, + ): + """Construct an :class:`xr.DataArray` from a :class:`torch.Tensor` + with coordinates for ``grid_index``, ``time`` and + ``{category}_feature`` matching the shape and number of times + provided, taking the per-grid coords from ``reference_dataarray``. - The number if times provided is expected to match the shape of the - tensor. For a 2D tensor, the dimensions are assumed to be (grid_index, - {category}_feature) and only a single time should be provided. For a 3D - tensor, the dimensions are assumed to be (time, grid_index, - {category}_feature) and a list of times should be provided. + Exposed as a staticmethod so callers that have a datastore but not + a full :class:`WeatherDataset` (e.g. the model in + :mod:`neural_lam.models.module`) can build dataarrays without + instantiating the dataset. Parameters ---------- + reference_dataarray : xr.DataArray + Source for ``grid_index``, ``{category}_feature``, and the + optional ``x``/``y`` coords. Typically what the datastore + returned from ``get_dataarray(category=...)``. tensor : torch.Tensor - The tensor to construct the DataArray from, this assumed to have - the same dimension ordering as returned by the __getitem__ method - (i.e. time, grid_index, {category}_feature). The tensor will be + The tensor to construct the DataArray from. For a 2D tensor + the dimensions are assumed to be + ``(grid_index, {category}_feature)`` and a single ``time`` + should be provided. For a 3D tensor the dimensions are + assumed to be ``(time, grid_index, {category}_feature)`` and + a list of times should be provided. The tensor will be copied to the CPU before constructing the DataArray. time : datetime.datetime or list[datetime.datetime] The time or times of the tensor. category : str - The category of the tensor, either "state", "forcing" or "static". + The category of the tensor, either ``"state"``, ``"forcing"`` + or ``"static"``. Returns ------- - da : xr.DataArray + xr.DataArray The constructed DataArray. """ @@ -589,9 +694,8 @@ def _is_listlike(obj): f"{len(tensor.shape)}" ) - da_datastore_state = getattr(self, f"da_{category}") - da_grid_index = da_datastore_state.grid_index - da_state_feature = da_datastore_state.state_feature + da_grid_index = reference_dataarray.grid_index + da_state_feature = reference_dataarray.state_feature coords = { f"{category}_feature": da_state_feature, @@ -608,10 +712,10 @@ def _is_listlike(obj): for grid_coord in ["x", "y"]: if ( - grid_coord in da_datastore_state.coords + grid_coord in reference_dataarray.coords and grid_coord not in da.coords ): - da.coords[grid_coord] = da_datastore_state[grid_coord] + da.coords[grid_coord] = reference_dataarray[grid_coord] if not add_time_as_dim: da.coords["time"] = time @@ -624,7 +728,8 @@ class WeatherDataModule(pl.LightningDataModule): def __init__( self, - datastore: BaseDatastore, + datastores: Dict[str, BaseDatastore], + selections: Dict[str, DatastoreSelection], ar_steps_train: int = 3, ar_steps_eval: int = 25, num_past_forcing_steps: int = 1, @@ -635,7 +740,8 @@ def __init__( eval_split: str = "test", ) -> None: super().__init__() - self._datastore = datastore + self._datastores = datastores + self._selections = selections self.num_past_forcing_steps = num_past_forcing_steps self.num_future_forcing_steps = num_future_forcing_steps self.ar_steps_train = ar_steps_train @@ -654,32 +760,30 @@ def __init__( self.multiprocessing_context = "spawn" def setup(self, stage: Optional[str] = None) -> None: + shared_kwargs: dict[str, Any] = dict( + datastores=self._datastores, + selections=self._selections, + num_past_forcing_steps=self.num_past_forcing_steps, + num_future_forcing_steps=self.num_future_forcing_steps, + load_single_member=self.load_single_member, + ) if stage == "fit" or stage is None: self.train_dataset = WeatherDataset( - datastore=self._datastore, split="train", ar_steps=self.ar_steps_train, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) self.val_dataset = WeatherDataset( - datastore=self._datastore, split="val", ar_steps=self.ar_steps_eval, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) if stage == "test" or stage is None: self.test_dataset = WeatherDataset( - datastore=self._datastore, split=self.eval_split, ar_steps=self.ar_steps_eval, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) def train_dataloader(self) -> torch.utils.data.DataLoader: diff --git a/tests/conftest.py b/tests/conftest.py index 47237ed55..9fd135c6a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,7 @@ from pytorch_lightning.utilities import rank_zero_only # First-party +from neural_lam.config import DatastoreSelection from neural_lam.datastore import DATASTORES, init_datastore from neural_lam.datastore.npyfilesmeps import ( compute_standardization_stats as compute_standardization_stats_meps, @@ -123,3 +124,17 @@ def init_datastore_example(datastore_kind): ) return datastore + + +def make_single_source_args(datastore, name="interior"): + """Wrap a single datastore in the multi-source dicts that + ``WeatherDataset`` / ``WeatherDataModule`` now expect. + + Returns ``(datastores, selections)`` ready to splat into the + constructors. + """ + selection = DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=str(datastore.root_path), + ) + return {name: datastore}, {name: selection} diff --git a/tests/datastore_examples/mdp/danra_100m_winds/config.yaml b/tests/datastore_examples/mdp/danra_100m_winds/config.yaml index 8b3362e0e..6c8091e1c 100644 --- a/tests/datastore_examples/mdp/danra_100m_winds/config.yaml +++ b/tests/datastore_examples/mdp/danra_100m_winds/config.yaml @@ -1,6 +1,7 @@ -datastore: - kind: mdp - config_path: danra.datastore.yaml +datastores: + danra: + kind: mdp + config_path: danra.datastore.yaml training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting diff --git a/tests/dummy_datastore.py b/tests/dummy_datastore.py index b269bfb2c..6d8be9a0d 100644 --- a/tests/dummy_datastore.py +++ b/tests/dummy_datastore.py @@ -489,6 +489,7 @@ class EnsembleDummyDatastore(BaseDatastore): assert exact numeric expectations. """ + SHORT_NAME = "dummydata" T0 = np.datetime64("2021-01-01T00:00:00") def __init__( diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148b..9a7328878 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -31,10 +31,12 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): ) config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ), + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + }, ) predictor = GraphLAM( diff --git a/tests/test_clamping.py b/tests/test_clamping.py index 8c44b5688..267d1dde8 100644 --- a/tests/test_clamping.py +++ b/tests/test_clamping.py @@ -46,9 +46,11 @@ class ModelArgs: model_args = ModelArgs() config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ), + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + }, training=nlconfig.TrainingConfig( output_clamping=nlconfig.OutputClamping( lower={"t2m": 0.0, "r2m": 0.0}, diff --git a/tests/test_config.py b/tests/test_config.py index 1ff40bc6a..e8a1daab9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -16,7 +16,9 @@ ) def test_config_serialization(state_weighting_config): c = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection(kind="mdp", config_path=""), + datastores={ + "danra": nlconfig.DatastoreSelection(kind="mdp", config_path=""), + }, training=nlconfig.TrainingConfig( state_feature_weighting=state_weighting_config ), @@ -27,22 +29,26 @@ def test_config_serialization(state_weighting_config): yaml_training_defaults = """ -datastore: - kind: mdp - config_path: "" +datastores: + danra: + kind: mdp + config_path: "" """ default_config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection(kind="mdp", config_path=""), + datastores={ + "danra": nlconfig.DatastoreSelection(kind="mdp", config_path=""), + }, training=nlconfig.TrainingConfig( state_feature_weighting=nlconfig.UniformFeatureWeighting() ), ) yaml_training_manual_weights = """ -datastore: - kind: mdp - config_path: "" +datastores: + danra: + kind: mdp + config_path: "" training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting @@ -52,7 +58,9 @@ def test_config_serialization(state_weighting_config): """ manual_weights_config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection(kind="mdp", config_path=""), + datastores={ + "danra": nlconfig.DatastoreSelection(kind="mdp", config_path=""), + }, training=nlconfig.TrainingConfig( state_feature_weighting=nlconfig.ManualStateFeatureWeighting( weights=dict(u100m=1.0, v100m=1.0) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 4b35840ec..4abbd882f 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -14,7 +14,7 @@ from neural_lam.datastore.base import BaseRegularGridDatastore from neural_lam.models import ForecasterModule from neural_lam.weather_dataset import WeatherDataset -from tests.conftest import init_datastore_example +from tests.conftest import init_datastore_example, make_single_source_args from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore @@ -36,8 +36,10 @@ def test_dataset_item_shapes(datastore_name): N_pred_steps = 4 num_past_forcing_steps = 1 num_future_forcing_steps = 1 + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=N_pred_steps, num_past_forcing_steps=num_past_forcing_steps, @@ -87,8 +89,10 @@ def test_dataset_item_create_dataarray_from_tensor(datastore_name): N_pred_steps = 4 num_past_forcing_steps = 1 num_future_forcing_steps = 1 + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=N_pred_steps, num_past_forcing_steps=num_past_forcing_steps, @@ -210,13 +214,15 @@ def _create_graph(): _create_graph() - config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) - ) + datastores, selections = make_single_source_args(datastore) + config = nlconfig.NeuralLAMConfig(datastores=selections) - dataset = WeatherDataset(datastore=datastore, split=split, ar_steps=2) + dataset = WeatherDataset( + datastores=datastores, + selections=selections, + split=split, + ar_steps=2, + ) # First-party from neural_lam.models import MODELS, ARForecaster @@ -278,8 +284,10 @@ def test_dataset_length(dataset_config): ds_len = 10 datastore = DummyDatastore(n_timesteps=ds_len) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=dataset_config["ar_steps"], num_past_forcing_steps=dataset_config["past"], @@ -301,8 +309,10 @@ def test_dataset_out_of_range_raises_index_error(): """`WeatherDataset.__getitem__` raises IndexError for out-of-range indices and supports Python-style negative indexing within bounds.""" datastore = DummyDatastore(n_timesteps=10) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=1, num_past_forcing_steps=0, @@ -329,8 +339,10 @@ def test_ensemble_len_scales_with_default_all_members(): n_timesteps=10, ) + datastores, selections = make_single_source_args(datastore) dataset_all = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -338,7 +350,8 @@ def test_ensemble_len_scales_with_default_all_members(): ) dataset_single = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -387,8 +400,10 @@ def test_ensemble_index_mapping_is_time_major(): n_ensemble_members=3, n_timesteps=10, ) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -411,8 +426,10 @@ def test_ensemble_forcing_uses_same_member_when_available(): n_ensemble_members=3, n_timesteps=10, ) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -434,8 +451,10 @@ def test_ensemble_forcing_without_member_dim_is_shared(): n_ensemble_members=3, n_timesteps=10, ) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -460,8 +479,10 @@ def test_forecast_ensemble_len_scales_with_default_all_members(): n_forecast_steps=6, ) + datastores, selections = make_single_source_args(datastore) dataset_all = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, @@ -470,7 +491,8 @@ def test_forecast_ensemble_len_scales_with_default_all_members(): with pytest.warns(UserWarning, match="only using first ensemble member"): dataset_single = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=1, diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 166a9b549..94258d2c4 100644 --- a/tests/test_gnn_layers.py +++ b/tests/test_gnn_layers.py @@ -94,10 +94,12 @@ def _get_datastore_and_config(graph_name): """Create a datastore with graph already built.""" datastore = init_datastore_example("mdp") config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ) + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + } ) # Ensure graph exists diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index 8d516bfbd..bb58278b0 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -6,7 +6,7 @@ from neural_lam import config as nlconfig from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor from neural_lam.weather_dataset import WeatherDataModule -from tests.conftest import init_datastore_example +from tests.conftest import init_datastore_example, make_single_source_args NUM_PAST_FORCING_STEPS = 1 NUM_FUTURE_FORCING_STEPS = 1 @@ -21,9 +21,11 @@ def forward(self, prev_state, prev_prev_state, forcing): def _build_module(datastore): config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + } ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) forecaster = ARForecaster(predictor, datastore) @@ -81,8 +83,10 @@ def test_normalization_applied_exactly_once(): datastore = init_datastore_example("mdp") model = _build_module(datastore) + datastores, selections = make_single_source_args(datastore) data_module = WeatherDataModule( - datastore=datastore, + datastores=datastores, + selections=selections, ar_steps_train=2, ar_steps_eval=2, batch_size=2, diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 616d563de..18f33f117 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -20,7 +20,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM from neural_lam.weather_dataset import WeatherDataset -from tests.conftest import init_datastore_example +from tests.conftest import init_datastore_example, make_single_source_args from tests.dummy_datastore import DummyDatastore # Create output directory for test figures @@ -442,10 +442,12 @@ class ModelArgs: # Create config. config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ), + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + }, ) # Create model @@ -483,8 +485,10 @@ class ModelArgs: ) # Create dataset to get a sample batch. + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, split="train", ar_steps=2, num_past_forcing_steps=0, @@ -551,7 +555,12 @@ def test_plot_examples_integration_saves_figure( time_slice = batch[3][0] # Create DataArrays. - dataset = WeatherDataset(datastore=datastore, split="train") + plot_datastores, plot_selections = make_single_source_args(datastore) + dataset = WeatherDataset( + datastores=plot_datastores, + selections=plot_selections, + split="train", + ) time = np.array(time_slice.cpu(), dtype="datetime64[ns]") @@ -715,10 +724,12 @@ def test_create_metric_log_dict_with_metrics_watch(tmp_path): ) config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ), + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + }, ) model = _build_metrics_watch_module(datastore, config) @@ -774,10 +785,12 @@ def test_aggregate_and_plot_metrics_with_metrics_watch(tmp_path): ) config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ), + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ), + }, ) model = _build_metrics_watch_module(datastore, config) diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index f081bbf9a..f3e8c18d9 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -75,9 +75,11 @@ def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + } ) # Build predictor and forecaster externally, then inject into @@ -173,9 +175,11 @@ def test_forecaster_module_old_checkpoint(tmp_path): datastore = init_datastore_example("mdp") config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "interior": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + } ) # First-party diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index a8b022eda..c27d7cb19 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -10,9 +10,11 @@ # First-party from neural_lam.datastore.base import BaseDatastore from neural_lam.weather_dataset import WeatherDataset +from tests.conftest import make_single_source_args class SinglePointDummyDatastore(BaseDatastore): + SHORT_NAME = "dummydata" config = {} coords_projection = None num_grid_points = 1 @@ -106,8 +108,10 @@ def test_time_slicing_analysis( is_forecast=False, ) + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, ar_steps=ar_steps, num_future_forcing_steps=num_future_forcing_steps, num_past_forcing_steps=num_past_forcing_steps, @@ -181,8 +185,10 @@ def test_step_length_timedeltas(step_length): assert datastore.step_length == step_length # Test that WeatherDataset can be created with this datastore + datastores, selections = make_single_source_args(datastore) dataset = WeatherDataset( - datastore=datastore, + datastores=datastores, + selections=selections, ar_steps=3, num_future_forcing_steps=0, num_past_forcing_steps=0, diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a9..59a805b9c 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -8,6 +8,21 @@ from neural_lam.train_model import main +def _make_fake_config_and_datastores(): + """Build a return value for ``load_config_and_datastore`` that survives + the multi-source role resolution in ``train_model.main`` without + needing a real datastore on disk. + """ + config = MagicMock() + selection = MagicMock() + selection.outputs = None + # Single source: ``_resolve_datastore_roles`` returns ("interior", None) + # without needing ``outputs`` to be declared explicitly. + config.datastores = {"interior": selection} + datastores = {"interior": MagicMock()} + return config, datastores + + @pytest.mark.parametrize( "eval_val,load_val,expect_warning", [ @@ -70,7 +85,7 @@ def capture_init(_self, **kwargs): ), patch( "neural_lam.train_model.load_config_and_datastore", - return_value=(MagicMock(), MagicMock()), + return_value=_make_fake_config_and_datastores(), ), patch("neural_lam.train_model.WeatherDataModule"), patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), diff --git a/tests/test_training.py b/tests/test_training.py index 141681fc0..032881d37 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -88,8 +88,12 @@ def run_simple_training( n_max_levels=1, ) + selection = nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) data_module = WeatherDataModule( - datastore=datastore, + datastores={"interior": datastore}, + selections={"interior": selection}, ar_steps_train=3, ar_steps_eval=5, batch_size=2, @@ -98,11 +102,7 @@ def run_simple_training( num_future_forcing_steps=1, ) - config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) - ) + config = nlconfig.NeuralLAMConfig(datastores={"interior": selection}) # Build predictor and forecaster externally, then inject into # ForecasterModule From 538a6ff0d8b0d0e3da4a7ea4f5f33e0bdafb0274 Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 9 Jun 2026 09:29:03 +0200 Subject: [PATCH 2/4] test: register `slow` pytest marker Match the marker registration used on #635/#651 so any future @pytest.mark.slow test on this branch is recognised by pytest without warnings. No tests currently use the marker on #656. Co-Authored-By: Claude Opus 4.7 --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index bccec0828..589c499e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,6 +108,11 @@ allow-any-import-level = "neural_lam" [tool.pylint.SIMILARITIES] min-similarity-lines = 10 +[tool.pytest.ini_options] +markers = [ + "slow: marks tests as slow (deselected by default, run with -m slow)", +] + [build-system] requires = ["hatchling>=1.27.0", "hatch-vcs"] build-backend = "hatchling.build" From fd155c3a541211e83383770799430aad19ad2247 Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 9 Jun 2026 09:39:11 +0200 Subject: [PATCH 3/4] config: correct collision-error fix-hint to in-place zarr op `xr.Dataset.assign_coords` returns a new dataset and does not touch disk; the fix-hint in `_validate_output_name_collisions` claimed it was an in-place rename, which would have led users astray. Replace with a small zarr-python snippet that overwrites the `{category}_feature` coord array directly, which is the actual milliseconds-scale in-place op the message intends. Co-Authored-By: Claude Opus 4.7 --- neural_lam/config.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/neural_lam/config.py b/neural_lam/config.py index 3e460200c..69544b996 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -190,9 +190,9 @@ def _validate_output_name_collisions( 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). + or for an existing zarr overwrite the small ``{category}_feature`` + coord array directly with zarr-python: + ``zarr.open_group(path, mode="r+")["{category}_feature"][:] = new_names``). """ seen: Dict[str, str] = {} for ds_name, sel in selections.items(): @@ -208,9 +208,10 @@ def _validate_output_name_collisions( 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). " + "builds, or by overwriting the `{category}_feature` " + "coord array in the existing zarr with zarr-python: " + '`zarr.open_group(path, mode="r+")' + '["{category}_feature"][:] = new_names`). ' "See mllam/neural-lam#652." ) seen[var] = ds_name From acc4e95019369182781e40bce3676292fdd77a8d Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 9 Jun 2026 10:06:34 +0200 Subject: [PATCH 4/4] config: scope #656 to dict shape only Strip the optional `inputs` / `outputs` fields from `DatastoreSelection` and the matching `_validate_output_name_collisions` validator. They were parsed but never consumed at runtime, so reviewers would have asked what they do and the honest answer was "nothing yet". Both pieces (per-category include-lists and the output-name collision validator) belong with the data-loader filtering follow-up, which must land in lockstep with @joeloskarsson's model-side adapter so feature dimensions agree. `_resolve_datastore_roles` also goes away. With no `outputs` field to distinguish interior from boundary, role resolution would be a guess anyway. Instead `load_config_and_datastore` and `WeatherDataset` now require exactly one entry in the `datastores:` dict; multi-source support comes back with the filtering follow-up. The single entry is picked as the legacy `self.datastore` interior view used by slicing/windowing/plotting. Net effect: #656 is now just the `datastore:` -> `datastores:` dict shape rename. The diagnostic / filtering / collision-validator work each land in their own follow-up PRs. Also gitignore `.github/draft-*.md` since the comment-draft files are local working notes, not part of the repo. Co-Authored-By: Claude Opus 4.7 --- .gitignore | 3 + neural_lam/config.py | 87 ++++++--------------------- neural_lam/create_graph.py | 14 ++--- neural_lam/plot_graph.py | 8 +-- neural_lam/train_model.py | 15 ++--- neural_lam/weather_dataset.py | 95 +++++++----------------------- tests/test_train_model_warnings.py | 12 ++-- 7 files changed, 57 insertions(+), 177 deletions(-) diff --git a/.gitignore b/.gitignore index ca528fa21..d7c046bf9 100644 --- a/.gitignore +++ b/.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 diff --git a/neural_lam/config.py b/neural_lam/config.py index 69544b996..e0ede22b6 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -1,7 +1,7 @@ # Standard library import dataclasses from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Dict, Union # Third-party import dataclass_wizard @@ -18,8 +18,7 @@ @dataclasses.dataclass class DatastoreSelection: """ - Configuration for selecting a datastore and declaring how its variables - are consumed by the model. + Configuration for selecting a datastore to use with neural-lam. Attributes ---------- @@ -29,31 +28,10 @@ class DatastoreSelection: 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. - 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: @@ -134,11 +112,11 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): Attributes ---------- 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 (e.g. as keys in future per-source - ``ForecastBatch`` field dicts, or to disambiguate weight / clamping - config when variable names collide between sources). + 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. """ @@ -179,44 +157,6 @@ 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 for an existing zarr overwrite the small ``{category}_feature`` - coord array directly with zarr-python: - ``zarr.open_group(path, mode="r+")["{category}_feature"][:] = new_names``). - """ - 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 by overwriting the `{category}_feature` " - "coord array in the existing zarr with zarr-python: " - '`zarr.open_group(path, mode="r+")' - '["{category}_feature"][:] = new_names`). ' - "See mllam/neural-lam#652." - ) - seen[var] = ds_name - - def load_config_and_datastore( config_path: str, ) -> tuple[ @@ -226,7 +166,10 @@ def load_config_and_datastore( """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. + 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 ---------- @@ -254,6 +197,13 @@ def load_config_and_datastore( 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]] = {} @@ -264,5 +214,4 @@ def load_config_and_datastore( config_path=datastore_config_path, ) - _validate_output_name_collisions(loaded, config.datastores) return config, loaded diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index 86f59fffb..004405b46 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -603,15 +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. The graph is - # built from the interior (output-producing) datastore; auxiliary - # input-only sources do not contribute graph nodes today. - # Local - from .weather_dataset import _resolve_datastore_roles - - config, datastores = load_config_and_datastore(config_path=args.config_path) - interior_name, _ = _resolve_datastore_roles(config.datastores) - datastore = datastores[interior_name] + # 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, diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py index 37b6c58db..72b4385b4 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -260,14 +260,10 @@ def main() -> None: ) args = parser.parse_args() - # Local - from .weather_dataset import _resolve_datastore_roles - - config, datastores = load_config_and_datastore( + _, datastores = load_config_and_datastore( config_path=args.datastore_config_path ) - interior_name, _ = _resolve_datastore_roles(config.datastores) - datastore = datastores[interior_name] + datastore = next(iter(datastores.values())) xy = datastore.get_xy("state", stacked=True) # (N_grid, 2) pos_max = np.max(np.abs(xy)) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index ff60b07c3..78976c6d5 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -377,17 +377,12 @@ def main(input_args=None): # Set seed seed.seed_everything(args.seed) - # Load neural-lam configuration and datastores to use + # 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) - - # Resolve the interior (output-producing) datastore for legacy - # single-source consumers (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 = 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 diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 1f0841dc4..80cbd31dc 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -10,73 +10,16 @@ import xarray as xr # First-party -from neural_lam.config import DatastoreSelection, InvalidConfigError +from neural_lam.config import DatastoreSelection from neural_lam.datastore.base import BaseDatastore -def _resolve_datastore_roles( - selections: Dict[str, DatastoreSelection], -) -> tuple[str, Optional[str]]: - """Identify the unique output-producing (interior) datastore and an - optional input-only (boundary) datastore from the multi-source - selections. - - Multi-source prediction (more than one output-producing datastore) and - multi-source inputs (more than one input-only datastore) are tracked in - `mllam/neural-lam#652 - `_ and are not - supported in this release; this function raises with a clear error - message if the configuration goes beyond the supported single-interior - + optional-single-boundary shape. - - Parameters - ---------- - selections : Dict[str, DatastoreSelection] - The datastore selections from ``NeuralLAMConfig.datastores``, - keyed by user-chosen names. - - Returns - ------- - interior_name : str - The name of the output-producing datastore. - boundary_name : str or None - The name of the input-only datastore, or ``None`` if no boundary - is configured. - """ - output_names = [ - name for name, sel in selections.items() if sel.outputs is not None - ] - if not output_names and len(selections) == 1: - # Single-source convenience: omitted `outputs` implies the lone - # datastore is the interior with all its state vars as outputs. - return next(iter(selections)), None - if len(output_names) != 1: - raise InvalidConfigError( - "Exactly one datastore must declare `outputs` in the current " - "release (the prognostic source). Multi-source prediction is " - f"tracked in #652. Got output-producing datastores: {output_names}." - ) - interior_name = output_names[0] - input_only = [n for n in selections if n != interior_name] - if len(input_only) > 1: - raise InvalidConfigError( - "At most one input-only (boundary) datastore is supported in " - "the current release. Multi-source inputs are tracked in #652. " - f"Got input-only datastores: {input_only}." - ) - return interior_name, (input_only[0] if input_only else None) - - class WeatherDataset(torch.utils.data.Dataset): - """Dataset class for weather data with multi-datastore inputs. - - The dataset takes a dict of loaded datastores and a parallel dict of - :class:`DatastoreSelection` configs declaring how each one is consumed. - Exactly one datastore must produce outputs (the "interior" / - prognostic source); zero or more may contribute inputs only. This - release still operates on the interior datastore only and ignores any - input-only datastores in the per-sample return - boundary forcing and - other auxiliary sources land on the model side via + """Dataset class for weather data. + + The dataset takes a single-entry dict of loaded datastores keyed by the + user-chosen name. Multi-source consumption (more than one datastore) + lands together with the per-category variable-filtering follow-up - see `mllam/neural-lam#652 `_. @@ -85,11 +28,11 @@ class WeatherDataset(torch.utils.data.Dataset): datastores : Dict[str, BaseDatastore] The loaded datastores, keyed by their user-chosen names. Typically the return value of - :func:`neural_lam.config.load_config_and_datastore`. + :func:`neural_lam.config.load_config_and_datastore`. Must contain + exactly one entry today. selections : Dict[str, DatastoreSelection] The matching :class:`DatastoreSelection` configs, with the same - keys. The ``outputs`` field on each selection determines which - datastore is the interior; the others are input-only. + keys. split : str, optional The data split to use ("train", "val" or "test"). Default is "train". ar_steps : int, optional @@ -123,17 +66,19 @@ def __init__( ) -> None: super().__init__() + if len(datastores) != 1: + raise ValueError( + "WeatherDataset expects exactly one datastore in the dict; " + "multi-source support lands together with the per-category " + "`inputs`/`outputs` filtering follow-up (mllam/neural-lam#652)." + ) + self._datastores = datastores self._selections = selections - self._interior_name, self._boundary_name = _resolve_datastore_roles( - selections - ) - - # The legacy single-source ``self.datastore`` attribute is kept as - # the interior alias so the within-class slicing/windowing code and - # external callers (model side, plotting) keep working unchanged. - # Multi-source consumption is the #652 follow-up. - datastore = datastores[self._interior_name] + # Take the only datastore as the interior alias used by the + # within-class slicing/windowing code and external callers (model + # side, plotting). + datastore = next(iter(datastores.values())) self.split = split self.ar_steps = ar_steps diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index 59a805b9c..03a486b55 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -9,16 +9,12 @@ def _make_fake_config_and_datastores(): - """Build a return value for ``load_config_and_datastore`` that survives - the multi-source role resolution in ``train_model.main`` without - needing a real datastore on disk. + """Build a return value for ``load_config_and_datastore`` that train_model + can consume without needing a real datastore on disk. The single-entry + dict matches the single-source enforcement on this branch. """ config = MagicMock() - selection = MagicMock() - selection.outputs = None - # Single source: ``_resolve_datastore_roles`` returns ("interior", None) - # without needing ``outputs`` to be declared explicitly. - config.datastores = {"interior": selection} + config.datastores = {"interior": MagicMock()} datastores = {"interior": MagicMock()} return config, datastores