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/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..e0ede22b6 100644 --- a/neural_lam/config.py +++ b/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 diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index c0f47f75d..004405b46 100644 --- a/neural_lam/create_graph.py +++ b/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, 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..72b4385b4 100644 --- a/neural_lam/plot_graph.py +++ b/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)) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index dfb199f2e..78976c6d5 100644 --- a/neural_lam/train_model.py +++ b/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, diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index bfc956250..80cbd31dc 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,29 @@ import xarray as xr # First-party +from neural_lam.config import DatastoreSelection from neural_lam.datastore.base import BaseDatastore class WeatherDataset(torch.utils.data.Dataset): """Dataset class for weather data. - This class loads and processes weather data from a given datastore. + 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 + `_. 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`. Must contain + exactly one entry today. + selections : Dict[str, DatastoreSelection] + The matching :class:`DatastoreSelection` configs, with the same + keys. split : str, optional The data split to use ("train", "val" or "test"). Default is "train". ar_steps : int, optional @@ -45,7 +56,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 +66,20 @@ 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 + # 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 self.datastore = datastore @@ -531,33 +557,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, + ) - 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. + @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``. + + 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 +639,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 +657,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 +673,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 +685,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 +705,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/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" 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..03a486b55 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -8,6 +8,17 @@ from neural_lam.train_model import main +def _make_fake_config_and_datastores(): + """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() + config.datastores = {"interior": MagicMock()} + datastores = {"interior": MagicMock()} + return config, datastores + + @pytest.mark.parametrize( "eval_val,load_val,expect_warning", [ @@ -70,7 +81,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