From dee15d15df919f9bbf7816c44050dfe2a5d6824f Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 16:00:11 +0200 Subject: [PATCH 01/35] feat: add optional boundary datastore support in WeatherDataset Add support for loading boundary forcing from a separate datastore, enabling LAM models to ingest boundary conditions from a different domain (e.g. ERA5 boundaries for a COSMO/DANRA interior). - NeuralLAMConfig accepts optional `datastore_boundary` field - load_config_and_datastore returns 3-tuple (config, datastore, datastore_boundary) - WeatherDataset loads, windows, and standardizes boundary forcing - __getitem__ returns 5-tuple (init_states, target_states, forcing, boundary, target_times) - New CLI args --num_past_boundary_steps / --num_future_boundary_steps - ForecasterModule.common_step unpacks boundary (not yet wired to forward) - 4 new boundary-specific tests, all 157 tests pass refs #108 Co-Authored-By: Claude Opus 4.6 --- AGENTS.md | 12 +-- CHANGELOG.md | 4 + README.md | 11 +- neural_lam/config.py | 26 ++++- neural_lam/create_graph.py | 2 +- neural_lam/models/module.py | 10 +- neural_lam/plot_graph.py | 2 +- neural_lam/train_model.py | 22 +++- neural_lam/weather_dataset.py | 166 +++++++++++++++++++++++++---- tests/dummy_datastore.py | 19 ++++ tests/test_datasets.py | 153 +++++++++++++++++++++++--- tests/test_plotting.py | 4 +- tests/test_time_slicing.py | 6 +- tests/test_train_model_warnings.py | 2 +- 14 files changed, 380 insertions(+), 59 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dcf05d563..c295c3fc0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,15 +9,15 @@ Mandatory rules for AI coding agents. Violations will result in rejected PRs. Neural-LAM: graph-based neural weather prediction for Limited Area Modeling. Models: `GraphLAM`, `HiLAM`, `HiLAMParallel`. -**Data flow:** Raw zarr/numpy → `Datastore` → `WeatherDataset` → `WeatherDataModule` → Model → -Predictions +**Data flow:** Raw zarr/numpy → `Datastore` (+ optional boundary `Datastore`) → `WeatherDataset` → +`WeatherDataModule` → Model → Predictions **Key modules:** - `datastore/` — `BaseDatastore` (abstract), `MDPDatastore` (zarr via mllam-data-prep) -- `models/` — `ARModel` (autoregressive base, Lightning) → `BaseGraphModel` (encode-process-decode) - → `GraphLAM` / `HiLAM` / `HiLAMParallel` -- `weather_dataset.py` — `WeatherDataset` + `WeatherDataModule` -- `config.py` — YAML config via dataclass-wizard +- `models/` — `ForecasterModule` (Lightning) → `ARForecaster` (Forecaster) → + `GraphLAM` / `HiLAM` / `HiLAMParallel` (StepPredictor) +- `weather_dataset.py` — `WeatherDataset` + `WeatherDataModule` (supports optional boundary datastore) +- `config.py` — YAML config via dataclass-wizard (`NeuralLAMConfig` with optional `datastore_boundary`) - `create_graph.py` — builds mesh graphs (must run before training) - `interaction_net.py` — `InteractionNet` GNN layer (PyG `MessagePassing`) - `utils.py` — `make_mlp`, normalization helpers diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d8d493ef..f961ae633 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [unreleased](https://github.com/mllam/neural-lam/compare/v0.6.0...HEAD) +### Added + +- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. [\#TODO](https://github.com/mllam/neural-lam/pull/TODO) @sadamov + ### Changed - Split `ARModel` into `ForecasterModule`, `Forecaster` and diff --git a/README.md b/README.md index 17e269841..81f0d0a6f 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,9 @@ And the content of `config.yaml` could in this case look like: datastore: kind: mdp config_path: danra.datastore.yaml +datastore_boundary: # optional, for boundary forcing from a separate domain + kind: mdp + config_path: era5_boundary.datastore.yaml training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting @@ -167,10 +170,14 @@ training: For now the neural-lam config only defines few things: 1. The kind of datastore and the path to its config -2. The weighting of different features in +2. (Optional) A boundary datastore (`datastore_boundary`) providing boundary + forcing from a separate domain (e.g. ERA5 for a LAM domain). When set, the + boundary forcing is windowed and included as an additional tensor in each + training sample. +3. The weighting of different features in the loss function. If you don't define the state feature weighting it will default to weighting all features equally. -3. Valid numerical range for output of each feature. The numerical range of all features default to $]-\infty, \infty[$. +4. Valid numerical range for output of each feature. The numerical range of all features default to $]-\infty, \infty[$. (This example is taken from the `tests/datastore_examples/mdp` directory.) diff --git a/neural_lam/config.py b/neural_lam/config.py index f4195ec36..e74abed0c 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -119,6 +119,7 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): """ datastore: DatastoreSelection + datastore_boundary: Union[DatastoreSelection, None] = None training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig) class _(dataclass_wizard.JSONWizard.Meta): @@ -155,9 +156,13 @@ class InvalidConfigError(Exception): def load_config_and_datastore( config_path: str, -) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]: +) -> tuple[ + NeuralLAMConfig, + Union[MDPDatastore, NpyFilesDatastoreMEPS], + Union[MDPDatastore, NpyFilesDatastoreMEPS, None], +]: """ - Load the neural-lam configuration and the datastore specified in the + Load the neural-lam configuration and the datastores specified in the configuration. Parameters @@ -167,8 +172,9 @@ def load_config_and_datastore( Returns ------- - tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]] - The Neural-LAM configuration and the loaded datastore. + tuple[NeuralLAMConfig, datastore, datastore_boundary] + The Neural-LAM configuration, the loaded (interior) datastore, + and the boundary datastore (or None if not configured). """ try: config = NeuralLAMConfig.from_yaml_file(config_path) @@ -185,4 +191,14 @@ def load_config_and_datastore( datastore_kind=config.datastore.kind, config_path=datastore_config_path ) - return config, datastore + datastore_boundary = None + if config.datastore_boundary is not None: + datastore_boundary_config_path = ( + Path(config_path).parent / config.datastore_boundary.config_path + ) + datastore_boundary = init_datastore( + datastore_kind=config.datastore_boundary.kind, + config_path=datastore_boundary_config_path, + ) + + return config, datastore, datastore_boundary diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index c0f47f75d..4497531a4 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -604,7 +604,7 @@ def cli(input_args: Optional[list[str]] = None) -> None: raise ValueError("Specify your config with --config_path") # Load neural-lam configuration and datastore to use - _, datastore = load_config_and_datastore(config_path=args.config_path) + _, datastore, _ = load_config_and_datastore(config_path=args.config_path) create_graph_from_datastore( datastore=datastore, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index ce6d06cfb..12f2d5544 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -176,7 +176,13 @@ def configure_optimizers(self): return opt def common_step(self, batch): - init_states, target_states, forcing_features, batch_times = batch + ( + init_states, + target_states, + forcing_features, + boundary_features, + batch_times, + ) = batch prediction, pred_std = self.forecaster( init_states, forcing_features, target_states ) @@ -374,7 +380,7 @@ def test_step(self, batch, batch_idx): def plot_examples(self, batch, n_examples, split, prediction): target = batch[1] - time = batch[3] + time = batch[4] da_state_stats = self.datastore.get_standardization_dataarray("state") state_std = torch.tensor( diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py index 6ed7d0268..99f4a0783 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -259,7 +259,7 @@ def main(): ) args = parser.parse_args() - _, datastore = load_config_and_datastore( + _, datastore, _ = load_config_and_datastore( config_path=args.datastore_config_path ) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 6db282136..ddcf72e50 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -270,6 +270,19 @@ def main(input_args=None): default=1, help="Number of future time steps to use as input for forcing data", ) + parser.add_argument( + "--num_past_boundary_steps", + type=int, + default=1, + help="Number of past time steps to use as input for boundary forcing", + ) + parser.add_argument( + "--num_future_boundary_steps", + type=int, + default=1, + help="Number of future time steps to use as input for boundary " + "forcing", + ) parser.add_argument( "--load_single_member", action="store_true", @@ -314,8 +327,10 @@ 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, datastore, datastore_boundary = load_config_and_datastore( + config_path=args.config_path + ) # Create datamodule data_module = WeatherDataModule( @@ -325,6 +340,9 @@ def main(input_args=None): standardize=True, num_past_forcing_steps=args.num_past_forcing_steps, num_future_forcing_steps=args.num_future_forcing_steps, + num_past_boundary_steps=args.num_past_boundary_steps, + num_future_boundary_steps=args.num_future_boundary_steps, + datastore_boundary=datastore_boundary, load_single_member=args.load_single_member, batch_size=args.batch_size, num_workers=args.num_workers, diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 91a371989..8a0bacacd 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -17,7 +17,8 @@ class WeatherDataset(torch.utils.data.Dataset): """Dataset class for weather data. - This class loads and processes weather data from a given datastore. + This class loads and processes weather data from a given datastore, + with optional boundary forcing from a separate boundary datastore. Parameters ---------- @@ -37,6 +38,15 @@ class WeatherDataset(torch.utils.data.Dataset): forcing from times t, t+1, ..., t+j-1, t+j (and potentially times before t, given num_past_forcing_steps) are included as forcing inputs at time t. Default is 1. + num_past_boundary_steps: int, optional + Number of past time steps to include in boundary forcing input. + Default is 1. + num_future_boundary_steps: int, optional + Number of future time steps to include in boundary forcing input. + Default is 1. + datastore_boundary : BaseDatastore, optional + A separate datastore providing boundary forcing data. If None, no + boundary forcing is used (boundary tensor will be empty). load_single_member : bool, optional If `False` and the datastore returns an ensemble of state realisations, treat each state ensemble member as an independent @@ -53,6 +63,9 @@ def __init__( ar_steps: int = 3, num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, + num_past_boundary_steps: int = 1, + num_future_boundary_steps: int = 1, + datastore_boundary: Union[BaseDatastore, None] = None, load_single_member: bool = False, standardize: bool = True, ): @@ -61,8 +74,11 @@ def __init__( self.split = split self.ar_steps = ar_steps self.datastore = datastore + self.datastore_boundary = datastore_boundary self.num_past_forcing_steps = num_past_forcing_steps self.num_future_forcing_steps = num_future_forcing_steps + self.num_past_boundary_steps = num_past_boundary_steps + self.num_future_boundary_steps = num_future_boundary_steps self.load_single_member = load_single_member self.da_state = self.datastore.get_dataarray( @@ -76,6 +92,14 @@ def __init__( "The datastore must provide state data for the WeatherDataset." ) + # Load boundary forcing from the boundary datastore + if self.datastore_boundary is not None: + self.da_boundary_forcing = self.datastore_boundary.get_dataarray( + category="forcing", split=self.split + ) + else: + self.da_boundary_forcing = None + if self.datastore.is_ensemble and self.load_single_member: warnings.warn( "only using first ensemble member, so dataset size is " @@ -138,6 +162,18 @@ def __init__( self.da_forcing_mean = None self.da_forcing_std = None + if self.datastore_boundary is not None: + self.ds_boundary_stats = ( + self.datastore_boundary.get_standardization_dataarray( + category="forcing" + ) + ) + self.da_boundary_mean = self.ds_boundary_stats.forcing_mean + self.da_boundary_std = self.ds_boundary_stats.forcing_std + else: + self.da_boundary_mean = None + self.da_boundary_std = None + self.state_std_safe = self._compute_std_safe( self.da_state_std, "state" ) @@ -149,6 +185,13 @@ def __init__( else: self.forcing_std_safe = None + if self.da_boundary_std is not None: + self.boundary_std_safe = self._compute_std_safe( + self.da_boundary_std, "boundary_forcing" + ) + else: + self.boundary_std_safe = None + def _compute_std_safe(self, std: xr.DataArray, feature: str): eps = np.finfo(std.dtype).eps if bool((std <= eps).any()): @@ -261,7 +304,14 @@ def _slice_state_time(self, da_state, idx, n_steps: int): da_sliced = da_state.isel(time=slice(start_idx, end_idx)) return da_sliced - def _slice_forcing_time(self, da_forcing, idx, n_steps: int): + def _slice_forcing_time( + self, + da_forcing, + idx, + n_steps: int, + num_past_steps: Union[int, None] = None, + num_future_steps: Union[int, None] = None, + ): """ Produce a time slice of the given dataarray `da_forcing` (forcing) starting at `idx` and with `n_steps` steps. An `offset` is calculated @@ -296,6 +346,12 @@ def _slice_forcing_time(self, da_forcing, idx, n_steps: int): init_steps = 2 da_list = [] + # Allow overriding the window size (used for boundary forcing) + if num_past_steps is None: + num_past_steps = self.num_past_forcing_steps + if num_future_steps is None: + num_future_steps = self.num_future_forcing_steps + if self.datastore.is_forecast: # This implies that the data will have both `analysis_time` and # `elapsed_forecast_duration` dimensions for forecasts. We for now @@ -303,10 +359,10 @@ def _slice_forcing_time(self, da_forcing, idx, n_steps: int): # times (given no offset). Note that this means that we get one # sample per forecast. # Add a 'time' dimension using the actual forecast times - offset = max(init_steps, self.num_past_forcing_steps) + offset = max(init_steps, num_past_steps) for step in range(n_steps): - start_idx = offset + step - self.num_past_forcing_steps - end_idx = offset + step + self.num_future_forcing_steps + start_idx = offset + step - num_past_steps + end_idx = offset + step + num_future_steps current_time = ( da_forcing.analysis_time[idx] @@ -340,10 +396,10 @@ def _slice_forcing_time(self, da_forcing, idx, n_steps: int): # For analysis data, we slice the time dimension directly. The # offset is only relevant for the very first (and last) samples in # the dataset. - offset = idx + max(init_steps, self.num_past_forcing_steps) + offset = idx + max(init_steps, num_past_steps) for step in range(n_steps): - start_idx = offset + step - self.num_past_forcing_steps - end_idx = offset + step + self.num_future_forcing_steps + start_idx = offset + step - num_past_steps + end_idx = offset + step + num_future_steps # Slice the data over the desired time window da_sliced = da_forcing.isel(time=slice(start_idx, end_idx + 1)) @@ -371,8 +427,8 @@ def _slice_forcing_time(self, da_forcing, idx, n_steps: int): def _build_item_dataarrays(self, idx): """ - Create the dataarrays for the initial states, target states and forcing - data for the sample at index `idx`. + Create the dataarrays for the initial states, target states, forcing + and boundary data for the sample at index `idx`. Parameters ---------- @@ -387,6 +443,9 @@ def _build_item_dataarrays(self, idx): The dataarray for the target states. da_forcing_windowed : xr.DataArray The dataarray for the forcing data, windowed for the sample. + da_boundary_windowed : xr.DataArray + The dataarray for the boundary forcing data, windowed for the + sample. da_target_times : xr.DataArray The dataarray for the target times. """ @@ -421,10 +480,24 @@ def _build_item_dataarrays(self, idx): da_forcing=da_forcing, idx=sample_idx, n_steps=self.ar_steps ) + # Slice boundary forcing if available + if self.da_boundary_forcing is not None: + da_boundary_windowed = self._slice_forcing_time( + da_forcing=self.da_boundary_forcing, + idx=sample_idx, + n_steps=self.ar_steps, + num_past_steps=self.num_past_boundary_steps, + num_future_steps=self.num_future_boundary_steps, + ) + else: + da_boundary_windowed = None + # load the data into memory da_state.load() if da_forcing is not None: da_forcing_windowed.load() + if da_boundary_windowed is not None: + da_boundary_windowed.load() da_init_states = da_state.isel(time=slice(0, 2)) da_target_states = da_state.isel(time=slice(2, None)) @@ -447,6 +520,11 @@ def _build_item_dataarrays(self, idx): da_forcing_windowed - self.da_forcing_mean ) / self.forcing_std_safe + if da_boundary_windowed is not None: + da_boundary_windowed = ( + da_boundary_windowed - self.da_boundary_mean + ) / self.boundary_std_safe + if da_forcing is not None: # stack the `forcing_feature` and `window_sample` dimensions into a # single `forcing_feature` dimension @@ -467,10 +545,43 @@ def _build_item_dataarrays(self, idx): }, ) + if da_boundary_windowed is not None: + da_boundary_windowed = da_boundary_windowed.stack( + forcing_feature_windowed=("forcing_feature", "window") + ) + else: + # create an empty boundary tensor with the right shape + # Use the boundary datastore's grid_index if available, otherwise + # fall back to state grid_index (for the no-boundary case the + # last dim is 0 anyway) + if self.datastore_boundary is not None: + da_boundary_ref = self.datastore_boundary.get_dataarray( + category="forcing", split=self.split + ) + boundary_grid_index = ( + da_boundary_ref.grid_index + if da_boundary_ref is not None + else da_state.grid_index + ) + else: + boundary_grid_index = da_state.grid_index + da_boundary_windowed = xr.DataArray( + data=np.empty( + (self.ar_steps, boundary_grid_index.size, 0), + ), + dims=("time", "grid_index", "forcing_feature"), + coords={ + "time": da_target_times, + "grid_index": boundary_grid_index, + "forcing_feature": [], + }, + ) + return ( da_init_states, da_target_states, da_forcing_windowed, + da_boundary_windowed, da_target_times, ) @@ -505,6 +616,7 @@ def __getitem__(self, idx): da_init_states, da_target_states, da_forcing_windowed, + da_boundary_windowed, da_target_times, ) = self._build_item_dataarrays(idx=idx) @@ -521,13 +633,15 @@ def __getitem__(self, idx): ) forcing = torch.tensor(da_forcing_windowed.values, dtype=tensor_dtype) + boundary = torch.tensor(da_boundary_windowed.values, dtype=tensor_dtype) # init_states: (2, N_grid, d_features) # target_states: (ar_steps, N_grid, d_features) # forcing: (ar_steps, N_grid, d_windowed_forcing) + # boundary: (ar_steps, N_boundary_grid, d_windowed_boundary) # target_times: (ar_steps,) - return init_states, target_states, forcing, target_times + return init_states, target_states, forcing, boundary, target_times def __iter__(self): """ @@ -645,6 +759,9 @@ def __init__( standardize: bool = True, num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, + num_past_boundary_steps: int = 1, + num_future_boundary_steps: int = 1, + datastore_boundary: Union[BaseDatastore, None] = None, load_single_member: bool = False, batch_size: int = 4, num_workers: int = 16, @@ -652,8 +769,11 @@ def __init__( ): super().__init__() self._datastore = datastore + self._datastore_boundary = datastore_boundary self.num_past_forcing_steps = num_past_forcing_steps self.num_future_forcing_steps = num_future_forcing_steps + self.num_past_boundary_steps = num_past_boundary_steps + self.num_future_boundary_steps = num_future_boundary_steps self.ar_steps_train = ar_steps_train self.ar_steps_eval = ar_steps_eval self.standardize = standardize @@ -671,24 +791,27 @@ def __init__( self.multiprocessing_context = "spawn" def setup(self, stage=None): + shared_kwargs = dict( + num_past_forcing_steps=self.num_past_forcing_steps, + num_future_forcing_steps=self.num_future_forcing_steps, + num_past_boundary_steps=self.num_past_boundary_steps, + num_future_boundary_steps=self.num_future_boundary_steps, + datastore_boundary=self._datastore_boundary, + load_single_member=self.load_single_member, + standardize=self.standardize, + ) if stage == "fit" or stage is None: self.train_dataset = WeatherDataset( datastore=self._datastore, split="train", ar_steps=self.ar_steps_train, - standardize=self.standardize, - 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, - standardize=self.standardize, - 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: @@ -696,10 +819,7 @@ def setup(self, stage=None): datastore=self._datastore, split=self.eval_split, ar_steps=self.ar_steps_eval, - standardize=self.standardize, - 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): diff --git a/tests/dummy_datastore.py b/tests/dummy_datastore.py index 3a844d6d9..a2c80b975 100644 --- a/tests/dummy_datastore.py +++ b/tests/dummy_datastore.py @@ -473,6 +473,25 @@ def grid_shape_state(self) -> CartesianGridShape: return CartesianGridShape(x=n_points_1d, y=n_points_1d) +class BoundaryDummyDatastore(DummyDatastore): + """DummyDatastore acting as a boundary forcing provider. + + Uses a different (smaller) grid and different forcing features to simulate + a boundary datastore from a separate domain. Only the ``forcing`` + category is meaningful; ``state`` and ``static`` are inherited but unused. + """ + + SHORT_NAME = "dummydata_boundary" + N_FEATURES = dict(state=5, forcing=3, static=1) + + def __init__(self, n_grid_points=400, n_timesteps=10, step_length=None): + super().__init__( + n_grid_points=n_grid_points, + n_timesteps=n_timesteps, + step_length=step_length, + ) + + class EnsembleDummyDatastore(BaseDatastore): """Small offline datastore for ensemble WeatherDataset tests. diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 3ca9b31b3..ee73cbb3d 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -15,7 +15,11 @@ from neural_lam.models import ForecasterModule from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example -from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore +from tests.dummy_datastore import ( + BoundaryDummyDatastore, + DummyDatastore, + EnsembleDummyDatastore, +) @pytest.mark.parametrize("datastore_name", DATASTORES.keys()) @@ -48,7 +52,7 @@ def test_dataset_item_shapes(datastore_name): # unpack the item, this is the current return signature for # WeatherDataset.__getitem__ - init_states, target_states, forcing, target_times = item + init_states, target_states, forcing, boundary, target_times = item # initial states assert init_states.ndim == 3 @@ -99,9 +103,9 @@ def test_dataset_item_create_dataarray_from_tensor(datastore_name): # unpack the item, this is the current return signature for # WeatherDataset.__getitem__ - _, target_states, _, target_times_arr = dataset[idx] - _, da_target_true, _, da_target_times_true = dataset._build_item_dataarrays( - idx=idx + _, target_states, _, _, target_times_arr = dataset[idx] + _, da_target_true, _, _, da_target_times_true = ( + dataset._build_item_dataarrays(idx=idx) ) target_times = np.array(target_times_arr, dtype="datetime64[ns]") @@ -374,8 +378,8 @@ def test_ensemble_index_mapping_is_time_major(): standardize=False, ) - init_states_0, _, _, target_times_0 = dataset[0] - init_states_1, _, _, target_times_1 = dataset[1] + init_states_0, _, _, _, target_times_0 = dataset[0] + init_states_1, _, _, _, target_times_1 = dataset[1] # Adjacent flat indices correspond to same sample_idx and different member. assert torch.equal(target_times_0, target_times_1) @@ -399,8 +403,8 @@ def test_ensemble_forcing_uses_same_member_when_available(): standardize=False, ) - _, _, forcing_0, target_times_0 = dataset[0] - _, _, forcing_1, target_times_1 = dataset[1] + _, _, forcing_0, _, target_times_0 = dataset[0] + _, _, forcing_1, _, target_times_1 = dataset[1] assert torch.equal(target_times_0, target_times_1) assert not torch.equal(forcing_0, forcing_1) @@ -423,8 +427,8 @@ def test_ensemble_forcing_without_member_dim_is_shared(): standardize=False, ) - init_states_0, _, forcing_0, target_times_0 = dataset[0] - init_states_1, _, forcing_1, target_times_1 = dataset[1] + init_states_0, _, forcing_0, _, target_times_0 = dataset[0] + init_states_1, _, forcing_1, _, target_times_1 = dataset[1] assert torch.equal(target_times_0, target_times_1) assert not torch.equal(init_states_0, init_states_1) @@ -527,7 +531,132 @@ def get_dataarray(self, category, split, **kwargs): assert dataset.da_forcing_std is None # Ensure we can still retrieve a sample (forcing tensor should be empty) - init_states, target_states, forcing, target_times = dataset[0] + init_states, target_states, forcing, boundary, target_times = dataset[0] assert ( forcing.shape[-1] == 0 ), "Expected zero forcing features when forcing is None" + + +def test_boundary_datastore_shapes(): + """WeatherDataset with a boundary datastore should return a 5-tuple where + the boundary tensor has the boundary grid and windowed features.""" + n_timesteps = 20 + ar_steps = 3 + num_past_boundary = 1 + num_future_boundary = 1 + boundary_window = num_past_boundary + num_future_boundary + 1 + + datastore = DummyDatastore(n_grid_points=100, n_timesteps=n_timesteps) + boundary_ds = BoundaryDummyDatastore( + n_grid_points=25, n_timesteps=n_timesteps + ) + + dataset = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=ar_steps, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + num_past_boundary_steps=num_past_boundary, + num_future_boundary_steps=num_future_boundary, + datastore_boundary=boundary_ds, + standardize=True, + ) + + init_states, target_states, forcing, boundary, target_times = dataset[0] + + assert init_states.shape == (2, 100, datastore.N_FEATURES["state"]) + assert target_states.shape == (ar_steps, 100, datastore.N_FEATURES["state"]) + assert forcing.ndim == 3 + assert boundary.ndim == 3 + assert boundary.shape[0] == ar_steps + assert boundary.shape[1] == 25 # boundary grid + n_boundary_features = boundary_ds.N_FEATURES["forcing"] + assert boundary.shape[2] == n_boundary_features * boundary_window + assert target_times.shape == (ar_steps,) + + # Verify no NaN from standardization + assert not torch.isnan(boundary).any() + + +def test_boundary_datastore_none_gives_empty_boundary(): + """Without a boundary datastore the boundary tensor should have zero + features (last dim == 0).""" + datastore = DummyDatastore(n_grid_points=100, n_timesteps=20) + + dataset = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=3, + standardize=False, + ) + + _, _, _, boundary, _ = dataset[0] + assert boundary.shape[-1] == 0 + + +def test_boundary_datastore_standardization(): + """Boundary forcing should be standardized when standardize=True and + left raw when standardize=False.""" + n_timesteps = 20 + datastore = DummyDatastore(n_grid_points=100, n_timesteps=n_timesteps) + boundary_ds = BoundaryDummyDatastore( + n_grid_points=25, n_timesteps=n_timesteps + ) + + dataset_std = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=2, + num_past_boundary_steps=0, + num_future_boundary_steps=0, + datastore_boundary=boundary_ds, + standardize=True, + ) + + dataset_raw = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=2, + num_past_boundary_steps=0, + num_future_boundary_steps=0, + datastore_boundary=boundary_ds, + standardize=False, + ) + + _, _, _, boundary_std, _ = dataset_std[0] + _, _, _, boundary_raw, _ = dataset_raw[0] + + # Standardized and raw should differ (unless data happens to have + # mean=0 and std=1, which is essentially impossible for random data) + assert not torch.equal(boundary_std, boundary_raw) + + +def test_boundary_dataset_length_unchanged(): + """Adding a boundary datastore should not change the dataset length.""" + n_timesteps = 20 + datastore = DummyDatastore(n_grid_points=100, n_timesteps=n_timesteps) + boundary_ds = BoundaryDummyDatastore( + n_grid_points=25, n_timesteps=n_timesteps + ) + + dataset_no_boundary = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=3, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + ) + + dataset_with_boundary = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=3, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + datastore_boundary=boundary_ds, + ) + + assert len(dataset_no_boundary) == len(dataset_with_boundary) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 962eddcc6..9805d6653 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -316,7 +316,7 @@ def test_plot_examples_integration_saves_figure( ), f"Expected time_step_unit={time_unit}, got {model.time_step_unit}" # Generate prediction - (init_states, target, forcing_features, _batch_times) = batch + (init_states, target, forcing_features, _boundary, _batch_times) = batch prediction, _ = model.forecaster(init_states, forcing_features, target) # Rescale to original data scale @@ -338,7 +338,7 @@ def test_plot_examples_integration_saves_figure( # Get first example pred_slice = prediction_rescaled[0].detach() # Detach from graph target_slice = target_rescaled[0].detach() - time_slice = batch[3][0] + time_slice = batch[4][0] # Create DataArrays dataset = WeatherDataset(datastore=datastore, split="train") diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index 91aa01578..2ff70bc68 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -116,7 +116,7 @@ def test_time_slicing_analysis( sample = dataset[0] - init_states, target_states, forcing, _ = [ + init_states, target_states, forcing, _boundary, _ = [ tensor.numpy() for tensor in sample ] @@ -192,4 +192,6 @@ def test_step_length_timedeltas(step_length): # Test that we can get a sample sample = dataset[0] - assert len(sample) == 4 # init_states, target_states, forcing, target_times + assert ( + len(sample) == 5 + ) # init_states, target_states, forcing, boundary, target_times diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a9..6e8e95f88 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -70,7 +70,7 @@ def capture_init(_self, **kwargs): ), patch( "neural_lam.train_model.load_config_and_datastore", - return_value=(MagicMock(), MagicMock()), + return_value=(MagicMock(), MagicMock(), None), ), patch("neural_lam.train_model.WeatherDataModule"), patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), From 3dc0045bf1e615aba5b609cfe22768b359e258a2 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 16:02:00 +0200 Subject: [PATCH 02/35] Update CHANGELOG PR link to #635 Co-Authored-By: Claude Opus 4.6 --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f961ae633..8262bbc31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. [\#TODO](https://github.com/mllam/neural-lam/pull/TODO) @sadamov +- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov ### Changed From e945a4de804163fd0ec3ab095acd680efc4e78b5 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 17:17:03 +0200 Subject: [PATCH 03/35] feat: add ERA5 boundary datastore test configs and conftest fixture Add MDP-based ERA5 boundary example at tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/ with config.yaml, era5.datastore.yaml (WeatherBench2 64x32 equiangular), and danra.datastore.yaml (DANRA 100m winds interior). Add DATASTORES_BOUNDARY_EXAMPLES dict and init_datastore_boundary_example() to conftest.py for use in boundary integration tests. Co-Authored-By: Claude Opus 4.6 --- tests/conftest.py | 17 +++ .../era5_1000hPa_danra_100m_winds/.gitignore | 2 + .../era5_1000hPa_danra_100m_winds/config.yaml | 12 ++ .../danra.datastore.yaml | 99 ++++++++++++++++ .../era5.datastore.yaml | 106 ++++++++++++++++++ 5 files changed, 236 insertions(+) create mode 100644 tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/.gitignore create mode 100644 tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml create mode 100644 tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml create mode 100644 tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml diff --git a/tests/conftest.py b/tests/conftest.py index 47237ed55..bfae0cf62 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -105,6 +105,15 @@ def download_meps_example_reduced_dataset(): dummydata=None, ) +DATASTORES_BOUNDARY_EXAMPLES = { + "mdp": ( + DATASTORE_EXAMPLES_ROOT_PATH + / "mdp" + / "era5_1000hPa_danra_100m_winds" + / "era5.datastore.yaml" + ), +} + DATASTORES[DummyDatastore.SHORT_NAME] = DummyDatastore @@ -123,3 +132,11 @@ def init_datastore_example(datastore_kind): ) return datastore + + +def init_datastore_boundary_example(datastore_kind): + datastore_boundary = init_datastore( + datastore_kind=datastore_kind, + config_path=DATASTORES_BOUNDARY_EXAMPLES[datastore_kind], + ) + return datastore_boundary diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/.gitignore b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/.gitignore new file mode 100644 index 000000000..f2828f46c --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/.gitignore @@ -0,0 +1,2 @@ +*.zarr/ +graph/ diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml new file mode 100644 index 000000000..a158bee3b --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml @@ -0,0 +1,12 @@ +datastore: + kind: mdp + config_path: danra.datastore.yaml +datastore_boundary: + kind: mdp + config_path: era5.datastore.yaml +training: + state_feature_weighting: + __config_class__: ManualStateFeatureWeighting + weights: + u100m: 1.0 + v100m: 1.0 diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml new file mode 100644 index 000000000..3edf12673 --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml @@ -0,0 +1,99 @@ +schema_version: v0.5.0 +dataset_version: v0.1.0 + +output: + variables: + static: [grid_index, static_feature] + state: [time, grid_index, state_feature] + forcing: [time, grid_index, forcing_feature] + coord_ranges: + time: + start: 1990-09-03T00:00 + end: 1990-09-09T00:00 + step: PT3H + chunking: + time: 1 + splitting: + dim: time + splits: + train: + start: 1990-09-03T00:00 + end: 1990-09-06T00:00 + compute_statistics: + ops: [mean, std, diff_mean, diff_std] + dims: [grid_index, time] + val: + start: 1990-09-06T00:00 + end: 1990-09-07T00:00 + test: + start: 1990-09-07T00:00 + end: 1990-09-09T00:00 + +inputs: + danra_height_levels: + path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/height_levels.zarr + dims: [time, x, y, altitude] + variables: + u: + altitude: + values: [100,] + units: m + v: + altitude: + values: [100, ] + units: m + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + dims: [altitude] + name_format: "{var_name}{altitude}m" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state + + danra_surface: + path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr + dims: [time, x, y] + variables: + # use surface incoming shortwave radiation as forcing + - swavr0m + dim_mapping: + time: + method: rename + dim: time + grid_index: + method: stack + dims: [x, y] + forcing_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + target_output_variable: forcing + + danra_lsm: + path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/lsm.zarr + dims: [x, y] + variables: + - lsm + dim_mapping: + grid_index: + method: stack + dims: [x, y] + static_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + target_output_variable: static + +extra: + projection: + class_name: LambertConformal + kwargs: + central_longitude: 25.0 + central_latitude: 56.7 + standard_parallels: [56.7, 56.7] + globe: + semimajor_axis: 6367470.0 + semiminor_axis: 6367470.0 diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml new file mode 100644 index 000000000..c83489c62 --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml @@ -0,0 +1,106 @@ +schema_version: v0.5.0 +dataset_version: v1.0.0 + +output: + variables: + static: [grid_index, static_feature] + forcing: [time, grid_index, forcing_feature] + coord_ranges: + time: + start: 1990-09-01T00:00 + end: 2022-09-30T00:00 + step: PT6H + chunking: + time: 1 + splitting: + dim: time + splits: + train: + start: 1990-09-01T00:00 + end: 2022-09-30T00:00 + compute_statistics: + ops: [mean, std, diff_mean, diff_std] + dims: [grid_index, time] + val: + start: 1990-09-01T00:00 + end: 2022-09-30T00:00 + test: + start: 1990-09-01T00:00 + end: 2022-09-30T00:00 + +inputs: + era_height_levels: + path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' + dims: [time, longitude, latitude, level] + variables: + u_component_of_wind: + level: + values: [1000,] + units: hPa + dim_mapping: + time: + method: rename + dim: time + x: + method: rename + dim: longitude + y: + method: rename + dim: latitude + forcing_feature: + method: stack_variables_by_var_name + dims: [level] + name_format: "{var_name}{level}hPa" + grid_index: + method: stack + dims: [x, y] + target_output_variable: forcing + + era5_surface: + path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' + dims: [time, longitude, latitude, level] + variables: + - mean_sea_level_pressure + dim_mapping: + time: + method: rename + dim: time + x: + method: rename + dim: longitude + y: + method: rename + dim: latitude + forcing_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [x, y] + target_output_variable: forcing + + era5_static: + path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' + dims: [time, longitude, latitude, level] + variables: + - land_sea_mask + dim_mapping: + x: + method: rename + dim: longitude + y: + method: rename + dim: latitude + static_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [x, y] + target_output_variable: static + +extra: + projection: + class_name: PlateCarree + kwargs: + central_longitude: 0.0 From 10e18777f546ee10afdfd22b5482127cb0ed3fa7 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 17:35:20 +0200 Subject: [PATCH 04/35] fix: handle boundary-only MDP datastores without state variables MDPDatastore.__init__ crashed with KeyError when loading a datastore that has only forcing+static (no state), e.g. ERA5 boundary data. Fix is_ensemble check to guard against missing state, and grid_shape_state to fall back to forcing/static categories. Co-Authored-By: Claude Opus 4.6 --- neural_lam/datastore/mdp.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 114550e6d..ad97c6f87 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -75,7 +75,9 @@ def __init__(self, config_path, n_boundary_points=30, reuse_existing=True): self._ds = mdp.create_dataset(config=self._config) self._ds.to_zarr(fp_ds) self._n_boundary_points = n_boundary_points - self.is_ensemble = "ensemble_member" in self._ds["state"].dims + self.is_ensemble = ( + "state" in self._ds and "ensemble_member" in self._ds["state"].dims + ) self.has_ensemble_forcing = ( "forcing" in self._ds and "ensemble_member" in self._ds["forcing"].dims @@ -453,9 +455,16 @@ def grid_shape_state(self): The shape of the cartesian grid for the state variables. """ - ds_state = self.unstack_grid_coords(self._ds["state"]) + # Use state if available, otherwise fall back to forcing (e.g. + # boundary-only datastores that have no state variables). + for category in ("state", "forcing", "static"): + if category in self._ds: + ds_cat = self.unstack_grid_coords(self._ds[category]) + break + else: + raise ValueError("Dataset has no state, forcing, or static data") xdim, ydim = self.spatial_coordinates - da_x, da_y = ds_state[xdim], ds_state[ydim] + da_x, da_y = ds_cat[xdim], ds_cat[ydim] assert da_x.ndim == da_y.ndim == 1 return CartesianGridShape(x=da_x.size, y=da_y.size) From affb44199e97a32599ea5177dd568ae5381f0972 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 17:42:59 +0200 Subject: [PATCH 05/35] docs: mention MDP boundary-only datastore support in CHANGELOG Co-Authored-By: Claude Opus 4.6 --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8262bbc31..eda128ef1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` now supports boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov ### Changed From c3af69eb70dfa79355d99d2eda5886eada7562e2 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 17:53:54 +0200 Subject: [PATCH 06/35] fix: handle boundary-only NpyFilesDatastoreMEPS without state variables Make _get_analysis_times fall back to forcing file patterns when no state files exist, guard get_dataarray("state") against empty var_names, and prevent empty feature list from matching state loading path. Co-Authored-By: Claude Opus 4.6 --- CHANGELOG.md | 2 +- neural_lam/datastore/npyfilesmeps/store.py | 44 ++++++++++++++-------- 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eda128ef1..c68954398 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` now supports boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads and standardizes boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov ### Changed diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py index f65fd8aef..d28a3941d 100644 --- a/neural_lam/datastore/npyfilesmeps/store.py +++ b/neural_lam/datastore/npyfilesmeps/store.py @@ -234,11 +234,17 @@ def get_dataarray( """ if category == "state": + state_vars = self.get_vars_names(category="state") + if not state_vars: + raise ValueError( + "No state variables configured. This datastore may " + "be a boundary-only datastore without state data." + ) das = [] # for the state category, we need to load all ensemble members for member in range(self._num_ensemble_members): da_member = self._get_single_timeseries_dataarray( - features=self.get_vars_names(category="state"), + features=state_vars, split=split, member=member, ) @@ -382,7 +388,8 @@ def _get_single_timeseries_dataarray( features_vary_with_analysis_time = True feature_dim_mask = None if ( - features == self.get_vars_names(category="state") + features + and features == self.get_vars_names(category="state") and split is not None ): filename_format = STATE_FILENAME_FORMAT @@ -533,22 +540,29 @@ def _get_analysis_times(self, split) -> List[np.datetime64]: The analysis times for the given split, sorted in ascending order. """ - pattern = re.sub(r"{analysis_time:[^}]*}", "*", STATE_FILENAME_FORMAT) - pattern = re.sub(r"{member_id:[^}]*}", "*", pattern) - sample_dir = self.root_path / "samples" / split - sample_files = sample_dir.glob(pattern) - times = [] - for fp in sample_files: - name_parts = parse.parse(STATE_FILENAME_FORMAT, fp.name) - times.append(name_parts["analysis_time"]) - if len(times) == 0: - raise ValueError( - f"No files found in {sample_dir} with pattern {pattern}" - ) + # Try state files first, then fall back to forcing files + # (boundary-only datastores may not have state files) + formats_to_try = [ + (STATE_FILENAME_FORMAT, [r"{member_id:[^}]*}"]), + (TOA_SW_DOWN_FLUX_FILENAME_FORMAT, []), + ] + + for filename_format, extra_wildcards in formats_to_try: + pattern = re.sub(r"{analysis_time:[^}]*}", "*", filename_format) + for wc in extra_wildcards: + pattern = re.sub(wc, "*", pattern) + + sample_files = list(sample_dir.glob(pattern)) + if sample_files: + times = [] + for fp in sample_files: + name_parts = parse.parse(filename_format, fp.name) + times.append(name_parts["analysis_time"]) + return sorted(times) - return sorted(times) + raise ValueError(f"No state or forcing files found in {sample_dir}") def _calc_datetime_forcing_features(self, da_time: xr.DataArray): da_hour_angle = da_time.dt.hour / 12 * np.pi From e5d8061d03368bf1923e544a18ccead54b8bbc80 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 18:06:13 +0200 Subject: [PATCH 07/35] fix: MDP get_vars_names/get_vars_long_names for any missing category Guard against missing state/static feature keys in the zarr, not just forcing. Boundary-only datastores (e.g. ERA5) may lack state_feature entirely. Co-Authored-By: Claude Opus 4.6 --- neural_lam/datastore/mdp.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index ad97c6f87..8e69f1f34 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -193,10 +193,12 @@ def get_vars_names(self, category: str) -> List[str]: The names of the variables in the given category. """ - if category not in self._ds and category == "forcing": - warnings.warn("no forcing data found in datastore") + feature_key = f"{category}_feature" + if feature_key not in self._ds: + if category == "forcing": + warnings.warn("no forcing data found in datastore") return [] - return self._ds[f"{category}_feature"].values.tolist() + return self._ds[feature_key].values.tolist() def get_vars_long_names(self, category: str) -> List[str]: """ @@ -213,10 +215,12 @@ def get_vars_long_names(self, category: str) -> List[str]: The long names of the variables in the given category. """ - if category not in self._ds and category == "forcing": - warnings.warn("no forcing data found in datastore") + long_name_key = f"{category}_feature_long_name" + if long_name_key not in self._ds: + if category == "forcing": + warnings.warn("no forcing data found in datastore") return [] - return self._ds[f"{category}_feature_long_name"].values.tolist() + return self._ds[long_name_key].values.tolist() def get_num_data_vars(self, category: str) -> int: """Return the number of variables in the given category. From 8d8bb77eb61a397380cc8b3fb19c73ed9e3048ac Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 11 May 2026 18:41:49 +0200 Subject: [PATCH 08/35] fix: add num_grid_points override to MDPDatastore Return len(grid_index) directly instead of computing from grid_shape_state, which is more robust for boundary-only datastores. Co-Authored-By: Claude Opus 4.6 --- neural_lam/datastore/mdp.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 8e69f1f34..126e89dc3 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -556,3 +556,7 @@ def get_lat_lon(self, category: str) -> np.ndarray: coords = np.stack((lon.values, lat.values), axis=1) return coords + + @property + def num_grid_points(self) -> int: + return len(self._ds.grid_index) From 4c94adc7fcd2c311cdbfd72d59932db7c96810b7 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 14 May 2026 14:47:47 +0200 Subject: [PATCH 09/35] test: tighten ERA5 boundary test fixture and skip slow by default - Shrink the ERA5 boundary test dataset to 2022-03-30..2022-04-12 (was 1990-2022), add per-input lat/lon coord_ranges and enable mllam's convex-hull domain_cropping with include_interior_points=true. Stats computation drops from minutes to seconds and the cached zarr stays under 1 MB. - Stack [longitude, latitude] directly into grid_index in the era5 dim_mapping (per mllam's example.era5_cropped.yaml) so the original coord names survive for the convex-hull crop -- removes the need for any rename-preserve workaround in neural-lam. - Generalise the MDPDatastore units loop over self.spatial_coordinates with sensible defaults for x/y (m) and longitude/latitude/lon/lat (degrees_*) so ERA5-style geographic datastores work. - Register a pytest `slow` marker and a `--run-slow` CLI flag so the ERA5 boundary integration test (added in this PR) is skipped by default and can be run on demand. Co-Authored-By: Claude Opus 4.7 --- neural_lam/datastore/mdp.py | 19 +++++- pyproject.toml | 5 ++ tests/conftest.py | 17 ++++++ .../era5.datastore.yaml | 60 ++++++++++--------- 4 files changed, 69 insertions(+), 32 deletions(-) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 126e89dc3..53c200c2f 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -288,10 +288,23 @@ def get_dataarray( da_category = self._ds[category] - # set units on x y coordinates if missing - for coord in ["x", "y"]: + # Set units on spatial coordinates if missing. Use the dim names + # actually declared in the config's grid_index stacking (so this + # works for both projected (x, y) and geographic (longitude, + # latitude) source datasets like ERA5). + _UNITS_BY_COORD = { + "x": "m", + "y": "m", + "longitude": "degrees_east", + "latitude": "degrees_north", + "lon": "degrees_east", + "lat": "degrees_north", + } + for coord in self.spatial_coordinates: if "units" not in da_category[coord].attrs: - da_category[coord].attrs["units"] = "m" + da_category[coord].attrs["units"] = _UNITS_BY_COORD.get( + coord, "" + ) # set multi-index for grid-index da_category = da_category.set_index(grid_index=self.spatial_coordinates) diff --git a/pyproject.toml b/pyproject.toml index bccec0828..77a49d792 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 --run-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 bfae0cf62..876acc025 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,6 +23,23 @@ os.environ["WANDB_MODE"] = "disabled" +def pytest_addoption(parser): + parser.addoption( + "--run-slow", + action="store_true", + default=False, + help="Run tests marked as slow", + ) + + +def pytest_collection_modifyitems(config, items): + if not config.getoption("--run-slow"): + skip_slow = pytest.mark.skip(reason="need --run-slow option to run") + for item in items: + if "slow" in item.keywords: + item.add_marker(skip_slow) + + @pytest.fixture(autouse=True) def ensure_rank_zero(monkeypatch): """Ensure rank_zero_only.rank == 0 so @rank_zero_only-decorated functions diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml index c83489c62..9a03970bb 100644 --- a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml @@ -7,8 +7,8 @@ output: forcing: [time, grid_index, forcing_feature] coord_ranges: time: - start: 1990-09-01T00:00 - end: 2022-09-30T00:00 + start: 2022-03-30T00:00 + end: 2022-04-12T00:00 step: PT6H chunking: time: 1 @@ -16,22 +16,36 @@ output: dim: time splits: train: - start: 1990-09-01T00:00 - end: 2022-09-30T00:00 + start: 2022-03-30T00:00 + end: 2022-04-12T00:00 compute_statistics: ops: [mean, std, diff_mean, diff_std] dims: [grid_index, time] val: - start: 1990-09-01T00:00 - end: 2022-09-30T00:00 + start: 2022-03-30T00:00 + end: 2022-04-12T00:00 test: - start: 1990-09-01T00:00 - end: 2022-09-30T00:00 + start: 2022-03-30T00:00 + end: 2022-04-12T00:00 + # Second-pass crop: keep only ERA5 points within `margin_width_degrees` + # of the interior dataset's convex hull. Follows the documented mllam + # pattern in `example.era5_cropped.yaml` -- the dim_mapping below stacks + # `[longitude, latitude]` directly into `grid_index` (no rename) so that + # convex-hull cropping can find lat/lon coords by name. + domain_cropping: + margin_width_degrees: 10 + interior_dataset_config_path: danra.datastore.yaml + include_interior_points: true inputs: era_height_levels: path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' dims: [time, longitude, latitude, level] + # First-pass crop on the source lat/lon dims to keep the test dataset + # small. Region covers DANRA + a ~10 degree margin. + coord_ranges: + latitude: {start: 40, end: 75} + longitude: {start: 0, end: 30} variables: u_component_of_wind: level: @@ -41,62 +55,50 @@ inputs: time: method: rename dim: time - x: - method: rename - dim: longitude - y: - method: rename - dim: latitude forcing_feature: method: stack_variables_by_var_name dims: [level] name_format: "{var_name}{level}hPa" grid_index: method: stack - dims: [x, y] + dims: [longitude, latitude] target_output_variable: forcing era5_surface: path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' dims: [time, longitude, latitude, level] + coord_ranges: + latitude: {start: 40, end: 75} + longitude: {start: 0, end: 30} variables: - mean_sea_level_pressure dim_mapping: time: method: rename dim: time - x: - method: rename - dim: longitude - y: - method: rename - dim: latitude forcing_feature: method: stack_variables_by_var_name name_format: "{var_name}" grid_index: method: stack - dims: [x, y] + dims: [longitude, latitude] target_output_variable: forcing era5_static: path: 'gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr' dims: [time, longitude, latitude, level] + coord_ranges: + latitude: {start: 40, end: 75} + longitude: {start: 0, end: 30} variables: - land_sea_mask dim_mapping: - x: - method: rename - dim: longitude - y: - method: rename - dim: latitude static_feature: method: stack_variables_by_var_name name_format: "{var_name}" grid_index: method: stack - dims: [x, y] + dims: [longitude, latitude] target_output_variable: static extra: From 3e6e246e07c3e42ba161e71af4945cd534cb834f Mon Sep 17 00:00:00 2001 From: sadamov Date: Fri, 15 May 2026 10:27:03 +0200 Subject: [PATCH 10/35] test: enforce stateless boundary datastore in BoundaryDummyDatastore Drop state metadata after init and raise KeyError on `state` lookups so plotting/model code that accidentally queries state on a boundary fails loudly. Real ERA5-style boundary datastores expose only forcing fields, and the existing boundary tests (test_datasets.py) only ever access forcing on the boundary, so making the dummy state-less brings it closer to real boundary semantics without changing test behaviour. --- tests/dummy_datastore.py | 43 +++++++++++++++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/tests/dummy_datastore.py b/tests/dummy_datastore.py index a2c80b975..5608e9ca1 100644 --- a/tests/dummy_datastore.py +++ b/tests/dummy_datastore.py @@ -472,17 +472,22 @@ def grid_shape_state(self) -> CartesianGridShape: n_points_1d = int(np.sqrt(self.num_grid_points)) return CartesianGridShape(x=n_points_1d, y=n_points_1d) + @cached_property + def state_feature_weights_values(self) -> List[float]: + return [1.0] * self.N_FEATURES["state"] + class BoundaryDummyDatastore(DummyDatastore): - """DummyDatastore acting as a boundary forcing provider. + """DummyDatastore acting as a boundary forcing provider with no state. - Uses a different (smaller) grid and different forcing features to simulate - a boundary datastore from a separate domain. Only the ``forcing`` - category is meaningful; ``state`` and ``static`` are inherited but unused. + Mimics a real ERA5-style boundary datastore that only supplies ``forcing`` + fields (no ``state`` variables). State metadata is dropped after init and + state-keyed lookups raise ``KeyError`` so any code path that accidentally + asks the boundary for ``state`` fails loudly in tests. """ SHORT_NAME = "dummydata_boundary" - N_FEATURES = dict(state=5, forcing=3, static=1) + N_FEATURES = dict(state=0, forcing=3, static=1) def __init__(self, n_grid_points=400, n_timesteps=10, step_length=None): super().__init__( @@ -490,6 +495,34 @@ def __init__(self, n_grid_points=400, n_timesteps=10, step_length=None): n_timesteps=n_timesteps, step_length=step_length, ) + state_vars = [ + v + for v in ( + "state", + "state_feature", + "state_feature_units", + "state_feature_long_name", + ) + if v in self.ds.variables + ] + if state_vars: + self.ds = self.ds.drop_vars(state_vars) + + def get_xy(self, category: str, stacked: bool) -> ndarray: + if category == "state": + raise KeyError( + "BoundaryDummyDatastore has no state category; " + "use 'forcing' instead." + ) + return super().get_xy(category=category, stacked=stacked) + + def get_vars_names(self, category: str) -> list[str]: + if category == "state": + raise KeyError( + "BoundaryDummyDatastore has no state category; " + "use 'forcing' instead." + ) + return super().get_vars_names(category=category) class EnsembleDummyDatastore(BaseDatastore): From cf323a871bf6f408dd8c05997973d3ff95464969 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 17:07:20 +0200 Subject: [PATCH 11/35] fix(mdp): require state or forcing in datastore and simplify grid_shape_state Previously a datastore with no state, forcing or static would silently reach `grid_shape_state` and fail with a confusing fallback error. Now the constructor raises immediately if neither state nor forcing is present, and the fallback in `grid_shape_state` collapses to a simple `state if present else forcing` pick. Co-Authored-By: Claude Opus 4.7 --- neural_lam/datastore/mdp.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 53c200c2f..4abf20d49 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -74,6 +74,15 @@ def __init__(self, config_path, n_boundary_points=30, reuse_existing=True): if self._ds is None: self._ds = mdp.create_dataset(config=self._config) self._ds.to_zarr(fp_ds) + + # A valid datastore must provide at least state or forcing data; + # otherwise downstream code has nothing to slice or grid-shape against. + if "state" not in self._ds and "forcing" not in self._ds: + raise ValueError( + f"Datastore at {self._config_path} contains neither 'state' " + "nor 'forcing' data. At least one is required." + ) + self._n_boundary_points = n_boundary_points self.is_ensemble = ( "state" in self._ds and "ensemble_member" in self._ds["state"].dims @@ -472,14 +481,11 @@ def grid_shape_state(self): The shape of the cartesian grid for the state variables. """ - # Use state if available, otherwise fall back to forcing (e.g. - # boundary-only datastores that have no state variables). - for category in ("state", "forcing", "static"): - if category in self._ds: - ds_cat = self.unstack_grid_coords(self._ds[category]) - break - else: - raise ValueError("Dataset has no state, forcing, or static data") + # Use state if available, otherwise fall back to forcing (for + # boundary-only datastores with no state variables). The presence + # of at least one is guaranteed by __init__. + category = "state" if "state" in self._ds else "forcing" + ds_cat = self.unstack_grid_coords(self._ds[category]) xdim, ydim = self.spatial_coordinates da_x, da_y = ds_cat[xdim], ds_cat[ydim] assert da_x.ndim == da_y.ndim == 1 From 0ec56f56fa2ea511e6f7bfbf2b7340f3dadebe38 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 17:07:53 +0200 Subject: [PATCH 12/35] docs(models): mark boundary_features as not yet used in common_step Boundary features are unpacked from the batch but the forecaster forward pass does not consume them yet; the model-side wiring lands in a follow-up PR (#108). Co-Authored-By: Claude Opus 4.7 --- neural_lam/models/module.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index dacefff42..8380d7b4a 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -279,6 +279,8 @@ def common_step(self, batch): boundary_features, batch_times, ) = batch + # NOTE: boundary_features is not yet consumed here. Model-side + # boundary handling is implemented in a follow-up PR (see #108). prediction, pred_std = self.forecaster( init_states, forcing_features, target_states ) From 140caf5659d53bed8b711ff2bd324cd6c9d4f642 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 17:08:13 +0200 Subject: [PATCH 13/35] feat(dataset): align boundary forcing to interior times across analysis/forecast modes Replace integer-idx boundary slicing with time-based nearest-neighbor (pad) lookup so the boundary datastore can have a different step length than the interior, and either side may be in analysis or forecast mode. - Add `get_time_step`, `check_time_overlap`, `crop_time_if_needed` helpers in `neural_lam.utils` (ported from the research branch in joeloskarsson/neural-lam-dev, with an extra guard against silent argmax-on-all-false cropping). - Refactor `WeatherDataset`: precompute the within-sample state step and any forecast lead-time step in __init__; run `crop_time_if_needed` + `check_time_overlap` against the boundary so the first/last samples never fall outside boundary coverage; replace `_slice_forcing_time` with `_window_forcing_in_time` for time-aligned windowing of cross-datastore boundary; preserve the original integer- idx fast path as `_window_same_forecast_by_idx` for same-datastore forecast forcing (npyfilesmeps has non-unique analysis_time so the pandas pad-lookup cannot be used there). - Window alignment matches the existing forcing convention (target time, i.e. `state_times[init_steps + i]`). - Split `test_boundary_dataset_length_unchanged` into a no-crop and a cropping case to document the new behaviour. Co-Authored-By: Claude Opus 4.7 --- neural_lam/utils.py | 184 ++++++++++++++++ neural_lam/weather_dataset.py | 387 +++++++++++++++++++++------------- tests/test_datasets.py | 41 +++- 3 files changed, 463 insertions(+), 149 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index f3a937a7a..179af8f6e 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -10,8 +10,10 @@ from typing import Any, Iterator, Union # Third-party +import numpy as np import pytorch_lightning as pl import torch +import xarray as xr from loguru import logger from pytorch_lightning.loggers import MLFlowLogger, WandbLogger from pytorch_lightning.utilities import rank_zero_only @@ -644,3 +646,185 @@ def get_integer_time(tdelta: datetime.timedelta) -> tuple[int, str]: return int(total_seconds / unit_in_seconds), unit return 1, "unknown" + + +def get_time_step(times): + """Calculate the (constant) time step from a 1D time array. + + Parameters + ---------- + times : array-like + A 1D array of datetime64 (or timedelta64) values to compute the + step from. + + Returns + ------- + time_step : np.timedelta64 + The constant spacing between successive values. + + Raises + ------ + ValueError + If the spacing is not constant. + """ + time_diffs = np.diff(np.asarray(times)) + if not np.all(time_diffs == time_diffs[0]): + raise ValueError( + "Inconsistent time steps in data. " + f"Found different time steps: {np.unique(time_diffs)}" + ) + return time_diffs[0] + + +def check_time_overlap( + da1: xr.DataArray, + da2: xr.DataArray, + da1_is_forecast: bool = False, + da2_is_forecast: bool = False, + num_past_steps: int = 1, + num_future_steps: int = 1, +) -> None: + """Check that the time coverage of ``da2`` is wide enough to support + a windowed lookup driven by ``da1`` times with the given past/future + window sizes. + + Parameters + ---------- + da1 : xr.DataArray + Driving dataarray (typically interior state). + da2 : xr.DataArray + Dataarray to validate against (typically boundary forcing). + da1_is_forecast, da2_is_forecast : bool + Whether each side is in forecast mode (``analysis_time`` + + ``elapsed_forecast_duration`` dims) instead of plain ``time``. + num_past_steps, num_future_steps : int + Window size around each ``da1`` time, measured in ``da2`` steps. + + Raises + ------ + ValueError + If ``da2`` does not cover the required time range. + """ + if da1_is_forecast: + times_da1 = da1.analysis_time + else: + times_da1 = da1.time + time_min_da1 = times_da1.min().values + time_max_da1 = times_da1.max().values + + if da2_is_forecast: + times_da2 = da2.analysis_time + time_min_da2 = times_da2.min().values + time_max_da2 = times_da2.max().values + + time_step_da2 = get_time_step(times_da2.values) + time_step_da1 = get_time_step(times_da1.values) + + analysis_offset = max(time_step_da1, num_past_steps * time_step_da2) + da2_required_time_min = time_min_da1 - analysis_offset + da2_required_time_max = time_max_da1 - analysis_offset + else: + times_da2 = da2.time + time_min_da2 = times_da2.min().values + time_max_da2 = times_da2.max().values + time_step_da2 = get_time_step(times_da2.values) + + da2_required_time_min = time_min_da1 - num_past_steps * time_step_da2 + da2_required_time_max = time_max_da1 + num_future_steps * time_step_da2 + + if time_min_da2 > da2_required_time_min: + raise ValueError( + "The second DataArray (e.g. 'boundary forcing') starts too late. " + f"Required start: {da2_required_time_min}, " + f"but DataArray starts at {time_min_da2}." + ) + + if time_max_da2 < da2_required_time_max: + raise ValueError( + "The second DataArray (e.g. 'boundary forcing') ends too early. " + f"Required end: {da2_required_time_max}, " + f"but DataArray ends at {time_max_da2}." + ) + + +def crop_time_if_needed( + da1: xr.DataArray, + da2: xr.DataArray, + da1_is_forecast: bool = False, + da2_is_forecast: bool = False, + num_past_steps: int = 1, + num_future_steps: int = 1, +) -> xr.DataArray: + """Trim the leading/trailing times from ``da1`` so that ``da2`` covers + every needed window. If ``check_time_overlap`` already passes, ``da1`` + is returned unchanged. Cropping only applies to analysis-mode ``da1``; + for forecast-mode ``da1`` the input is returned as-is and overlap + failures will surface at sample-construction time. + + Parameters mirror :func:`check_time_overlap`. + + Returns + ------- + xr.DataArray + Possibly cropped ``da1``. + """ + if da1 is None or da2 is None: + return da1 + + try: + check_time_overlap( + da1, + da2, + da1_is_forecast, + da2_is_forecast, + num_past_steps, + num_future_steps, + ) + return da1 + except ValueError: + if da1_is_forecast: + # Cropping a forecast da1 by analysis time would change sample + # cardinality silently; leave alignment errors to surface later. + return da1 + + da1_tvals = da1.time.values + if da2_is_forecast: + da2_tvals = da2.analysis_time.values + else: + da2_tvals = da2.time.values + + da2_dt = get_time_step(da2_tvals) + if da2_is_forecast: + da1_dt = get_time_step(da1_tvals) + analysis_offset = max(da1_dt, num_past_steps * da2_dt) + required_min = da2_tvals[0] + analysis_offset + required_max = da2_tvals[-1] + analysis_offset + else: + required_min = da2_tvals[0] + num_past_steps * da2_dt + required_max = da2_tvals[-1] - num_future_steps * da2_dt + + begin_mask = da1_tvals >= required_min + if not begin_mask.any(): + raise ValueError( + "Boundary forcing ends before any interior time satisfies " + f"required_min={required_min}; cannot align." + ) + first_valid_idx = int(begin_mask.argmax()) + n_removed_begin = first_valid_idx + if da1_tvals[-1] > required_max: + end_mask = da1_tvals > required_max + last_valid_idx_plus_one = int(end_mask.argmax()) + n_removed_end = len(da1_tvals) - last_valid_idx_plus_one + else: + last_valid_idx_plus_one = None + n_removed_end = 0 + + if n_removed_begin > 0 or n_removed_end > 0: + log_on_rank_zero( + f"Cropping interior data to align with boundary: removed " + f"{n_removed_begin} steps at start and {n_removed_end} at " + "the end.", + level="warning", + ) + da1 = da1.isel(time=slice(first_valid_idx, last_valid_idx_plus_one)) + return da1 diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 434186f07..8ef2e986a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -11,6 +11,11 @@ # First-party from neural_lam.datastore.base import BaseDatastore +from neural_lam.utils import ( + check_time_overlap, + crop_time_if_needed, + get_time_step, +) class WeatherDataset(torch.utils.data.Dataset): @@ -18,6 +23,10 @@ class WeatherDataset(torch.utils.data.Dataset): This class loads and processes weather data from a given datastore, with optional boundary forcing from a separate boundary datastore. + Boundary windowing is aligned to interior state times by + nearest-neighbor lookup, so the interior and boundary datastores may + differ in step length and either side may be analysis or forecast + data. Parameters ---------- @@ -53,6 +62,8 @@ class WeatherDataset(torch.utils.data.Dataset): so all members are used when available. """ + INIT_STEPS = 2 + def __init__( self, datastore: BaseDatastore, @@ -88,7 +99,11 @@ def __init__( "The datastore must provide state data for the WeatherDataset." ) - # Load boundary forcing from the boundary datastore + # Load boundary forcing from the boundary datastore. Alignment to + # interior state times is done in `_window_forcing_in_time` via + # nearest-neighbor (pad) lookup on time coordinates, so the + # boundary datastore can have a different step length than the + # interior, and either side may be analysis or forecast. if self.datastore_boundary is not None: self.da_boundary_forcing = self.datastore_boundary.get_dataarray( category="forcing", split=self.split @@ -96,6 +111,55 @@ def __init__( else: self.da_boundary_forcing = None + # Within-sample time step for the state series: this is the step + # between consecutive state times that __getitem__ exposes, used + # below to advance the forcing/boundary window across AR steps. + if self.datastore.is_forecast: + self._time_step_state = get_time_step( + self.da_state.elapsed_forecast_duration.values + ) + else: + self._time_step_state = get_time_step(self.da_state.time.values) + + # Forecast lead-time step for forcing/boundary, only meaningful when + # the corresponding datastore is in forecast mode. + self._forecast_step_forcing = None + if self.da_forcing is not None and self.datastore.is_forecast: + self._forecast_step_forcing = get_time_step( + self.da_forcing.elapsed_forecast_duration.values + ) + self._forecast_step_boundary = None + if self.datastore_boundary is not None: + datastore_boundary = self.datastore_boundary + if ( + self.da_boundary_forcing is not None + and datastore_boundary.is_forecast + ): + self._forecast_step_boundary = get_time_step( + self.da_boundary_forcing.elapsed_forecast_duration.values + ) + + # Validate that the boundary covers the windows we will request, + # and if necessary crop the analysis-mode interior so that the + # very first/last samples don't fall outside boundary coverage. + if self.da_boundary_forcing is not None: + self.da_state = crop_time_if_needed( + self.da_state, + self.da_boundary_forcing, + da1_is_forecast=self.datastore.is_forecast, + da2_is_forecast=datastore_boundary.is_forecast, + num_past_steps=self.num_past_boundary_steps, + num_future_steps=self.num_future_boundary_steps, + ) + check_time_overlap( + self.da_state, + self.da_boundary_forcing, + da1_is_forecast=self.datastore.is_forecast, + da2_is_forecast=datastore_boundary.is_forecast, + num_past_steps=self.num_past_boundary_steps, + num_future_steps=self.num_future_boundary_steps, + ) + if self.datastore.is_ensemble and self.load_single_member: warnings.warn( "only using first ensemble member, so dataset size is " @@ -178,50 +242,28 @@ def __len__(self): return base_len def _slice_state_time(self, da_state, idx, n_steps: int): - """ - Produce a time slice of the given dataarray `da_state` (state) starting - at `idx` and with `n_steps` steps. An `offset`is calculated based on the - `num_past_forcing_steps` class attribute. `Offset` is used to offset the - start of the sample, to assert that enough previous time steps are - available for the 2 initial states and any corresponding forcings - (calculated in `_slice_forcing_time`). + """Slice ``da_state`` by integer ``idx`` into one training sample. - Parameters - ---------- - da_state : xr.DataArray - The dataarray to slice. This is expected to have a `time` dimension - if the datastore is providing analysis only data, and a - `analysis_time` and `elapsed_forecast_duration` dimensions if the - datastore is providing forecast data. - idx : int - The index of the time step to start the sample from. - n_steps : int - The number of time steps to include in the sample. + For analysis data the sample's ``time`` is contiguous; for forecast + data we pick a single ``analysis_time`` and walk its lead times. + The leading offset accounts for ``num_past_forcing_steps`` so the + forcing window of the very first sample is in-bounds. Returns ------- da_sliced : xr.DataArray - The sliced dataarray with dims ('time', 'grid_index', - 'state_feature'). + Sliced state with a single ``time`` dimension covering + ``INIT_STEPS + n_steps`` consecutive state times. """ - # The current implementation requires at least 2 time steps for the - # initial state (see GraphCast). - init_steps = 2 - # slice the dataarray to include the required number of time steps + init_steps = self.INIT_STEPS + n_total = init_steps + n_steps + offset = max(0, self.num_past_forcing_steps - init_steps) + if self.datastore.is_forecast: - start_idx = max(0, self.num_past_forcing_steps - init_steps) - end_idx = max(init_steps, self.num_past_forcing_steps) + n_steps - # this implies that the data will have both `analysis_time` and - # `elapsed_forecast_duration` dimensions for forecasts. We for now - # simply select a analysis time and the first `n_steps` forecast - # times (given no offset). Note that this means that we get one - # sample per forecast, always starting at forecast time 2. da_sliced = da_state.isel( analysis_time=idx, - elapsed_forecast_duration=slice(start_idx, end_idx), + elapsed_forecast_duration=slice(offset, offset + n_total), ) - # create a new time dimension so that the produced sample has a - # `time` dimension, similarly to the analysis only data da_sliced["time"] = ( da_sliced.analysis_time + da_sliced.elapsed_forecast_duration ) @@ -229,136 +271,170 @@ def _slice_state_time(self, da_state, idx, n_steps: int): {"elapsed_forecast_duration": "time"} ) else: - # For analysis data we slice the time dimension directly. The offset - # is only relevant for the very first (and last) samples in the - # dataset. - start_idx = idx + max(0, self.num_past_forcing_steps - init_steps) - end_idx = ( - idx + max(init_steps, self.num_past_forcing_steps) + n_steps + start_idx = idx + offset + da_sliced = da_state.isel( + time=slice(start_idx, start_idx + n_total) ) - da_sliced = da_state.isel(time=slice(start_idx, end_idx)) return da_sliced - def _slice_forcing_time( + def _window_same_forecast_by_idx( self, da_forcing, - idx, - n_steps: int, - num_past_steps: Union[int, None] = None, - num_future_steps: Union[int, None] = None, + idx: int, + state_times, + num_past_steps: int, + num_future_steps: int, ): + """Window forcing from the same forecast datastore as state. + + Uses integer ``analysis_time=idx`` indexing so it tolerates + repeated analysis_time values (e.g. npyfilesmeps duplicates the + analysis_time series). Walks lead times in lockstep with the + state slice; each window is centered on the corresponding target + state time. """ - Produce a time slice of the given dataarray `da_forcing` (forcing) - starting at `idx` and with `n_steps` steps. An `offset` is calculated - based on the `num_past_forcing_steps` class attribute. It is used to - offset the start of the sample, to ensure that enough previous time - steps are available for the forcing data. The forcing data is windowed - around the current autoregressive time step to include the past and - future forcings. + init_steps = self.INIT_STEPS + offset = max(0, self.num_past_forcing_steps - init_steps) + init_steps + da_list = [] + for step in range(self.ar_steps): + start_lead = offset + step - num_past_steps + end_lead = offset + step + num_future_steps + 1 + target_time = state_times[init_steps + step].values - Parameters - ---------- - da_forcing : xr.DataArray - The forcing dataarray to slice. This is expected to have a `time` - dimension if the datastore is providing analysis only data, and a - `analysis_time` and `elapsed_forecast_duration` dimensions if the - datastore is providing forecast data. - idx : int - The index of the time step to start the sample from. - n_steps : int - The number of time steps to include in the sample. + da_sliced = da_forcing.isel( + analysis_time=idx, + elapsed_forecast_duration=slice(start_lead, end_lead), + ).rename({"elapsed_forecast_duration": "window"}) + da_sliced = da_sliced.assign_coords( + window=np.arange(-num_past_steps, num_future_steps + 1) + ) + da_sliced = da_sliced.expand_dims(dim={"time": [target_time]}) + da_list.append(da_sliced) + return xr.concat(da_list, dim="time") + + def _window_forcing_in_time( + self, + da_forcing, + state_times, + num_past_steps: int, + num_future_steps: int, + forecast_step, + ): + """Window forcing/boundary in time, aligned to interior state times. + + ``state_times`` is the 1D ``time`` coordinate of the already-sliced + state sample. For each AR target step the matching forcing time is + picked by nearest-neighbor ``pad`` lookup (smallest forcing time + ``<=`` state time), and a window of + ``num_past_steps + num_future_steps + 1`` consecutive forcing + entries is taken around it. + + When ``da_forcing`` has an ``analysis_time`` dimension the same + logic is applied to forecast forcing/boundary: an analysis time is + chosen such that the lead times cover the requested window for + every AR step, then windows are walked across lead times. Returns ------- - da_concat : xr.DataArray - The sliced dataarray with dims ('time', 'grid_index', - 'window', 'forcing_feature'). + xr.DataArray + Concatenated windows with dims + ``('time', 'grid_index', 'window', 'forcing_feature')``. """ - # The current implementation requires at least 2 time steps for the - # initial state (see GraphCast). The forcing data is windowed around the - # current autregressive time step. The two `init_steps` can also be used - # as past forcings. - init_steps = 2 + init_steps = self.INIT_STEPS da_list = [] - # Allow overriding the window size (used for boundary forcing) - if num_past_steps is None: - num_past_steps = self.num_past_forcing_steps - if num_future_steps is None: - num_future_steps = self.num_future_forcing_steps - - if self.datastore.is_forecast: - # This implies that the data will have both `analysis_time` and - # `elapsed_forecast_duration` dimensions for forecasts. We for now - # simply select an analysis time and the first `n_steps` forecast - # times (given no offset). Note that this means that we get one - # sample per forecast. - # Add a 'time' dimension using the actual forecast times - offset = max(init_steps, num_past_steps) - for step in range(n_steps): - start_idx = offset + step - num_past_steps - end_idx = offset + step + num_future_steps - - current_time = ( - da_forcing.analysis_time[idx] - + da_forcing.elapsed_forecast_duration[offset + step] + if "analysis_time" in da_forcing.dims: + if forecast_step is None: + raise ValueError( + "forecast_step must be supplied when forcing/boundary " + "is in forecast mode." ) - - da_sliced = da_forcing.isel( - analysis_time=idx, - elapsed_forecast_duration=slice(start_idx, end_idx + 1), + # Choose a single analysis_time for this sample. We want it to + # be the latest analysis_time strictly less than the first + # target time (so the first target sits at lead >= 1, leaving + # room for past-window lookups), shifted further back if a + # larger num_past_steps requires more lead headroom. + first_target_time = state_times[init_steps].values + + analysis_index = da_forcing.analysis_time.get_index("analysis_time") + forcing_at_idx = analysis_index.get_indexer( + [first_target_time], method="pad" + )[0] + if forcing_at_idx < 0: + raise ValueError( + "Boundary/forcing analysis times start after the first " + f"target time ({first_target_time})." ) + forcing_at = da_forcing.analysis_time[forcing_at_idx] + if first_target_time == forcing_at.values: + if forcing_at_idx == 0: + raise ValueError( + "No earlier boundary/forcing analysis time is " + "available to provide lead headroom for target " + f"time {first_target_time}." + ) + forcing_at_idx -= 1 + forcing_at = da_forcing.analysis_time[forcing_at_idx] - da_sliced = da_sliced.rename( - {"elapsed_forecast_duration": "window"} + lead_at_first_target = int( + np.floor( + (first_target_time - forcing_at.values) / forecast_step ) + ) + past_analysis_offset = num_past_steps - lead_at_first_target + if past_analysis_offset > 0: + forcing_at_idx -= past_analysis_offset + if forcing_at_idx < 0: + raise ValueError( + "Boundary/forcing analysis times do not extend far " + "enough back to cover the requested past window." + ) + forcing_at = da_forcing.analysis_time[forcing_at_idx] - # Assign the 'window' coordinate to be relative positions - da_sliced = da_sliced.assign_coords( - window=np.arange(len(da_sliced.window)) + for step_idx in range(len(state_times) - init_steps): + target_time = state_times[init_steps + step_idx].values + lead = int( + np.floor((target_time - forcing_at.values) / forecast_step) ) + window_start = lead - num_past_steps + window_end = lead + num_future_steps + 1 - da_sliced = da_sliced.expand_dims( - dim={"time": [current_time.values]} + da_sliced = da_forcing.isel( + analysis_time=int(forcing_at_idx), + elapsed_forecast_duration=slice( + int(window_start), int(window_end) + ), + ).rename({"elapsed_forecast_duration": "window"}) + da_sliced = da_sliced.assign_coords( + window=np.arange(-num_past_steps, num_future_steps + 1) ) - + da_sliced = da_sliced.expand_dims(dim={"time": [target_time]}) da_list.append(da_sliced) - - # Concatenate the list of DataArrays along the 'time' dimension - da_concat = xr.concat(da_list, dim="time") - else: - # For analysis data, we slice the time dimension directly. The - # offset is only relevant for the very first (and last) samples in - # the dataset. - offset = idx + max(init_steps, num_past_steps) - for step in range(n_steps): - start_idx = offset + step - num_past_steps - end_idx = offset + step + num_future_steps - - # Slice the data over the desired time window - da_sliced = da_forcing.isel(time=slice(start_idx, end_idx + 1)) - - da_sliced = da_sliced.rename({"time": "window"}) + forcing_time_index = da_forcing.time.get_index("time") + for step_idx in range(init_steps, len(state_times)): + state_time = state_times[step_idx].values + forcing_time_idx = forcing_time_index.get_indexer( + [state_time], method="pad" + )[0] + if forcing_time_idx < 0: + raise ValueError( + f"No boundary/forcing time at or before {state_time}." + ) - # Assign the 'window' coordinate to be relative positions - da_sliced = da_sliced.assign_coords( - window=np.arange(len(da_sliced.window)) - ) + window_start = forcing_time_idx - num_past_steps + window_end = forcing_time_idx + num_future_steps + 1 - # Add a 'time' dimension to keep track of steps using actual - # time coordinates - current_time = da_forcing.time[offset + step] - da_sliced = da_sliced.expand_dims( - dim={"time": [current_time.values]} + da_window = da_forcing.isel( + time=slice(int(window_start), int(window_end)) + ).rename({"time": "window"}) + da_window = da_window.assign_coords( + window=np.arange(-num_past_steps, num_future_steps + 1) ) + da_window = da_window.expand_dims(dim={"time": [state_time]}) + da_list.append(da_window) - da_list.append(da_sliced) - - # Concatenate the list of DataArrays along the 'time' dimension - da_concat = xr.concat(da_list, dim="time") - - return da_concat + return xr.concat(da_list, dim="time") def _build_item_dataarrays(self, idx): """ @@ -405,24 +481,43 @@ def _build_item_dataarrays(self, idx): else: da_forcing = None - # handle time sampling in a way that is compatible with both analysis - # and forecast data + # Slice the state once, then window forcing and boundary against + # the resulting state times. Forcing is windowed by integer + # `analysis_time` index when it comes from the same forecast + # datastore as state (the analysis_time series can have repeats + # there, e.g. npyfilesmeps); boundary always comes from a + # different datastore so it is windowed by time-based + # nearest-neighbor lookup. da_state = self._slice_state_time( da_state=da_state, idx=sample_idx, n_steps=self.ar_steps ) + state_times = da_state["time"] + if da_forcing is not None: - da_forcing_windowed = self._slice_forcing_time( - da_forcing=da_forcing, idx=sample_idx, n_steps=self.ar_steps - ) + if self.datastore.is_forecast: + da_forcing_windowed = self._window_same_forecast_by_idx( + da_forcing=da_forcing, + idx=sample_idx, + state_times=state_times, + num_past_steps=self.num_past_forcing_steps, + num_future_steps=self.num_future_forcing_steps, + ) + else: + da_forcing_windowed = self._window_forcing_in_time( + da_forcing=da_forcing, + state_times=state_times, + num_past_steps=self.num_past_forcing_steps, + num_future_steps=self.num_future_forcing_steps, + forecast_step=None, + ) - # Slice boundary forcing if available if self.da_boundary_forcing is not None: - da_boundary_windowed = self._slice_forcing_time( + da_boundary_windowed = self._window_forcing_in_time( da_forcing=self.da_boundary_forcing, - idx=sample_idx, - n_steps=self.ar_steps, + state_times=state_times, num_past_steps=self.num_past_boundary_steps, num_future_steps=self.num_future_boundary_steps, + forecast_step=self._forecast_step_boundary, ) else: da_boundary_windowed = None diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 42c3818b9..ebc2dbaf0 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -516,10 +516,13 @@ def test_boundary_datastore_none_gives_empty_boundary(): assert boundary.shape[-1] == 0 -def test_boundary_dataset_length_unchanged(): - """Adding a boundary datastore should not change the dataset length.""" +def test_boundary_dataset_length_unchanged_when_boundary_covers(): + """Adding a boundary datastore that covers the requested past/future + window does not change the dataset length.""" n_timesteps = 20 datastore = DummyDatastore(n_grid_points=100, n_timesteps=n_timesteps) + # num_past_boundary=num_future_boundary=0 means the boundary only + # needs to cover the interior times themselves, no padding. boundary_ds = BoundaryDummyDatastore( n_grid_points=25, n_timesteps=n_timesteps ) @@ -532,6 +535,36 @@ def test_boundary_dataset_length_unchanged(): num_future_forcing_steps=1, ) + dataset_with_boundary = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=3, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + num_past_boundary_steps=0, + num_future_boundary_steps=0, + datastore_boundary=boundary_ds, + ) + + assert len(dataset_no_boundary) == len(dataset_with_boundary) + + +def test_boundary_crops_interior_when_window_overflows(): + """When the boundary does not cover the requested past/future window, + interior is cropped at start/end and the dataset shrinks accordingly.""" + n_timesteps = 20 + datastore = DummyDatastore(n_grid_points=100, n_timesteps=n_timesteps) + boundary_ds = BoundaryDummyDatastore( + n_grid_points=25, n_timesteps=n_timesteps + ) + + dataset_no_boundary = WeatherDataset( + datastore=datastore, + split="train", + ar_steps=3, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + ) dataset_with_boundary = WeatherDataset( datastore=datastore, split="train", @@ -543,4 +576,6 @@ def test_boundary_dataset_length_unchanged(): datastore_boundary=boundary_ds, ) - assert len(dataset_no_boundary) == len(dataset_with_boundary) + # Boundary spans the same range as interior, so a (past=1, future=1) + # window forces 1 step of cropping at each end. + assert len(dataset_with_boundary) == len(dataset_no_boundary) - 2 From cbc1c147af5c2c16ffb11c67a14321823c2895ff Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 17:08:28 +0200 Subject: [PATCH 14/35] test(time-slicing): cover boundary windowing and mixed forecast modes Port a slimmed-down version of the alignment tests from joeloskarsson/neural-lam-dev to exercise the new time-based boundary windowing in WeatherDataset. - Extend `SinglePointDummyDatastore` with forecast-mode support. - Add a boundary-only variant whose state lookup raises KeyError, to catch any path that accidentally asks the boundary for state. - `test_time_slicing_boundary_analysis`: parametrised over past/future window sizes, asserts exact boundary window values around each target state time. - `test_boundary_step_length_mismatch_supported`: 1h interior with a 6h boundary, verifies the pad-matched lookup. - `test_forecast_interior_with_analysis_boundary` and `test_analysis_interior_with_forecast_boundary`: the two mixed analysis/forecast combinations. - `test_check_time_overlap_insufficient_raises`: surface the cropping failure path with a clear error. Co-Authored-By: Claude Opus 4.7 --- tests/test_time_slicing.py | 344 ++++++++++++++++++++++++++++++++++++- 1 file changed, 341 insertions(+), 3 deletions(-) diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index a7ffd2b80..4575ebf0d 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -13,6 +13,17 @@ class SinglePointDummyDatastore(BaseDatastore): + """One-grid-point datastore in either analysis or forecast mode. + + Analysis mode: ``time_values`` is a 1D datetime array, ``state_data`` + and ``forcing_data`` are 1D arrays aligned to it. + + Forecast mode: ``time_values`` is the pair + ``(analysis_times, elapsed_forecast_durations)``, and ``state_data`` + / ``forcing_data`` are 2D arrays shaped + ``(n_analysis_times, n_forecast_steps)``. + """ + config = {} coords_projection = None num_grid_points = 1 @@ -27,14 +38,16 @@ def __init__( step_length=timedelta(hours=1), ): self._step_length = step_length - self._time_values = np.array(time_values) self._state_data = np.array(state_data) self._forcing_data = np.array(forcing_data) self.is_forecast = is_forecast if is_forecast: + self._analysis_times = np.array(time_values[0]) + self._forecast_times = np.array(time_values[1]) assert self._state_data.ndim == 2 else: + self._time_values = np.array(time_values) assert self._state_data.ndim == 1 @property @@ -53,12 +66,18 @@ def get_dataarray(self, category, split): raise NotImplementedError(category) if self.is_forecast: - raise NotImplementedError() + da = xr.DataArray( + values, + dims=["analysis_time", "elapsed_forecast_duration"], + coords={ + "analysis_time": self._analysis_times, + "elapsed_forecast_duration": self._forecast_times, + }, + ) else: da = xr.DataArray( values, dims=["time"], coords={"time": self._time_values} ) - # add `{category}_feature` and `grid_index` dimensions da = da.expand_dims("grid_index") da = da.expand_dims(f"{category}_feature") @@ -84,6 +103,56 @@ def get_vars_long_names(self, category): ANALYSIS_STATE_VALUES = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] FORCING_VALUES = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19] +# Boundary spans 4 extra steps on each side of the interior so windowing +# with up to num_past/num_future = 4 can be tested without cropping. +BOUNDARY_PAD = 4 +BOUNDARY_FORCING_VALUES = list(range(20, 20 + 10 + 2 * BOUNDARY_PAD)) + +FORECAST_STATE_VALUES = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], + [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], + [30, 31, 32, 33, 34, 35, 36, 37, 38, 39], +] +FORECAST_FORCING_VALUES = [ + [100, 101, 102, 103, 104, 105, 106, 107, 108, 109], + [110, 111, 112, 113, 114, 115, 116, 117, 118, 119], + [120, 121, 122, 123, 124, 125, 126, 127, 128, 129], + [130, 131, 132, 133, 134, 135, 136, 137, 138, 139], +] + + +class BoundaryOnlyDummyDatastore(SinglePointDummyDatastore): + """Boundary-only variant providing forcing but no state. + + State-keyed lookups raise KeyError to mirror real boundary datastores + (e.g. ERA5) and to catch any path that accidentally asks the boundary + for state. + """ + + def __init__( + self, + time_values, + forcing_data, + is_forecast=False, + step_length=timedelta(hours=1), + ): + # state_data is a dummy zeros array of the right shape so the + # parent constructor accepts it; the override below blocks state + # access. + forcing_arr = np.asarray(forcing_data) + super().__init__( + time_values=time_values, + state_data=np.zeros_like(forcing_arr), + forcing_data=forcing_arr, + is_forecast=is_forecast, + step_length=step_length, + ) + + def get_dataarray(self, category, split): + if category == "state": + raise KeyError("BoundaryOnlyDummyDatastore has no state category.") + return super().get_dataarray(category=category, split=split) @pytest.mark.parametrize( @@ -193,3 +262,272 @@ def test_step_length_timedeltas(step_length): assert ( len(sample) == 5 ) # init_states, target_states, forcing, boundary, target_times + + +def _interior_times(): + return np.datetime64("2020-01-01") + np.arange(len(ANALYSIS_STATE_VALUES)) + + +def _boundary_times_aligned(): + """Boundary times surrounding the interior on both sides so that + windows up to BOUNDARY_PAD steps don't trigger cropping.""" + return ( + np.datetime64("2020-01-01") + - BOUNDARY_PAD + + np.arange(len(BOUNDARY_FORCING_VALUES)) + ) + + +@pytest.mark.parametrize( + "ar_steps,num_past_boundary_steps,num_future_boundary_steps", + [ + [3, 0, 0], + [3, 1, 0], + [3, 0, 1], + [3, 1, 1], + [3, 2, 2], + [3, 3, 1], + [3, 1, 3], + ], +) +def test_time_slicing_boundary_analysis( + ar_steps, num_past_boundary_steps, num_future_boundary_steps +): + """Boundary windowing for analysis-interior + analysis-boundary. + + Boundary spans BOUNDARY_PAD extra steps on each side of the interior + so no cropping kicks in; the exact window values around each state + time are checked.""" + interior_datastore = SinglePointDummyDatastore( + state_data=ANALYSIS_STATE_VALUES, + forcing_data=FORCING_VALUES, + time_values=_interior_times(), + is_forecast=False, + ) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=BOUNDARY_FORCING_VALUES, + time_values=_boundary_times_aligned(), + is_forecast=False, + ) + + dataset = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=ar_steps, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=num_past_boundary_steps, + num_future_boundary_steps=num_future_boundary_steps, + ) + + _, _, _, boundary, _ = [tensor.numpy() for tensor in dataset[0]] + + # Interior sample idx=0 has state slice [t_0..t_4] (no past-forcing + # offset since num_past_forcing=0). Target states start at t_2; the + # boundary index for t_2 in BOUNDARY_FORCING_VALUES is BOUNDARY_PAD+2. + boundary_center = BOUNDARY_PAD + 2 + window_size = num_past_boundary_steps + num_future_boundary_steps + 1 + assert boundary.shape == (ar_steps, 1, window_size) + for i in range(ar_steps): + start = boundary_center + i - num_past_boundary_steps + end = boundary_center + i + num_future_boundary_steps + 1 + expected = BOUNDARY_FORCING_VALUES[start:end] + np.testing.assert_array_equal(boundary[i, 0, :], expected) + + +def test_boundary_step_length_mismatch_supported(): + """Interior and boundary with different step lengths align by time: + a 6h boundary still produces correctly-windowed slices around the + 1h interior times.""" + interior_times = np.datetime64("2020-01-01") + np.arange( + 24 + ) * np.timedelta64(1, "h") + interior_values = np.arange(24, dtype=float) + + # Boundary every 6h, covering the same calendar span plus a 6h pad + # on each end so the past/future window stays in-bounds. + boundary_times = np.datetime64("2019-12-31T18:00") + np.arange( + 7 + ) * np.timedelta64(6, "h") + boundary_values = np.arange(100, 107, dtype=float) + + interior_datastore = SinglePointDummyDatastore( + state_data=interior_values, + forcing_data=interior_values, + time_values=interior_times, + is_forecast=False, + step_length=timedelta(hours=1), + ) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=boundary_values, + time_values=boundary_times, + is_forecast=False, + step_length=timedelta(hours=6), + ) + + dataset = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=2, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + ) + + _, _, _, boundary, _ = [tensor.numpy() for tensor in dataset[0]] + # First target state is at hour 2; nearest boundary <= hour 2 is hour 0 + # (= boundary_values[1] = 101). Window [past=1, future=1] takes + # boundary_values[0], boundary_values[1], boundary_values[2]. + assert boundary.shape == (2, 1, 3) + np.testing.assert_array_equal(boundary[0, 0, :], [100, 101, 102]) + np.testing.assert_array_equal(boundary[1, 0, :], [100, 101, 102]) + + +def test_forecast_interior_with_analysis_boundary(): + """Forecast-mode interior + analysis-mode boundary: boundary windows + around each lead-time of the forecast pick the corresponding boundary + times.""" + analysis_times = np.datetime64("2020-01-01") + np.arange( + len(FORECAST_STATE_VALUES) + ) * np.timedelta64(1, "D") + forecast_durations = np.arange( + len(FORECAST_STATE_VALUES[0]) + ) * np.timedelta64(1, "D") + + interior_datastore = SinglePointDummyDatastore( + state_data=FORECAST_STATE_VALUES, + forcing_data=FORECAST_FORCING_VALUES, + time_values=(analysis_times, forecast_durations), + is_forecast=True, + step_length=timedelta(days=1), + ) + + # Boundary covers analysis_time[0] + leads, padded on both sides. + boundary_times = np.datetime64("2019-12-30") + np.arange( + 12 + ) * np.timedelta64(1, "D") + boundary_values = np.arange(200, 212, dtype=float) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=boundary_values, + time_values=boundary_times, + is_forecast=False, + step_length=timedelta(days=1), + ) + + dataset = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=3, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + ) + + init_states, target_states, _, boundary, _ = [t.numpy() for t in dataset[0]] + # Sample idx=0: pick analysis_time[0] (2020-01-01), state at lead + # 0..4 = [0,1,2,3,4]. Init=[0,1], target=[2,3,4]. State times are + # 2020-01-01 + (0..4) days = 01..05. + np.testing.assert_array_equal(init_states[:, 0, 0], [0, 1]) + np.testing.assert_array_equal(target_states[:, 0, 0], [2, 3, 4]) + # Boundary starts at 2019-12-30 (idx 0). Target state times 03..05 + # correspond to boundary idx 4..6, with past/future windows of 1. + assert boundary.shape == (3, 1, 3) + np.testing.assert_array_equal(boundary[0, 0, :], [203, 204, 205]) + np.testing.assert_array_equal(boundary[1, 0, :], [204, 205, 206]) + np.testing.assert_array_equal(boundary[2, 0, :], [205, 206, 207]) + + +def test_analysis_interior_with_forecast_boundary(): + """Analysis-mode interior + forecast-mode boundary: an analysis time + of the boundary forecast is picked so the requested past/future + window around each target state time stays in lead-range, then + lead-time windows are walked across AR steps.""" + interior_times = np.datetime64("2020-01-05") + np.arange( + 8 + ) * np.timedelta64(1, "D") + interior_values = np.arange(8, dtype=float) + interior_datastore = SinglePointDummyDatastore( + state_data=interior_values, + forcing_data=interior_values, + time_values=interior_times, + is_forecast=False, + step_length=timedelta(days=1), + ) + + # Boundary: 6 analysis times, 8 lead-day steps each. Analysis times + # 2020-01-04..09 so coverage extends past the latest interior + # target times after cropping. + n_analysis = 6 + n_leads = 8 + boundary_analysis = np.datetime64("2020-01-04") + np.arange( + n_analysis + ) * np.timedelta64(1, "D") + boundary_leads = np.arange(n_leads) * np.timedelta64(1, "D") + boundary_values = ( + np.arange(n_analysis).reshape(-1, 1) * 1000 + + np.arange(n_leads).reshape(1, -1) * 10 + ).astype(float) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=boundary_values, + time_values=(boundary_analysis, boundary_leads), + is_forecast=True, + step_length=timedelta(days=1), + ) + + dataset = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=2, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + ) + + _, _, _, boundary, _ = [t.numpy() for t in dataset[0]] + # Sample idx=0: state slice = interior[0:4] = times 2020-01-05..08. + # Targets are 2020-01-07 and 2020-01-08 (first_target=07). Boundary + # analysis_time pad-pick for 07 = idx 3 (07); equals first_target so + # decrement to idx 2 (06). lead_at_first_target = (07-06)/1d = 1, + # which already covers num_past=1, so no further shift. Window at + # target 07: lead 1, [0..2]. Window at target 08: lead 2, [1..3]. + expected_analysis_idx = 2 + assert boundary.shape == (2, 1, 3) + np.testing.assert_array_equal( + boundary[0, 0, :], boundary_values[expected_analysis_idx, 0:3] + ) + np.testing.assert_array_equal( + boundary[1, 0, :], boundary_values[expected_analysis_idx, 1:4] + ) + + +def test_check_time_overlap_insufficient_raises(): + """If the boundary cannot be cropped enough to cover the requested + past-window, ``check_time_overlap`` surfaces a clear error.""" + interior_datastore = SinglePointDummyDatastore( + state_data=ANALYSIS_STATE_VALUES, + forcing_data=FORCING_VALUES, + time_values=_interior_times(), + is_forecast=False, + ) + # Boundary covers the same range as interior but no padding, so + # any non-zero past/future window forces cropping; with a huge past + # window the boundary cannot cover even a single sample. + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=BOUNDARY_FORCING_VALUES[:10], + time_values=_interior_times(), + is_forecast=False, + ) + + with pytest.raises(ValueError): + WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=3, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=20, + num_future_boundary_steps=20, + ) From 9a9af315d9d9eeb9345cac7a0715e9125b2b0458 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 17:08:53 +0200 Subject: [PATCH 15/35] feat(models): standardize boundary forcing on-device in ForecasterModule Boundary forcing was previously passed through unchanged in `on_after_batch_transfer`, leaving the only normalization step on a separate code path and inconsistent with how interior state/forcing are handled. Wire it through the same on-device hook. - `ForecasterModule.__init__` takes a new optional `datastore_boundary` arg (excluded from `save_hyperparameters` alongside `datastore` and `forecaster`, so it must be passed at `load_from_checkpoint` time). - Register `boundary_mean` / `boundary_std` from `datastore_boundary.get_standardization_dataarray("forcing")` when a boundary datastore is provided; otherwise leave both as None. - `on_after_batch_transfer` standardizes the boundary tensor the same way it standardizes forcing: feature-major `(feature, window)` per-feature stats are tiled once on the first batch and cached. - Update the NOTE in `common_step` to reflect that boundary is now standardized but still not consumed by the forecaster (#108). - Pass `datastore_boundary` from `train_model.main` to the module. - New tests: `test_boundary_standardized_when_datastore_provided` and `test_boundary_passthrough_when_no_boundary_datastore`. Co-Authored-By: Claude Opus 4.7 --- neural_lam/models/module.py | 64 +++++++++++++++++++++++++--- neural_lam/train_model.py | 1 + tests/test_gpu_normalization.py | 75 ++++++++++++++++++++++++++++++++- 3 files changed, 132 insertions(+), 8 deletions(-) diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 8380d7b4a..a70a35808 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -36,6 +36,7 @@ def __init__( forecaster: Forecaster, config: NeuralLAMConfig, datastore: BaseDatastore, + datastore_boundary: Optional[BaseDatastore] = None, loss: str = "wmse", lr: float = 1e-3, restore_opt: bool = False, @@ -82,8 +83,11 @@ def __init__( # makes the checkpoint self-describing: it carries model, graph_name, # hidden_dim, etc. so the caller can reconstruct the exact forecaster # architecture from the checkpoint alone. - self.save_hyperparameters(ignore=["datastore", "forecaster"]) + self.save_hyperparameters( + ignore=["datastore", "datastore_boundary", "forecaster"] + ) self.datastore = datastore + self.datastore_boundary = datastore_boundary self.forecaster = forecaster self.matched_metrics: set = set() @@ -168,6 +172,38 @@ def __init__( self.forcing_mean = None self.forcing_std = None + # Boundary standardization: registered only when a boundary + # datastore is wired in. The same feature-major window tiling as + # forcing applies, so we cache the tiled mean/std on first batch. + if ( + datastore_boundary is not None + and datastore_boundary.get_num_data_vars(category="forcing") > 0 + ): + da_boundary_stats = ( + datastore_boundary.get_standardization_dataarray( + category="forcing" + ) + ) + self.register_buffer( + "boundary_mean", + torch.tensor( + da_boundary_stats.forcing_mean.values, dtype=torch.float32 + ), + persistent=False, + ) + self.register_buffer( + "boundary_std", + self._safe_std( + da_boundary_stats.forcing_std.values, eps, "boundary" + ), + persistent=False, + ) + self.register_buffer("boundary_mean_tiled", None, persistent=False) + self.register_buffer("boundary_std_tiled", None, persistent=False) + else: + self.boundary_mean = None + self.boundary_std = None + # Instantiate loss function self.loss = metrics.get_metric(loss) @@ -242,9 +278,11 @@ def on_after_batch_transfer(self, batch, dataloader_idx): """Standardize a batch on-device after transfer to the accelerator. Lightning calls this for every train/val/test/predict batch. - WeatherDataset returns unstandardized state and forcing; both are - normalized here so the work runs on the accelerator. The boundary - forcing is passed through unchanged. + WeatherDataset returns unstandardized state, forcing and boundary; + all three are normalized here so the work runs on the accelerator. + Forcing and boundary share the same feature-major + ``(feature, window)`` stacking, so per-feature mean/std are tiled + once and cached. """ init_states, target_states, forcing, boundary, batch_times = batch @@ -269,6 +307,19 @@ def on_after_batch_transfer(self, batch, dataloader_idx): forcing - self.forcing_mean_tiled ) / self.forcing_std_tiled + if boundary.shape[-1] > 0 and self.boundary_mean is not None: + if self.boundary_mean_tiled is None: + window_size = boundary.shape[-1] // self.boundary_mean.shape[-1] + self.boundary_mean_tiled = self.boundary_mean.repeat_interleave( + window_size + ) + self.boundary_std_tiled = self.boundary_std.repeat_interleave( + window_size + ) + boundary = ( + boundary - self.boundary_mean_tiled + ) / self.boundary_std_tiled + return init_states, target_states, forcing, boundary, batch_times def common_step(self, batch): @@ -279,8 +330,9 @@ def common_step(self, batch): boundary_features, batch_times, ) = batch - # NOTE: boundary_features is not yet consumed here. Model-side - # boundary handling is implemented in a follow-up PR (see #108). + # NOTE: boundary_features is standardized in + # `on_after_batch_transfer` but not yet consumed by the forecaster. + # The model-side boundary handling lands in a follow-up PR (#108). prediction, pred_std = self.forecaster( init_states, forcing_features, target_states ) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index f714a4522..117c6b007 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -422,6 +422,7 @@ def main(input_args=None): forecaster=forecaster, config=config, datastore=datastore, + datastore_boundary=datastore_boundary, loss=args.loss, lr=args.lr, restore_opt=args.restore_opt, diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index dfa201ef1..5fd018efd 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -7,6 +7,7 @@ from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example +from tests.dummy_datastore import BoundaryDummyDatastore, DummyDatastore NUM_PAST_FORCING_STEPS = 1 NUM_FUTURE_FORCING_STEPS = 1 @@ -19,7 +20,7 @@ def forward(self, prev_state, prev_prev_state, forcing): return torch.zeros_like(prev_state), None -def _build_module(datastore): +def _build_module(datastore, datastore_boundary=None): config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path @@ -28,7 +29,10 @@ def _build_module(datastore): predictor = _MockStepPredictor(datastore=datastore, output_std=False) forecaster = ARForecaster(predictor, datastore) return ForecasterModule( - forecaster=forecaster, config=config, datastore=datastore + forecaster=forecaster, + config=config, + datastore=datastore, + datastore_boundary=datastore_boundary, ) @@ -106,6 +110,73 @@ def test_normalization_applied_exactly_once(): assert not torch.allclose(norm_init, twice) # not twice +def test_boundary_standardized_when_datastore_provided(): + """Boundary forcing is standardized by ForecasterModule when a + boundary datastore is wired in, using its own forcing mean/std.""" + datastore = DummyDatastore(n_grid_points=100, n_timesteps=20) + datastore_boundary = BoundaryDummyDatastore( + n_grid_points=25, n_timesteps=20 + ) + model = _build_module( + datastore=datastore, datastore_boundary=datastore_boundary + ) + assert model.boundary_mean is not None + assert model.boundary_std is not None + + num_boundary_grid = datastore_boundary.num_grid_points + num_boundary_vars = datastore_boundary.get_num_data_vars("forcing") + window_size = 3 # arbitrary; must match the stacked feature axis below + ar_steps = 2 + + init_states = torch.randn(1, 2, datastore.num_grid_points, 5) + target_states = torch.randn(1, ar_steps, datastore.num_grid_points, 5) + forcing = torch.randn( + 1, ar_steps, datastore.num_grid_points, 2 * window_size + ) + boundary = torch.randn( + 1, ar_steps, num_boundary_grid, num_boundary_vars * window_size + ) + target_times = torch.randint(0, 1000000, (1, ar_steps)) + + _, _, _, norm_boundary, _ = model.on_after_batch_transfer( + (init_states, target_states, forcing, boundary, target_times), 0 + ) + + boundary_mean_tiled = model.boundary_mean.repeat_interleave(window_size) + boundary_std_tiled = model.boundary_std.repeat_interleave(window_size) + expected = (boundary - boundary_mean_tiled) / boundary_std_tiled + assert torch.allclose(norm_boundary, expected) + # Tiled buffers should now be cached. + assert model.boundary_mean_tiled is not None + assert model.boundary_std_tiled is not None + + +def test_boundary_passthrough_when_no_boundary_datastore(): + """Without a boundary datastore, boundary is passed through unchanged + even if the tensor has non-zero last dim.""" + datastore = init_datastore_example("mdp") + model = _build_module(datastore) + assert model.boundary_mean is None + + num_state = datastore.get_num_data_vars("state") + boundary = torch.randn(1, 2, datastore.num_grid_points, 3) + init_states = torch.randn(1, 2, datastore.num_grid_points, num_state) + target_states = torch.randn(1, 2, datastore.num_grid_points, num_state) + forcing = torch.randn( + 1, + 2, + datastore.num_grid_points, + datastore.get_num_data_vars("forcing") + * (NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1), + ) + target_times = torch.randint(0, 1000000, (1, 2)) + + _, _, _, norm_boundary, _ = model.on_after_batch_transfer( + (init_states, target_states, forcing, boundary, target_times), 0 + ) + assert torch.equal(norm_boundary, boundary) + + def test_safe_std_clamps_near_zero(): """Regression test for https://github.com/mllam/neural-lam/issues/136: near-zero std is clamped to machine epsilon (with a warning) so From a543ce9e7a73e8cde193cbd0ed1b91a0aae5b7d7 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 28 May 2026 21:22:40 +0200 Subject: [PATCH 16/35] fix(dataset): anchor forecast boundary on model init time The forecast-boundary path selected its analysis_time by pad-matching the first target time, which can pick a boundary forecast launched after model init - unavailable operationally. Anchor on the model init time (state_times[init_steps - 1], strictly before) instead, matching the research branch, and assert the per-step boundary valid time never runs ahead of the target. Co-Authored-By: Claude Opus 4.8 --- neural_lam/weather_dataset.py | 31 ++++++++------ tests/test_time_slicing.py | 78 ++++++++++++++++++++++++++++++++--- 2 files changed, 91 insertions(+), 18 deletions(-) diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 8ef2e986a..555de5ec1 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -349,29 +349,31 @@ def _window_forcing_in_time( "forecast_step must be supplied when forcing/boundary " "is in forecast mode." ) - # Choose a single analysis_time for this sample. We want it to - # be the latest analysis_time strictly less than the first - # target time (so the first target sits at lead >= 1, leaving - # room for past-window lookups), shifted further back if a - # larger num_past_steps requires more lead headroom. + # Choose a single analysis_time (launch) for this sample. We + # anchor on the model init time (the last input state), not the + # first target, so we never select a boundary forecast launched + # after init - that forecast would be unavailable operationally. + # A launch exactly at init is also rejected (strictly before), + # then shifted further back if a larger num_past_steps requires + # more lead headroom. + model_init_time = state_times[init_steps - 1].values first_target_time = state_times[init_steps].values analysis_index = da_forcing.analysis_time.get_index("analysis_time") forcing_at_idx = analysis_index.get_indexer( - [first_target_time], method="pad" + [model_init_time], method="pad" )[0] if forcing_at_idx < 0: raise ValueError( - "Boundary/forcing analysis times start after the first " - f"target time ({first_target_time})." + "Boundary/forcing analysis times start after the model " + f"init time ({model_init_time})." ) forcing_at = da_forcing.analysis_time[forcing_at_idx] - if first_target_time == forcing_at.values: + if model_init_time == forcing_at.values: if forcing_at_idx == 0: raise ValueError( - "No earlier boundary/forcing analysis time is " - "available to provide lead headroom for target " - f"time {first_target_time}." + "No boundary/forcing analysis time strictly before " + f"the model init time ({model_init_time}) is available." ) forcing_at_idx -= 1 forcing_at = da_forcing.analysis_time[forcing_at_idx] @@ -396,6 +398,11 @@ def _window_forcing_in_time( lead = int( np.floor((target_time - forcing_at.values) / forecast_step) ) + center_time = forcing_at.values + lead * forecast_step + assert center_time <= target_time, ( + "Boundary forecast valid time runs ahead of the interior " + f"target time ({center_time} > {target_time})." + ) window_start = lead - num_past_steps window_end = lead + num_future_steps + 1 diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index 4575ebf0d..1f9c851c2 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -488,18 +488,84 @@ def test_analysis_interior_with_forecast_boundary(): _, _, _, boundary, _ = [t.numpy() for t in dataset[0]] # Sample idx=0: state slice = interior[0:4] = times 2020-01-05..08. - # Targets are 2020-01-07 and 2020-01-08 (first_target=07). Boundary - # analysis_time pad-pick for 07 = idx 3 (07); equals first_target so - # decrement to idx 2 (06). lead_at_first_target = (07-06)/1d = 1, + # Model init is the last input state 2020-01-06; targets are 07 and 08. + # Boundary analysis_time pad-pick for the init 06 = idx 2 (06); equals + # init so decrement to idx 1 (05). lead_at_first_target = (07-05)/1d = 2, # which already covers num_past=1, so no further shift. Window at - # target 07: lead 1, [0..2]. Window at target 08: lead 2, [1..3]. + # target 07: lead 2, [1..3]. Window at target 08: lead 3, [2..4]. + expected_analysis_idx = 1 + assert boundary.shape == (2, 1, 3) + np.testing.assert_array_equal( + boundary[0, 0, :], boundary_values[expected_analysis_idx, 1:4] + ) + np.testing.assert_array_equal( + boundary[1, 0, :], boundary_values[expected_analysis_idx, 2:5] + ) + + +def test_forecast_boundary_anchors_on_init_not_target(): + """A boundary forecast launched after model init (between the last + input state and the first target) must not be selected - operationally + it would be unavailable. The analysis_time is anchored on the model + init time, so the latest launch at or before init is used instead.""" + # Interior analysis, 2h step. Sample idx=0 state = 00,02,04,06: + # model init = 02, first target = 04, second target = 06. + interior_times = np.datetime64("2020-01-01T00") + np.arange( + 8 + ) * np.timedelta64(2, "h") + interior_values = np.arange(8, dtype=float) + interior_datastore = SinglePointDummyDatastore( + state_data=interior_values, + forcing_data=interior_values, + time_values=interior_times, + is_forecast=False, + step_length=timedelta(hours=2), + ) + + # Boundary launches at odd hours (2019-12-31T21, 23, 01, 03, ...), + # spanning wide enough that no interior cropping is triggered. Launch + # 01 (idx 2) is the latest <= init (02); launch 03 (idx 3) sits + # strictly between init (02) and the first target (04). The buggy + # target-time anchor would pick 03 (a future launch); the fixed + # init-time anchor picks 01. + n_analysis = 9 + n_leads = 16 + boundary_analysis = np.datetime64("2019-12-31T21") + np.arange( + n_analysis + ) * np.timedelta64(2, "h") + boundary_leads = np.arange(n_leads) * np.timedelta64(1, "h") + boundary_values = ( + np.arange(n_analysis).reshape(-1, 1) * 1000 + + np.arange(n_leads).reshape(1, -1) * 10 + ).astype(float) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=boundary_values, + time_values=(boundary_analysis, boundary_leads), + is_forecast=True, + step_length=timedelta(hours=1), + ) + + dataset = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=2, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + ) + + _, _, _, boundary, _ = [t.numpy() for t in dataset[0]] + # Launch at 01 = analysis idx 2 (not 03 = idx 3). From 01: target 04 + # is lead (04-01)/1h = 3 -> window [2,5); target 06 is lead 5 -> + # window [4,7). expected_analysis_idx = 2 assert boundary.shape == (2, 1, 3) np.testing.assert_array_equal( - boundary[0, 0, :], boundary_values[expected_analysis_idx, 0:3] + boundary[0, 0, :], boundary_values[expected_analysis_idx, 2:5] ) np.testing.assert_array_equal( - boundary[1, 0, :], boundary_values[expected_analysis_idx, 1:4] + boundary[1, 0, :], boundary_values[expected_analysis_idx, 4:7] ) From 685e2575c777f1b179bde024f9a276c37d8d44c0 Mon Sep 17 00:00:00 2001 From: sadamov Date: Sun, 7 Jun 2026 16:48:19 +0200 Subject: [PATCH 17/35] ci: pass --run-slow to pytest so slow tests run in CI Matches the CI invocation introduced in #651 so this branch's eventual ERA5 boundary integration test (and any other @pytest.mark.slow on this branch) is exercised by CI, not just by local --run-slow invocations. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/install-and-test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/install-and-test.yml b/.github/workflows/install-and-test.yml index d9f4e3f6e..21d908b8c 100644 --- a/.github/workflows/install-and-test.yml +++ b/.github/workflows/install-and-test.yml @@ -128,7 +128,7 @@ jobs: - name: Run tests if: matrix.package_manager == 'uv' run: | - pytest -vv -s --doctest-modules + pytest -vv -s --doctest-modules --run-slow - name: Upload test figures if: matrix.package_manager == 'uv' From 78a4d52d731d3ee410012e882a22f30a2657e31e Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 8 Jun 2026 20:29:35 +0200 Subject: [PATCH 18/35] test/ci: drop --run-slow plumbing in favour of pytest -m Per the simplification in #651, the custom --run-slow flag is being dropped in favour of relying on pytest's native -m marker selection. Remove the corresponding pytest_addoption + pytest_collection_modifyitems from conftest and the --run-slow flag from CI on this branch too, so this PR doesn't re-introduce conflicts when #651 lands. The slow marker registration in pyproject.toml stays, ready for use on any future @pytest.mark.slow test on this branch. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/install-and-test.yml | 2 +- tests/conftest.py | 17 ----------------- 2 files changed, 1 insertion(+), 18 deletions(-) diff --git a/.github/workflows/install-and-test.yml b/.github/workflows/install-and-test.yml index 21d908b8c..d9f4e3f6e 100644 --- a/.github/workflows/install-and-test.yml +++ b/.github/workflows/install-and-test.yml @@ -128,7 +128,7 @@ jobs: - name: Run tests if: matrix.package_manager == 'uv' run: | - pytest -vv -s --doctest-modules --run-slow + pytest -vv -s --doctest-modules - name: Upload test figures if: matrix.package_manager == 'uv' diff --git a/tests/conftest.py b/tests/conftest.py index 876acc025..bfae0cf62 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,23 +23,6 @@ os.environ["WANDB_MODE"] = "disabled" -def pytest_addoption(parser): - parser.addoption( - "--run-slow", - action="store_true", - default=False, - help="Run tests marked as slow", - ) - - -def pytest_collection_modifyitems(config, items): - if not config.getoption("--run-slow"): - skip_slow = pytest.mark.skip(reason="need --run-slow option to run") - for item in items: - if "slow" in item.keywords: - item.add_marker(skip_slow) - - @pytest.fixture(autouse=True) def ensure_rank_zero(monkeypatch): """Ensure rank_zero_only.rank == 0 so @rank_zero_only-decorated functions From 6e0a2c6d9a8bdee0013ea7334da3250dc90f7b13 Mon Sep 17 00:00:00 2001 From: sadamov Date: Mon, 15 Jun 2026 15:12:05 +0200 Subject: [PATCH 19/35] refactor(config): use named datastores mapping with role implied by category Replace the `datastore` + optional `datastore_boundary` config keys with a single `datastores` mapping keyed by a user-chosen name. A datastore's role is now implied by the categories it provides rather than by a dedicated key: a datastore with `state` data is the interior (model input and output), one without `state` data is used for input only (e.g. boundary forcing). At least one datastore must provide `state`; at most one boundary datastore is currently supported. `load_config_and_datastore` classifies the datastores and keeps returning the same `(config, interior, boundary)` tuple, so downstream callers are unchanged. Implements the design agreed in the discussion on #652 (leifdenby's named `datastores` proposal, endorsed by joeloskarsson). refs #108 refs #652 Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 4 +- README.md | 37 ++++++----- neural_lam/config.py | 62 +++++++++++++------ .../mdp/danra_100m_winds/config.yaml | 7 ++- .../era5_1000hPa_danra_100m_winds/config.yaml | 13 ++-- tests/test_checkpoint.py | 10 +-- tests/test_clamping.py | 8 ++- tests/test_config.py | 26 +++++--- tests/test_datasets.py | 8 ++- tests/test_gnn_layers.py | 10 +-- tests/test_gpu_normalization.py | 8 ++- tests/test_plotting.py | 30 +++++---- tests/test_prediction_model_classes.py | 16 +++-- tests/test_training.py | 8 ++- 14 files changed, 154 insertions(+), 93 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a1dc5875f..bf786c382 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional boundary datastore support: `NeuralLAMConfig` accepts a `datastore_boundary` field, `WeatherDataset` loads boundary forcing from a separate domain, and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Add optional boundary datastore support: `NeuralLAMConfig` now takes a named `datastores` mapping where each datastore's role is implied by the categories it provides (a datastore with `state` data is the interior input/output domain, one without `state` data is used for input only, e.g. boundary forcing from a separate domain). `WeatherDataset` loads boundary forcing from such a datastore and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov - Add `PropagationNet` GNN layer that incentivises directional message propagation from sender to receiver nodes, and expose it alongside @@ -25,6 +25,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:`. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov + - Move data normalization from CPU (`WeatherDataset`) to GPU (`ForecasterModule.on_after_batch_transfer`) for improved performance and multi-GPU compatibility. `WeatherDataset` / `WeatherDataModule` no longer diff --git a/README.md b/README.md index 12e0cea27..cbee13f82 100644 --- a/README.md +++ b/README.md @@ -145,12 +145,13 @@ data/ And the content of `config.yaml` could in this case look like: ```yaml -datastore: - kind: mdp - config_path: danra.datastore.yaml -datastore_boundary: # optional, for boundary forcing from a separate domain - kind: mdp - config_path: era5_boundary.datastore.yaml +datastores: + danra: + kind: mdp + config_path: danra.datastore.yaml + era5_boundary: # optional; no `state` data, so used for input (boundary) only + kind: mdp + config_path: era5_boundary.datastore.yaml training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting @@ -169,15 +170,18 @@ training: For now the neural-lam config only defines few things: -1. The kind of datastore and the path to its config -2. (Optional) A boundary datastore (`datastore_boundary`) providing boundary - forcing from a separate domain (e.g. ERA5 for a LAM domain). When set, the - boundary forcing is windowed and included as an additional tensor in each - training sample. -3. The weighting of different features in +1. A named mapping of datastores (`datastores`), each giving the kind of + datastore and the path to its config. The role of each datastore is implied + by the categories of data it provides: a datastore that contains `state` + data is used for both model input and output (the interior domain), while a + datastore without `state` data is used for input only (e.g. boundary forcing + from a separate domain such as ERA5 for a LAM domain). At least one datastore + must provide `state` data. When a boundary datastore is present its forcing + is windowed and included as an additional tensor in each training sample. +2. The weighting of different features in the loss function. If you don't define the state feature weighting it will default to weighting all features equally. -4. Valid numerical range for output of each feature. The numerical range of all features default to $]-\infty, \infty[$. +3. Valid numerical range for output of each feature. The numerical range of all features default to $]-\infty, \infty[$. (This example is taken from the `tests/datastore_examples/mdp` directory.) @@ -378,9 +382,10 @@ Which you can then use in a neural-lam configuration file like this: ```yaml # config.yaml -datastore: - kind: npyfilesmeps - config_path: meps.datastore.yaml +datastores: + meps: + kind: npyfilesmeps + config_path: meps.datastore.yaml training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting diff --git a/neural_lam/config.py b/neural_lam/config.py index e74abed0c..00098494e 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -112,14 +112,19 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): Attributes ---------- - datastore : DatastoreSelection - The configuration for the datastore to use. + datastores : Dict[str, DatastoreSelection] + Mapping from a user-chosen datastore name to its selection config. The + role of each datastore is implied by the categories of data it + provides rather than by a dedicated config key: a datastore that + contains `state` data is used for both model input and output (the + interior domain), while a datastore without `state` data is used for + input only (e.g. boundary forcing from a separate domain). At least + one datastore must provide `state` data. training : TrainingConfig The configuration for training the model. """ - datastore: DatastoreSelection - datastore_boundary: Union[DatastoreSelection, None] = None + datastores: Dict[str, DatastoreSelection] training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig) class _(dataclass_wizard.JSONWizard.Meta): @@ -173,8 +178,9 @@ def load_config_and_datastore( Returns ------- tuple[NeuralLAMConfig, datastore, datastore_boundary] - The Neural-LAM configuration, the loaded (interior) datastore, - and the boundary datastore (or None if not configured). + The Neural-LAM configuration, the loaded interior datastore (the one + providing `state` data), and the boundary datastore (the one without + `state` data, or None if no such datastore is configured). """ try: config = NeuralLAMConfig.from_yaml_file(config_path) @@ -183,22 +189,38 @@ def load_config_and_datastore( "There was an error loading the configuration file at " f"{config_path}. " ) from ex - # datastore config is assumed to be relative to the config file - datastore_config_path = ( - Path(config_path).parent / config.datastore.config_path - ) - datastore = init_datastore( - datastore_kind=config.datastore.kind, config_path=datastore_config_path - ) - datastore_boundary = None - if config.datastore_boundary is not None: - datastore_boundary_config_path = ( - Path(config_path).parent / config.datastore_boundary.config_path + # datastore configs are assumed to be relative to the config file. The + # role of each datastore is implied by the categories of data it provides: + # a datastore with `state` data is the interior (input and output), one + # without `state` data is used for input only (e.g. boundary forcing). + config_dir = Path(config_path).parent + interior_datastores = {} + boundary_datastores = {} + for name, selection in config.datastores.items(): + datastore = init_datastore( + datastore_kind=selection.kind, + config_path=config_dir / selection.config_path, ) - datastore_boundary = init_datastore( - datastore_kind=config.datastore_boundary.kind, - config_path=datastore_boundary_config_path, + if datastore.get_num_data_vars(category="state") > 0: + interior_datastores[name] = datastore + else: + boundary_datastores[name] = datastore + + if len(interior_datastores) != 1: + raise InvalidConfigError( + "Exactly one datastore must provide `state` data (the interior " + f"domain), but {len(interior_datastores)} were found in " + f"{config_path}: {sorted(interior_datastores)}." ) + if len(boundary_datastores) > 1: + raise InvalidConfigError( + "At most one boundary datastore (a datastore without `state` " + f"data) is currently supported, but {len(boundary_datastores)} " + f"were found in {config_path}: {sorted(boundary_datastores)}." + ) + + (datastore,) = interior_datastores.values() + datastore_boundary = next(iter(boundary_datastores.values()), None) return config, datastore, datastore_boundary 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/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml index a158bee3b..750ebcce2 100644 --- a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml @@ -1,9 +1,10 @@ -datastore: - kind: mdp - config_path: danra.datastore.yaml -datastore_boundary: - kind: mdp - config_path: era5.datastore.yaml +datastores: + danra: + kind: mdp + config_path: danra.datastore.yaml + era5: + kind: mdp + config_path: era5.datastore.yaml training: state_feature_weighting: __config_class__: ManualStateFeatureWeighting diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148b..45bdda47d 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={ + "main": 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..95bd3919d 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={ + "main": 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..1b3680245 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={ + "main": 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: + main: + kind: mdp + config_path: "" """ default_config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection(kind="mdp", config_path=""), + datastores={ + "main": nlconfig.DatastoreSelection(kind="mdp", config_path="") + }, training=nlconfig.TrainingConfig( state_feature_weighting=nlconfig.UniformFeatureWeighting() ), ) yaml_training_manual_weights = """ -datastore: - kind: mdp - config_path: "" +datastores: + main: + 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={ + "main": 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 5aa437ad5..a877dea0a 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -215,9 +215,11 @@ def _create_graph(): _create_graph() config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + } ) dataset = WeatherDataset(datastore=datastore, split=split, ar_steps=2) diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 166a9b549..acab63611 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={ + "main": 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 5fd018efd..cf4772c8a 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -22,9 +22,11 @@ def forward(self, prev_state, prev_prev_state, forcing): def _build_module(datastore, datastore_boundary=None): config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + } ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) forecaster = ARForecaster(predictor, datastore) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index ef69abe5f..4915bebc7 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -442,10 +442,12 @@ class ModelArgs: # Create config. config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, - config_path=datastore.root_path, - ), + datastores={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ) + }, ) # Create model @@ -715,10 +717,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={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ) + }, ) model = _build_metrics_watch_module(datastore, config) @@ -774,10 +778,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={ + "main": 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..45860a7f1 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={ + "main": 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={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + } ) # First-party diff --git a/tests/test_training.py b/tests/test_training.py index 141681fc0..45da78d76 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -99,9 +99,11 @@ def run_simple_training( ) config = nlconfig.NeuralLAMConfig( - datastore=nlconfig.DatastoreSelection( - kind=datastore.SHORT_NAME, config_path=datastore.root_path - ) + datastores={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + } ) # Build predictor and forecaster externally, then inject into From f8d43ff3d58b9dec06112562470f3e728ee9f45d Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 16 Jun 2026 13:03:05 +0200 Subject: [PATCH 20/35] docs(pytest): correct slow-marker description to run-by-default The `slow` marker registration still described the markers as "deselected by default" (a leftover from the removed `--run-slow` flag). Slow tests actually run by default and are skipped with `-m "not slow"`, matching the behaviour agreed in #651. Description-only change; no behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 113329e13..b3a05c563 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,7 +160,7 @@ exclude = ["tests", "docs", "build"] [tool.pytest.ini_options] markers = [ - "slow: marks tests as slow (deselected by default, run with -m slow)", + "slow: marks tests as slow (run by default; skip with -m 'not slow')", ] [build-system] From 720530fd93f9acbfebf1bc94c548c8caef5db66e Mon Sep 17 00:00:00 2001 From: sadamov Date: Tue, 16 Jun 2026 20:59:19 +0200 Subject: [PATCH 21/35] feat(config): raise migration error for legacy datastore key Loading a config that still uses the old top-level `datastore:` key now raises a clear InvalidConfigError pointing to the named `datastores:` mapping. Add a regression test and note the behaviour in the CHANGELOG. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 2 +- neural_lam/config.py | 18 ++++++++++++++++++ tests/test_config.py | 9 +++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c9e17d224..2d75438e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed -- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:`. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:`; loading a config that still uses the old `datastore:` key now raises a clear migration error. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov - Move data normalization from CPU (`WeatherDataset`) to GPU (`ForecasterModule.on_after_batch_transfer`) for improved performance and diff --git a/neural_lam/config.py b/neural_lam/config.py index 36fe77c8a..2133b8770 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -7,6 +7,7 @@ # Third-party import dataclass_wizard +import yaml # Local from .datastore import ( @@ -200,6 +201,23 @@ def load_config_and_datastore( providing `state` data), and the boundary datastore (the one without `state` data, or None if no such datastore is configured). """ + with open(config_path, encoding="utf-8") as f: + raw_config = yaml.safe_load(f) + if isinstance(raw_config, dict) and ( + "datastore" in raw_config and "datastores" not in raw_config + ): + raise InvalidConfigError( + "The `datastore:` config key has been replaced by a named " + "`datastores:` mapping (the role of each datastore is now implied " + "by the categories it provides). Move your datastore under a " + "named entry, e.g.:\n" + "datastores:\n" + " :\n" + " kind: ...\n" + " config_path: ...\n" + "See the README and CHANGELOG for details." + ) + try: config = NeuralLAMConfig.from_yaml_file(config_path) except dataclass_wizard.errors.UnknownJSONKey as ex: diff --git a/tests/test_config.py b/tests/test_config.py index 1b3680245..096c15928 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -78,3 +78,12 @@ def test_config_serialization(state_weighting_config): def test_config_load_from_yaml(yaml_str, config_expected): c = nlconfig.NeuralLAMConfig.from_yaml(yaml_str) assert c == config_expected + + +def test_legacy_datastore_key_raises_migration_error(tmp_path): + config_path = tmp_path / "config.yaml" + config_path.write_text( + "datastore:\n kind: mdp\n config_path: ''\n", encoding="utf-8" + ) + with pytest.raises(nlconfig.InvalidConfigError, match="datastores:"): + nlconfig.load_config_and_datastore(str(config_path)) From 22310b6866c57071078b37715b76d7f3eea0eac8 Mon Sep 17 00:00:00 2001 From: sadamov Date: Wed, 17 Jun 2026 03:54:11 +0200 Subject: [PATCH 22/35] linting formatting and detailed review --- CHANGELOG.md | 4 +- README.md | 7 +- neural_lam/config.py | 13 ++- neural_lam/datastore/mdp.py | 13 +++ .../compute_standardization_stats.py | 34 +++--- neural_lam/datastore/npyfilesmeps/store.py | 8 +- neural_lam/utils.py | 6 +- neural_lam/vis.py | 6 +- neural_lam/weather_dataset.py | 64 +++++------ .../era5_1000hPa_danra_100m_winds/config.yaml | 9 ++ .../danra.datastore.yaml | 100 +----------------- tests/test_datasets.py | 40 ++++++- tests/test_training.py | 23 +++- 13 files changed, 155 insertions(+), 172 deletions(-) mode change 100644 => 120000 tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d75438e8..83c66415c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional boundary datastore support: `NeuralLAMConfig` now takes a named `datastores` mapping where each datastore's role is implied by the categories it provides (a datastore with `state` data is the interior input/output domain, one without `state` data is used for input only, e.g. boundary forcing from a separate domain). `WeatherDataset` loads boundary forcing from such a datastore and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores (forcing + static, no state variables). [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Add optional boundary datastore support: `NeuralLAMConfig` now takes a named `datastores` dict. `WeatherDataset` loads boundary forcing from such a datastore and `__getitem__` returns a 5-tuple `(init_states, target_states, forcing, boundary, target_times)`. New CLI args `--num_past_boundary_steps` / `--num_future_boundary_steps` control the boundary forcing window. `MDPDatastore` and `NpyFilesDatastoreMEPS` now support boundary-only datastores. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov - Add `PropagationNet` GNN layer that incentivises directional message propagation from sender to receiver nodes, and expose it alongside @@ -27,7 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed -- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:`; loading a config that still uses the old `datastore:` key now raises a clear migration error. [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov +- Replace the single `datastore` (and optional `datastore_boundary`) keys in the neural-lam config with a named `datastores` mapping. Each datastore's role is now implied by the categories it provides rather than by a dedicated config key. Existing configs must move their datastore under a named entry in `datastores:` [\#635](https://github.com/mllam/neural-lam/pull/635) @sadamov - Move data normalization from CPU (`WeatherDataset`) to GPU (`ForecasterModule.on_after_batch_transfer`) for improved performance and diff --git a/README.md b/README.md index 9051f6c63..5bfcfb7bc 100644 --- a/README.md +++ b/README.md @@ -184,7 +184,7 @@ For now the neural-lam config only defines few things: by the categories of data it provides: a datastore that contains `state` data is used for both model input and output (the interior domain), while a datastore without `state` data is used for input only (e.g. boundary forcing - from a separate domain such as ERA5 for a LAM domain). At least one datastore + from a separate domain such as ERA5 for a LAM domain). Exactly one datastore must provide `state` data. When a boundary datastore is present its forcing is windowed and included as an additional tensor in each training sample. 2. The weighting of different features in @@ -596,13 +596,18 @@ Canonical dimension names used in tensor shape annotations throughout the codeba - `B` - batch size - `pred_steps` - number of autoregressive prediction steps +- `num_times` - number of time steps along the time axis of a raw or batched timeseries (a trailing `'`, e.g. `num_times'`, denotes a pre-subsampling / pre-differencing variant) - `num_grid_nodes` - number of nodes in the flattened spatial grid +- `num_boundary_grid_nodes` - number of nodes in the flattened boundary spatial grid - `num_mesh_nodes` - number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l` - `num_state_vars` - number of atmospheric state variables - `num_forcing_vars` - number of forcing input variables +- `num_windowed_forcing_vars` - forcing variables stacked over the past/future forcing window +- `num_windowed_boundary_vars` - boundary forcing variables stacked over the past/future boundary window - `num_variables` - generic variable dimension used in metric functions - `hidden_dim` - internal hidden representation size in GNN layers and MLPs - `input_dim` - input feature dimensionality to a layer before transformation +- `d_mesh_static` - number of static features per mesh node - `num_edges` - number of edges in a graph (g2m, m2g, same-level, up, down) - `num_send` - number of sender nodes in a message-passing step - `num_rec` - number of receiver nodes in a message-passing step diff --git a/neural_lam/config.py b/neural_lam/config.py index 2133b8770..fd33073b3 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -134,8 +134,10 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): provides rather than by a dedicated config key: a datastore that contains `state` data is used for both model input and output (the interior domain), while a datastore without `state` data is used for - input only (e.g. boundary forcing from a separate domain). At least - one datastore must provide `state` data. + input only (e.g. boundary forcing from a separate domain). Exactly + one datastore must provide `state` data, and at most one datastore + may omit it (the boundary); both constraints are enforced in + :func:`load_config_and_datastore`. training : TrainingConfig Configuration for training the model, including loss function and feature-weighting strategy. Defaults to ``TrainingConfig()``. @@ -200,6 +202,13 @@ def load_config_and_datastore( The Neural-LAM configuration, the loaded interior datastore (the one providing `state` data), and the boundary datastore (the one without `state` data, or None if no such datastore is configured). + + Raises + ------ + InvalidConfigError + If not exactly one datastore provides `state` data, or if more than + one datastore omits it (only a single boundary datastore is currently + supported). """ with open(config_path, encoding="utf-8") as f: raw_config = yaml.safe_load(f) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 9876c4a72..8612d20d3 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -74,6 +74,19 @@ def __init__( ".yaml", ".zarr" ) + # `domain_cropping.interior_dataset_config_path` references a sibling + # config and is written relative to this config file, but + # mllam-data-prep resolves it against the process CWD. Make only that + # path absolute so cropped datastores build from any directory while + # all other (CWD-relative) input paths keep their default behaviour. + domain_cropping = self._config.output.domain_cropping + if domain_cropping is not None: + interior_path = Path(domain_cropping.interior_dataset_config_path) + if not interior_path.is_absolute(): + domain_cropping.interior_dataset_config_path = str( + self._root_path / interior_path + ) + _ds = None if reuse_existing and fp_ds.exists(): # check that the zarr directory is newer than the config file diff --git a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py index 9e694f141..82da7dde9 100644 --- a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py +++ b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py @@ -176,13 +176,13 @@ def save_stats( """ means = ( torch.stack(means) if len(means) > 1 else means[0] - ) # (B, d_features,) + ) # (B, num_state_vars,) squares = ( torch.stack(squares) if len(squares) > 1 else squares[0] - ) # (B, d_features,) - mean = torch.mean(means, dim=0) # (d_features,) - second_moment = torch.mean(squares, dim=0) # (d_features,) - std = torch.sqrt(second_moment - mean**2) # (d_features,) + ) # (B, num_state_vars,) + mean = torch.mean(means, dim=0) # (num_state_vars,) + second_moment = torch.mean(squares, dim=0) # (num_state_vars,) + std = torch.sqrt(second_moment - mean**2) # (num_state_vars,) print( f"Saving {filename_prefix} mean and std.-dev. to " f"{filename_prefix}_mean.pt and {filename_prefix}_std.pt" @@ -289,15 +289,15 @@ def main( target_batch.to(device), forcing_batch.to(device), ) - # (B, N_t, num_grid_nodes, d_features) + # (B, num_times, num_grid_nodes, num_state_vars) batch = torch.cat((init_batch, target_batch), dim=1) # Flux at 1st windowed position is index 0 in forcing flux_batch = forcing_batch[:, :, :, 0] - # (B, d_features,) + # (B, num_state_vars,) means.append(torch.mean(batch, dim=(1, 2)).cpu()) squares.append( torch.mean(batch**2, dim=(1, 2)).cpu() - ) # (B, d_features,) + ) # (B, num_state_vars,) flux_means.append(torch.mean(flux_batch).cpu()) # (,) flux_squares.append(torch.mean(flux_batch**2).cpu()) # (,) @@ -342,8 +342,8 @@ def main( ) ] else: - means = [torch.cat(means, dim=0)] # (B, d_features,) - squares = [torch.cat(squares, dim=0)] # (B, d_features,) + means = [torch.cat(means, dim=0)] # (B, num_state_vars,) + squares = [torch.cat(squares, dim=0)] # (B, num_state_vars,) flux_means = [torch.tensor(flux_means)] # (B,) flux_squares = [torch.tensor(flux_squares)] # (B,) @@ -417,7 +417,7 @@ def main( ) init_batch = (init_batch - state_mean) / state_std target_batch = (target_batch - state_mean) / state_std - # (B, N_t', num_grid_nodes, num_state_vars) + # (B, num_times', num_grid_nodes, num_state_vars) batch = torch.cat((init_batch, target_batch), dim=1) # Note: batch contains only 1h-steps stepped_batch = torch.cat( @@ -427,14 +427,14 @@ def main( ], dim=0, ) - # (B', N_t, num_grid_nodes, d_features), + # (B', num_times, num_grid_nodes, num_state_vars), # B' = step_length*B batch_diffs = stepped_batch[:, 1:] - stepped_batch[:, :-1] - # (B', N_t-1, num_grid_nodes, d_features) + # (B', num_times-1, num_grid_nodes, num_state_vars) diff_means.append(torch.mean(batch_diffs, dim=(1, 2)).cpu()) - # (B', d_features,) + # (B', num_state_vars,) diff_squares.append(torch.mean(batch_diffs**2, dim=(1, 2)).cpu()) - # (B', d_features,) + # (B', num_state_vars,) if distributed and world_size > 1: dist.barrier() @@ -458,8 +458,8 @@ def main( diff_means = [diff_means_gathered[:n_original_windows]] diff_squares = [diff_squares_gathered[:n_original_windows]] - diff_means = [torch.cat(diff_means, dim=0)] # (B', d_features,) - diff_squares = [torch.cat(diff_squares, dim=0)] # (B', d_features,) + diff_means = [torch.cat(diff_means, dim=0)] # (B', num_state_vars,) + diff_squares = [torch.cat(diff_squares, dim=0)] # (B', num_state_vars,) if rank == 0: save_stats(static_dir_path, diff_means, diff_squares, [], [], "diff") diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py index 079ea4aad..69fc85362 100644 --- a/neural_lam/datastore/npyfilesmeps/store.py +++ b/neural_lam/datastore/npyfilesmeps/store.py @@ -140,8 +140,8 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore): └── surface_geopotential.npy For the MEPS dataset: - N_t' = 65 - N_t = 65//subsample_step (= 21 for 3h steps) + num_times' = 65 + num_times = 65//subsample_step (= 21 for 3h steps) dim_y = 268 dim_x = 238 num_grid_nodes = 268x238 = 63784 @@ -149,8 +149,8 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore): num_forcing_vars = 5 For the MEPS reduced dataset: - N_t' = 65 - N_t = 65//subsample_step (= 21 for 3h steps) + num_times' = 65 + num_times = 65//subsample_step (= 21 for 3h steps) dim_y = 134 dim_x = 119 num_grid_nodes = 134x119 = 15946 diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 4e5ec399b..04f3daf26 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -327,13 +327,13 @@ def loads_file(fn: str) -> Any: # Load static node features mesh_static_features = loads_file( "mesh_features.pt" - ) # List of (N_mesh[l], d_mesh_static) + ) # List of (num_mesh_nodes[l], d_mesh_static) # Load edges (edge_index) m2m_edge_index = BufferList( [zero_index_edge_index(ei) for ei in loads_file("m2m_edge_index.pt")], persistent=False, - ) # List of (2, M_m2m[l]) + ) # List of (2, num_edges[l]) g2m_edge_index = loads_file("g2m_edge_index.pt") # (2, num_edges) m2g_edge_index = loads_file("m2g_edge_index.pt") # (2, num_edges) @@ -356,7 +356,7 @@ def loads_file(fn: str) -> Any: hierarchical = n_levels > 1 # Not just single level mesh graph # Load static edge features - # List of (M_m2m[l], input_dim) + # List of (num_edges[l], input_dim) m2m_features = loads_file("m2m_features.pt") g2m_features = loads_file("g2m_features.pt") # (num_edges, input_dim) m2g_features = loads_file("m2g_features.pt") # (num_edges, input_dim) diff --git a/neural_lam/vis.py b/neural_lam/vis.py index 1e3dd415f..08793a7d4 100644 --- a/neural_lam/vis.py +++ b/neural_lam/vis.py @@ -502,11 +502,11 @@ def plot_error_heatmap( unavailable the colorbar label includes "[fallback]". """ errors_np = _to_heatmap_matrix(errors) - d_f, pred_steps = errors_np.shape + num_variables, pred_steps = errors_np.shape step_length = datastore.step_length time_step_int, time_step_unit = utils.get_integer_time(step_length) - layout = _compute_heatmap_layout(n_rows=d_f, n_cols=pred_steps) + layout = _compute_heatmap_layout(n_rows=num_variables, n_cols=pred_steps) color_values_np, colorbar_label, heatmap_cmap = _get_heatmap_color_values( errors_np, datastore, normalization ) @@ -571,7 +571,7 @@ def plot_error_heatmap( f"Lead time ({time_step_unit[0]})", size=layout["tick_label_size"] ) - ax.set_yticks(np.arange(d_f)) + ax.set_yticks(np.arange(num_variables)) ax.set_yticklabels( _get_heatmap_var_labels(datastore=datastore), size=layout["tick_label_size"], diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 411ab37f0..e4b297b5e 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -79,26 +79,8 @@ def __init__( load_single_member: bool = False, ) -> None: """ - Parameters - ---------- - datastore : BaseDatastore - Datastore providing access to state/forcing/static arrays. - split : str, optional - Data split (``"train"``, ``"val"``, or ``"test"``). - Default ``"train"``. - ar_steps : int, optional - Number of autoregressive steps per training sample. Default ``3``. - num_past_forcing_steps : int, optional - Past forcing window length ``i`` so that ``[t-i, ..., t]`` forcings - are concatenated. Default ``1``. - num_future_forcing_steps : int, optional - Future forcing window length ``j`` so that ``[t, ..., t+j]`` - forcings are available. Default ``1``. - load_single_member : bool, optional - If ``False`` and the datastore returns an ensemble of state - realisations, treat each state ensemble member as an independent - sample. If ``True``, only ensemble member 0 is used. Default - ``False``. + Construct a ``WeatherDataset``. See the class docstring for the + constructor parameters. Raises ------ @@ -699,16 +681,19 @@ def __getitem__( init_states : torch.Tensor Initial states, shape ``(2, num_grid_nodes, num_state_vars)``. target_states : torch.Tensor - Target states, shape ``(ar_steps, num_grid_nodes, num_state_vars)``. + Target states, shape + ``(pred_steps, num_grid_nodes, num_state_vars)``. forcing : torch.Tensor - Windowed forcing, shape ``(ar_steps, num_grid_nodes, F)`` where - ``F = num_forcing_vars * (num_past_forcing_steps`` - ``+ num_future_forcing_steps + 1)``. + Windowed forcing, shape + ``(pred_steps, num_grid_nodes, num_windowed_forcing_vars)`` where + ``num_windowed_forcing_vars = num_forcing_vars`` + ``* (num_past_forcing_steps + num_future_forcing_steps + 1)``. boundary : torch.Tensor Windowed boundary forcing, shape - ``(ar_steps, N_boundary_grid, d_windowed_boundary)``. + ``(pred_steps, num_boundary_grid_nodes,`` + ``num_windowed_boundary_vars)``. target_times : torch.Tensor - Times of the target steps, shape ``(ar_steps,)``. + Times of the target steps, shape ``(pred_steps,)``. """ n_samples = len(self) @@ -743,11 +728,12 @@ def __getitem__( forcing = torch.tensor(da_forcing_windowed.values, dtype=tensor_dtype) boundary = torch.tensor(da_boundary_windowed.values, dtype=tensor_dtype) - # init_states: (2, N_grid, d_features) - # target_states: (ar_steps, N_grid, d_features) - # forcing: (ar_steps, N_grid, d_windowed_forcing) - # boundary: (ar_steps, N_boundary_grid, d_windowed_boundary) - # target_times: (ar_steps,) + # init_states: (2, num_grid_nodes, num_state_vars) + # target_states: (pred_steps, num_grid_nodes, num_state_vars) + # forcing: (pred_steps, num_grid_nodes, num_windowed_forcing_vars) + # boundary: (pred_steps, num_boundary_grid_nodes, + # num_windowed_boundary_vars) + # target_times: (pred_steps,) return init_states, target_states, forcing, boundary, target_times @@ -937,14 +923,14 @@ def setup(self, stage: Optional[str] = None) -> None: ``None``, both the training split and the validation/test evaluation splits are prepared. """ - shared_kwargs: dict[str, Any] = dict( - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - num_past_boundary_steps=self.num_past_boundary_steps, - num_future_boundary_steps=self.num_future_boundary_steps, - datastore_boundary=self._datastore_boundary, - load_single_member=self.load_single_member, - ) + shared_kwargs: dict[str, Any] = { + "num_past_forcing_steps": self.num_past_forcing_steps, + "num_future_forcing_steps": self.num_future_forcing_steps, + "num_past_boundary_steps": self.num_past_boundary_steps, + "num_future_boundary_steps": self.num_future_boundary_steps, + "datastore_boundary": self._datastore_boundary, + "load_single_member": self.load_single_member, + } if stage == "fit" or stage is None: self.train_dataset = WeatherDataset( datastore=self._datastore, diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml index 750ebcce2..fa4bc8b80 100644 --- a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml @@ -11,3 +11,12 @@ training: weights: u100m: 1.0 v100m: 1.0 + t2m: 1.0 + r2m: 1.0 + output_clamping: + lower: + t2m: 0.0 + r2m: 0 + upper: + r2m: 1.0 + u100m: 100.0 diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml deleted file mode 100644 index 3edf12673..000000000 --- a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml +++ /dev/null @@ -1,99 +0,0 @@ -schema_version: v0.5.0 -dataset_version: v0.1.0 - -output: - variables: - static: [grid_index, static_feature] - state: [time, grid_index, state_feature] - forcing: [time, grid_index, forcing_feature] - coord_ranges: - time: - start: 1990-09-03T00:00 - end: 1990-09-09T00:00 - step: PT3H - chunking: - time: 1 - splitting: - dim: time - splits: - train: - start: 1990-09-03T00:00 - end: 1990-09-06T00:00 - compute_statistics: - ops: [mean, std, diff_mean, diff_std] - dims: [grid_index, time] - val: - start: 1990-09-06T00:00 - end: 1990-09-07T00:00 - test: - start: 1990-09-07T00:00 - end: 1990-09-09T00:00 - -inputs: - danra_height_levels: - path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/height_levels.zarr - dims: [time, x, y, altitude] - variables: - u: - altitude: - values: [100,] - units: m - v: - altitude: - values: [100, ] - units: m - dim_mapping: - time: - method: rename - dim: time - state_feature: - method: stack_variables_by_var_name - dims: [altitude] - name_format: "{var_name}{altitude}m" - grid_index: - method: stack - dims: [x, y] - target_output_variable: state - - danra_surface: - path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr - dims: [time, x, y] - variables: - # use surface incoming shortwave radiation as forcing - - swavr0m - dim_mapping: - time: - method: rename - dim: time - grid_index: - method: stack - dims: [x, y] - forcing_feature: - method: stack_variables_by_var_name - name_format: "{var_name}" - target_output_variable: forcing - - danra_lsm: - path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/lsm.zarr - dims: [x, y] - variables: - - lsm - dim_mapping: - grid_index: - method: stack - dims: [x, y] - static_feature: - method: stack_variables_by_var_name - name_format: "{var_name}" - target_output_variable: static - -extra: - projection: - class_name: LambertConformal - kwargs: - central_longitude: 25.0 - central_latitude: 56.7 - standard_parallels: [56.7, 56.7] - globe: - semimajor_axis: 6367470.0 - semiminor_axis: 6367470.0 diff --git a/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml new file mode 120000 index 000000000..8245a9701 --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/danra.datastore.yaml @@ -0,0 +1 @@ +../danra_100m_winds/danra.datastore.yaml \ No newline at end of file diff --git a/tests/test_datasets.py b/tests/test_datasets.py index a877dea0a..94dbe5667 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -14,7 +14,10 @@ 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_boundary_example, + init_datastore_example, +) from tests.dummy_datastore import ( BoundaryDummyDatastore, DummyDatastore, @@ -606,3 +609,38 @@ def test_boundary_crops_interior_when_window_overflows(): # Boundary spans the same range as interior, so a (past=1, future=1) # window forces 1 step of cropping at each end. assert len(dataset_with_boundary) == len(dataset_no_boundary) - 2 + + +@pytest.mark.slow +def test_boundary_datastore_example_shapes(): + """Build the real MDP interior (DANRA) and ERA5 boundary example + datastores and check WeatherDataset returns a coherent windowed boundary + tensor for a temporally overlapping interior/boundary pair.""" + datastore = init_datastore_example("mdp") + datastore_boundary = init_datastore_boundary_example("mdp") + + ar_steps = 3 + num_past_boundary = 1 + num_future_boundary = 1 + boundary_window = num_past_boundary + num_future_boundary + 1 + + dataset = WeatherDataset( + datastore=datastore, + datastore_boundary=datastore_boundary, + split="train", + ar_steps=ar_steps, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + num_past_boundary_steps=num_past_boundary, + num_future_boundary_steps=num_future_boundary, + ) + + _, _, _, boundary, target_times = dataset[0] + + n_boundary_forcing = datastore_boundary.get_num_data_vars("forcing") + assert boundary.ndim == 3 + assert boundary.shape[0] == ar_steps + assert boundary.shape[1] == datastore_boundary.num_grid_points + assert boundary.shape[2] == n_boundary_forcing * boundary_window + assert target_times.shape == (ar_steps,) + assert not torch.isnan(boundary).any() diff --git a/tests/test_training.py b/tests/test_training.py index 6e8c78d02..a0f979e8e 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -15,7 +15,10 @@ from neural_lam.datastore.base import BaseRegularGridDatastore from neural_lam.models import ForecasterModule from neural_lam.weather_dataset import WeatherDataModule -from tests.conftest import init_datastore_example +from tests.conftest import ( + init_datastore_boundary_example, + init_datastore_example, +) # Model architecture defaults for tests GRAPH = "1level" @@ -32,6 +35,7 @@ def run_simple_training( set_output_std, metrics_watch=None, var_leads_metrics_watch=None, + datastore_boundary=None, ): """ Run one epoch of a simple model training setup using the given datastore. @@ -42,6 +46,8 @@ def run_simple_training( Datastore to load data from for training set_output_std : bool If --output_std should be set during training + datastore_boundary : BaseDatastore, optional + Boundary datastore to load boundary forcing from during training """ if metrics_watch is None: metrics_watch = [] @@ -90,6 +96,7 @@ def run_simple_training( data_module = WeatherDataModule( datastore=datastore, + datastore_boundary=datastore_boundary, ar_steps_train=3, ar_steps_eval=5, batch_size=2, @@ -131,6 +138,7 @@ def run_simple_training( forecaster=forecaster, config=config, datastore=datastore, + datastore_boundary=datastore_boundary, loss="mse", lr=1.0e-3, restore_opt=False, @@ -163,6 +171,19 @@ def test_training_output_std(): run_simple_training(datastore, set_output_std=True) +@pytest.mark.slow +def test_training_with_boundary(): + """One epoch of training with a boundary datastore, exercising boundary + loading and standardization through the full Lightning training loop.""" + datastore = init_datastore_example("mdp") + datastore_boundary = init_datastore_boundary_example("mdp") + run_simple_training( + datastore, + set_output_std=False, + datastore_boundary=datastore_boundary, + ) + + def test_all_gather_cat_single_device(): """ Test that all_gather_cat preserves tensor shape on single-device runs. From 34ee3dfcac6c2eebf2026b54a87eb230ea366f19 Mon Sep 17 00:00:00 2001 From: sadamov Date: Wed, 17 Jun 2026 04:38:03 +0200 Subject: [PATCH 23/35] docs: document state/forcing/static category meaning Add an input/output table for the state/forcing/static categories in the Data section, as requested in #652, so the datastore role-by-category explanation has an explicit definition to reference. Co-authored-by: Claude Opus 4.8 (1M context) --- README.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.md b/README.md index 5bfcfb7bc..2d47b07a0 100644 --- a/README.md +++ b/README.md @@ -226,6 +226,21 @@ the input-data representation is split into two parts: `WeatherDataset` class is also responsible for normalising the values and returning `torch.Tensor`-objects. +Each variable in a datastore is assigned to one of three data *categories*, +which fix whether it is fed to the model as input, predicted as output, or both: + +| Category | Model input | Model output | Description | +|-----------|:-----------:|:------------:|-------------| +| `state` | ✓ | ✓ | Prognostic variables the model both reads and predicts (autoregressed forward in time). | +| `forcing` | ✓ | | Time-varying inputs known in advance (e.g. solar radiation, boundary forcing). | +| `static` | ✓ | | Time-invariant inputs (e.g. orography, land-sea mask). | + +(A `diagnostic` category, predicted but not fed back as input, is not used at +present but fits naturally as an output-only category.) These categories are +also what determine a datastore's role: a datastore that provides `state` data +is the interior domain (model input and output), while one without `state` data +is used for input only, e.g. boundary forcing from a separate domain. + There are currently two different datastores implemented in the codebase: 1. `neural_lam.datastore.MDPDatastore` which represents loading of From 059033dfba57020be96ce622b49ddb3e95ba3b61 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 07:29:13 +0200 Subject: [PATCH 24/35] Drop unimplemented diagnostic category note from README The parenthetical described a `diagnostic` (output-only) category that does not exist in the codebase; tracked separately instead. Co-Authored-By: Claude Opus 4.8 --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index 2d47b07a0..8b888dacc 100644 --- a/README.md +++ b/README.md @@ -235,8 +235,7 @@ which fix whether it is fed to the model as input, predicted as output, or both: | `forcing` | ✓ | | Time-varying inputs known in advance (e.g. solar radiation, boundary forcing). | | `static` | ✓ | | Time-invariant inputs (e.g. orography, land-sea mask). | -(A `diagnostic` category, predicted but not fed back as input, is not used at -present but fits naturally as an output-only category.) These categories are +These categories are also what determine a datastore's role: a datastore that provides `state` data is the interior domain (model input and output), while one without `state` data is used for input only, e.g. boundary forcing from a separate domain. From 8b08abd0e8aec360e5bb4fb9aa8ddff3d3009f0a Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 08:19:38 +0200 Subject: [PATCH 25/35] Document single-boundary-datastore constraint in README Make explicit that exactly one interior and at most one boundary datastore are supported, matching the constraint enforced in load_config_and_datastore. Co-Authored-By: Claude Opus 4.8 --- README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8b888dacc..8229f4b54 100644 --- a/README.md +++ b/README.md @@ -185,8 +185,10 @@ For now the neural-lam config only defines few things: data is used for both model input and output (the interior domain), while a datastore without `state` data is used for input only (e.g. boundary forcing from a separate domain such as ERA5 for a LAM domain). Exactly one datastore - must provide `state` data. When a boundary datastore is present its forcing - is windowed and included as an additional tensor in each training sample. + must provide `state` data, and at most one datastore may omit it, so there is + a single interior and a single (optional) boundary datastore. When a boundary + datastore is present its forcing is windowed and included as an additional + tensor in each training sample. 2. The weighting of different features in the loss function. If you don't define the state feature weighting it will default to weighting all features equally. From 682f792e6e760d2c9961765300e966749cf5dedc Mon Sep 17 00:00:00 2001 From: sadamov <45732287+sadamov@users.noreply.github.com> Date: Thu, 25 Jun 2026 08:34:34 +0200 Subject: [PATCH 26/35] Update neural_lam/train_model.py Co-authored-by: Joel Oskarsson --- neural_lam/train_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index ee76bd86b..02c116871 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -350,14 +350,15 @@ def main(input_args=None): "--num_past_boundary_steps", type=int, default=1, - help="Number of past time steps to use as input for boundary forcing", + help="Number of past time steps to use as input for boundary forcing, " + "when present", ) data_group.add_argument( "--num_future_boundary_steps", type=int, default=1, help="Number of future time steps to use as input for boundary " - "forcing", + "forcing, when present", ) data_group.add_argument( "--load_single_member", From dc0585ddd60b9f650df9294792729939efd2b72e Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 09:21:44 +0200 Subject: [PATCH 27/35] Rename da1/da2 to da_requested/da_available in time helpers Give the two dataarrays in check_time_overlap and crop_time_if_needed role-based names matching the da_ prefix convention used elsewhere: da_requested is the driving timeline, da_available must cover its windows. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 175 ++++++++++++++++++---------------- neural_lam/weather_dataset.py | 8 +- 2 files changed, 98 insertions(+), 85 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 04f3daf26..05f746605 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -858,133 +858,144 @@ def get_time_step(times): def check_time_overlap( - da1: xr.DataArray, - da2: xr.DataArray, - da1_is_forecast: bool = False, - da2_is_forecast: bool = False, + da_requested: xr.DataArray, + da_available: xr.DataArray, + da_requested_is_forecast: bool = False, + da_available_is_forecast: bool = False, num_past_steps: int = 1, num_future_steps: int = 1, ) -> None: - """Check that the time coverage of ``da2`` is wide enough to support - a windowed lookup driven by ``da1`` times with the given past/future - window sizes. + """Check that the time coverage of ``da_available`` is wide enough to + support a windowed lookup driven by ``da_requested`` times with the + given past/future window sizes. Parameters ---------- - da1 : xr.DataArray - Driving dataarray (typically interior state). - da2 : xr.DataArray - Dataarray to validate against (typically boundary forcing). - da1_is_forecast, da2_is_forecast : bool + da_requested : xr.DataArray + Driving dataarray whose times must be supported (typically interior + state). + da_available : xr.DataArray + Dataarray that must cover the requested windows (typically boundary + forcing). + da_requested_is_forecast, da_available_is_forecast : bool Whether each side is in forecast mode (``analysis_time`` + ``elapsed_forecast_duration`` dims) instead of plain ``time``. num_past_steps, num_future_steps : int - Window size around each ``da1`` time, measured in ``da2`` steps. + Window size around each ``da_requested`` time, measured in + ``da_available`` steps. Raises ------ ValueError - If ``da2`` does not cover the required time range. + If ``da_available`` does not cover the required time range. """ - if da1_is_forecast: - times_da1 = da1.analysis_time + if da_requested_is_forecast: + times_requested = da_requested.analysis_time else: - times_da1 = da1.time - time_min_da1 = times_da1.min().values - time_max_da1 = times_da1.max().values + times_requested = da_requested.time + time_min_requested = times_requested.min().values + time_max_requested = times_requested.max().values - if da2_is_forecast: - times_da2 = da2.analysis_time - time_min_da2 = times_da2.min().values - time_max_da2 = times_da2.max().values + if da_available_is_forecast: + times_available = da_available.analysis_time + time_min_available = times_available.min().values + time_max_available = times_available.max().values - time_step_da2 = get_time_step(times_da2.values) - time_step_da1 = get_time_step(times_da1.values) + time_step_available = get_time_step(times_available.values) + time_step_requested = get_time_step(times_requested.values) - analysis_offset = max(time_step_da1, num_past_steps * time_step_da2) - da2_required_time_min = time_min_da1 - analysis_offset - da2_required_time_max = time_max_da1 - analysis_offset + analysis_offset = max( + time_step_requested, num_past_steps * time_step_available + ) + required_time_min = time_min_requested - analysis_offset + required_time_max = time_max_requested - analysis_offset else: - times_da2 = da2.time - time_min_da2 = times_da2.min().values - time_max_da2 = times_da2.max().values - time_step_da2 = get_time_step(times_da2.values) + times_available = da_available.time + time_min_available = times_available.min().values + time_max_available = times_available.max().values + time_step_available = get_time_step(times_available.values) - da2_required_time_min = time_min_da1 - num_past_steps * time_step_da2 - da2_required_time_max = time_max_da1 + num_future_steps * time_step_da2 + required_time_min = ( + time_min_requested - num_past_steps * time_step_available + ) + required_time_max = ( + time_max_requested + num_future_steps * time_step_available + ) - if time_min_da2 > da2_required_time_min: + if time_min_available > required_time_min: raise ValueError( "The second DataArray (e.g. 'boundary forcing') starts too late. " - f"Required start: {da2_required_time_min}, " - f"but DataArray starts at {time_min_da2}." + f"Required start: {required_time_min}, " + f"but DataArray starts at {time_min_available}." ) - if time_max_da2 < da2_required_time_max: + if time_max_available < required_time_max: raise ValueError( "The second DataArray (e.g. 'boundary forcing') ends too early. " - f"Required end: {da2_required_time_max}, " - f"but DataArray ends at {time_max_da2}." + f"Required end: {required_time_max}, " + f"but DataArray ends at {time_max_available}." ) def crop_time_if_needed( - da1: xr.DataArray, - da2: xr.DataArray, - da1_is_forecast: bool = False, - da2_is_forecast: bool = False, + da_requested: xr.DataArray, + da_available: xr.DataArray, + da_requested_is_forecast: bool = False, + da_available_is_forecast: bool = False, num_past_steps: int = 1, num_future_steps: int = 1, ) -> xr.DataArray: - """Trim the leading/trailing times from ``da1`` so that ``da2`` covers - every needed window. If ``check_time_overlap`` already passes, ``da1`` - is returned unchanged. Cropping only applies to analysis-mode ``da1``; - for forecast-mode ``da1`` the input is returned as-is and overlap - failures will surface at sample-construction time. + """Trim the leading/trailing times from ``da_requested`` so that + ``da_available`` covers every needed window. If ``check_time_overlap`` + already passes, ``da_requested`` is returned unchanged. Cropping only + applies to analysis-mode ``da_requested``; for forecast-mode + ``da_requested`` the input is returned as-is and overlap failures will + surface at sample-construction time. Parameters mirror :func:`check_time_overlap`. Returns ------- xr.DataArray - Possibly cropped ``da1``. + Possibly cropped ``da_requested``. """ - if da1 is None or da2 is None: - return da1 + if da_requested is None or da_available is None: + return da_requested try: check_time_overlap( - da1, - da2, - da1_is_forecast, - da2_is_forecast, + da_requested, + da_available, + da_requested_is_forecast, + da_available_is_forecast, num_past_steps, num_future_steps, ) - return da1 + return da_requested except ValueError: - if da1_is_forecast: - # Cropping a forecast da1 by analysis time would change sample - # cardinality silently; leave alignment errors to surface later. - return da1 - - da1_tvals = da1.time.values - if da2_is_forecast: - da2_tvals = da2.analysis_time.values + if da_requested_is_forecast: + # Cropping a forecast da_requested by analysis time would change + # sample cardinality silently; leave alignment errors to surface + # later. + return da_requested + + requested_tvals = da_requested.time.values + if da_available_is_forecast: + available_tvals = da_available.analysis_time.values else: - da2_tvals = da2.time.values - - da2_dt = get_time_step(da2_tvals) - if da2_is_forecast: - da1_dt = get_time_step(da1_tvals) - analysis_offset = max(da1_dt, num_past_steps * da2_dt) - required_min = da2_tvals[0] + analysis_offset - required_max = da2_tvals[-1] + analysis_offset + available_tvals = da_available.time.values + + available_dt = get_time_step(available_tvals) + if da_available_is_forecast: + requested_dt = get_time_step(requested_tvals) + analysis_offset = max(requested_dt, num_past_steps * available_dt) + required_min = available_tvals[0] + analysis_offset + required_max = available_tvals[-1] + analysis_offset else: - required_min = da2_tvals[0] + num_past_steps * da2_dt - required_max = da2_tvals[-1] - num_future_steps * da2_dt + required_min = available_tvals[0] + num_past_steps * available_dt + required_max = available_tvals[-1] - num_future_steps * available_dt - begin_mask = da1_tvals >= required_min + begin_mask = requested_tvals >= required_min if not begin_mask.any(): raise ValueError( "Boundary forcing ends before any interior time satisfies " @@ -992,10 +1003,10 @@ def crop_time_if_needed( ) first_valid_idx = int(begin_mask.argmax()) n_removed_begin = first_valid_idx - if da1_tvals[-1] > required_max: - end_mask = da1_tvals > required_max + if requested_tvals[-1] > required_max: + end_mask = requested_tvals > required_max last_valid_idx_plus_one = int(end_mask.argmax()) - n_removed_end = len(da1_tvals) - last_valid_idx_plus_one + n_removed_end = len(requested_tvals) - last_valid_idx_plus_one else: last_valid_idx_plus_one = None n_removed_end = 0 @@ -1007,5 +1018,7 @@ def crop_time_if_needed( "the end.", level="warning", ) - da1 = da1.isel(time=slice(first_valid_idx, last_valid_idx_plus_one)) - return da1 + da_requested = da_requested.isel( + time=slice(first_valid_idx, last_valid_idx_plus_one) + ) + return da_requested diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index e4b297b5e..ef5ca6ff6 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -160,16 +160,16 @@ def __init__( self.da_state = crop_time_if_needed( self.da_state, self.da_boundary_forcing, - da1_is_forecast=self.datastore.is_forecast, - da2_is_forecast=datastore_boundary.is_forecast, + da_requested_is_forecast=self.datastore.is_forecast, + da_available_is_forecast=datastore_boundary.is_forecast, num_past_steps=self.num_past_boundary_steps, num_future_steps=self.num_future_boundary_steps, ) check_time_overlap( self.da_state, self.da_boundary_forcing, - da1_is_forecast=self.datastore.is_forecast, - da2_is_forecast=datastore_boundary.is_forecast, + da_requested_is_forecast=self.datastore.is_forecast, + da_available_is_forecast=datastore_boundary.is_forecast, num_past_steps=self.num_past_boundary_steps, num_future_steps=self.num_future_boundary_steps, ) From ec714028557da5af9fff42dbacbcca8d6427d955 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 09:41:00 +0200 Subject: [PATCH 28/35] Crop forecast interior along analysis_time to boundary coverage crop_time_if_needed previously returned a forecast da_requested unchanged, deferring alignment errors. Drop whole out-of-coverage analysis_time launches instead (logged like the analysis-mode crop), addressing Joel's review point that forecast cropping is no less safe than analysis cropping. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 23 ++++++--------- tests/test_time_slicing.py | 57 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 14 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 05f746605..7c677cd1b 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -947,10 +947,10 @@ def crop_time_if_needed( ) -> xr.DataArray: """Trim the leading/trailing times from ``da_requested`` so that ``da_available`` covers every needed window. If ``check_time_overlap`` - already passes, ``da_requested`` is returned unchanged. Cropping only - applies to analysis-mode ``da_requested``; for forecast-mode - ``da_requested`` the input is returned as-is and overlap failures will - surface at sample-construction time. + already passes, ``da_requested`` is returned unchanged. A forecast-mode + ``da_requested`` is cropped along ``analysis_time`` (dropping whole + launches), an analysis-mode one along ``time``; either way the removal + is logged. Parameters mirror :func:`check_time_overlap`. @@ -973,13 +973,8 @@ def crop_time_if_needed( ) return da_requested except ValueError: - if da_requested_is_forecast: - # Cropping a forecast da_requested by analysis time would change - # sample cardinality silently; leave alignment errors to surface - # later. - return da_requested - - requested_tvals = da_requested.time.values + crop_dim = "analysis_time" if da_requested_is_forecast else "time" + requested_tvals = da_requested[crop_dim].values if da_available_is_forecast: available_tvals = da_available.analysis_time.values else: @@ -1014,11 +1009,11 @@ def crop_time_if_needed( if n_removed_begin > 0 or n_removed_end > 0: log_on_rank_zero( f"Cropping interior data to align with boundary: removed " - f"{n_removed_begin} steps at start and {n_removed_end} at " - "the end.", + f"{n_removed_begin} {crop_dim} steps at start and " + f"{n_removed_end} at the end.", level="warning", ) da_requested = da_requested.isel( - time=slice(first_valid_idx, last_valid_idx_plus_one) + {crop_dim: slice(first_valid_idx, last_valid_idx_plus_one)} ) return da_requested diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index 1f9c851c2..ca05092b5 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -597,3 +597,60 @@ def test_check_time_overlap_insufficient_raises(): num_past_boundary_steps=20, num_future_boundary_steps=20, ) + + +def test_forecast_interior_cropped_along_analysis_time(): + """A forecast interior whose earliest launches fall outside the boundary + coverage is cropped along ``analysis_time`` (whole launches dropped), so + fewer samples remain and the survivors still build a boundary window.""" + n_analysis = 6 + n_leads = 5 + interior_analysis = np.datetime64("2020-01-01") + np.arange( + n_analysis + ) * np.timedelta64(1, "D") + interior_leads = np.arange(n_leads) * np.timedelta64(1, "D") + interior_values = ( + np.arange(n_analysis).reshape(-1, 1) * 100 + + np.arange(n_leads).reshape(1, -1) + ).astype(float) + interior_datastore = SinglePointDummyDatastore( + state_data=interior_values, + forcing_data=interior_values, + time_values=(interior_analysis, interior_leads), + is_forecast=True, + step_length=timedelta(days=1), + ) + + # Analysis boundary starts only at 2020-01-04, so the launches at + # analysis_time 01-01..01-03 have no boundary coverage and are dropped. + boundary_times = np.datetime64("2020-01-04") + np.arange( + 10 + ) * np.timedelta64(1, "D") + boundary_values = np.arange(300, 310, dtype=float) + boundary_datastore = BoundaryOnlyDummyDatastore( + forcing_data=boundary_values, + time_values=boundary_times, + is_forecast=False, + step_length=timedelta(days=1), + ) + + full = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=None, + ar_steps=2, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + ) + cropped = WeatherDataset( + datastore=interior_datastore, + datastore_boundary=boundary_datastore, + ar_steps=2, + num_past_forcing_steps=0, + num_future_forcing_steps=0, + num_past_boundary_steps=1, + num_future_boundary_steps=1, + ) + + assert len(cropped) < len(full) + _, _, _, boundary, _ = cropped[0] + assert boundary.shape[-1] == 3 From dccb7a6fc0fa76f7fd8ee2c59fe19191ddffffc8 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 10:07:23 +0200 Subject: [PATCH 29/35] Make time-overlap helper messages domain-agnostic Refer to da_requested/da_available in the error and log messages of check_time_overlap and crop_time_if_needed instead of "boundary forcing"/"interior", matching the agnostic signatures. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 7c677cd1b..8859a2a47 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -924,16 +924,16 @@ def check_time_overlap( if time_min_available > required_time_min: raise ValueError( - "The second DataArray (e.g. 'boundary forcing') starts too late. " + "`da_available` starts too late to cover the requested window. " f"Required start: {required_time_min}, " - f"but DataArray starts at {time_min_available}." + f"but `da_available` starts at {time_min_available}." ) if time_max_available < required_time_max: raise ValueError( - "The second DataArray (e.g. 'boundary forcing') ends too early. " + "`da_available` ends too early to cover the requested window. " f"Required end: {required_time_max}, " - f"but DataArray ends at {time_max_available}." + f"but `da_available` ends at {time_max_available}." ) @@ -993,7 +993,7 @@ def crop_time_if_needed( begin_mask = requested_tvals >= required_min if not begin_mask.any(): raise ValueError( - "Boundary forcing ends before any interior time satisfies " + "`da_available` covers no `da_requested` time at or after " f"required_min={required_min}; cannot align." ) first_valid_idx = int(begin_mask.argmax()) @@ -1008,8 +1008,8 @@ def crop_time_if_needed( if n_removed_begin > 0 or n_removed_end > 0: log_on_rank_zero( - f"Cropping interior data to align with boundary: removed " - f"{n_removed_begin} {crop_dim} steps at start and " + f"Cropping `da_requested` to align with `da_available`: " + f"removed {n_removed_begin} {crop_dim} steps at start and " f"{n_removed_end} at the end.", level="warning", ) From ddf882c5ee9b910137ef07a051aceb8b25fdc6d0 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 10:24:56 +0200 Subject: [PATCH 30/35] Use forecast max lead time in boundary coverage bound crop_time_if_needed and check_time_overlap bounded a forecast boundary's reach by a single analysis step, over-cropping the interior for boundaries with few launches but long forecasts. Use the boundary's max elapsed_forecast_duration as its reach past the last launch in required_max, keeping the two functions consistent. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 8859a2a47..5470bef40 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -903,12 +903,15 @@ def check_time_overlap( time_step_available = get_time_step(times_available.values) time_step_requested = get_time_step(times_requested.values) + max_lead = da_available.elapsed_forecast_duration.values.max() analysis_offset = max( time_step_requested, num_past_steps * time_step_available ) required_time_min = time_min_requested - analysis_offset - required_time_max = time_max_requested - analysis_offset + required_time_max = time_max_requested - ( + max_lead - num_future_steps * time_step_available + ) else: times_available = da_available.time time_min_available = times_available.min().values @@ -983,9 +986,12 @@ def crop_time_if_needed( available_dt = get_time_step(available_tvals) if da_available_is_forecast: requested_dt = get_time_step(requested_tvals) + max_lead = da_available.elapsed_forecast_duration.values.max() analysis_offset = max(requested_dt, num_past_steps * available_dt) required_min = available_tvals[0] + analysis_offset - required_max = available_tvals[-1] + analysis_offset + required_max = ( + available_tvals[-1] + max_lead - num_future_steps * available_dt + ) else: required_min = available_tvals[0] + num_past_steps * available_dt required_max = available_tvals[-1] - num_future_steps * available_dt From 441ce6196ca2eb402927fe6f505e48ce77084953 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 10:32:51 +0200 Subject: [PATCH 31/35] Use binary search for the crop cut points Replace the O(T) boolean-mask scan in crop_time_if_needed with two np.searchsorted calls on the already-sorted time axis, dropping the mask allocation for an O(log T) lookup. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 5470bef40..f45968dc1 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -996,21 +996,19 @@ def crop_time_if_needed( required_min = available_tvals[0] + num_past_steps * available_dt required_max = available_tvals[-1] - num_future_steps * available_dt - begin_mask = requested_tvals >= required_min - if not begin_mask.any(): + first_valid_idx = int( + np.searchsorted(requested_tvals, required_min, side="left") + ) + if first_valid_idx == len(requested_tvals): raise ValueError( "`da_available` covers no `da_requested` time at or after " f"required_min={required_min}; cannot align." ) - first_valid_idx = int(begin_mask.argmax()) + last_valid_idx_plus_one = int( + np.searchsorted(requested_tvals, required_max, side="right") + ) n_removed_begin = first_valid_idx - if requested_tvals[-1] > required_max: - end_mask = requested_tvals > required_max - last_valid_idx_plus_one = int(end_mask.argmax()) - n_removed_end = len(requested_tvals) - last_valid_idx_plus_one - else: - last_valid_idx_plus_one = None - n_removed_end = 0 + n_removed_end = len(requested_tvals) - last_valid_idx_plus_one if n_removed_begin > 0 or n_removed_end > 0: log_on_rank_zero( From a75a1ad81cc2f98c8ad909f0e0609dd255f38691 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 10:40:49 +0200 Subject: [PATCH 32/35] Remove dead time-step attributes from WeatherDataset _time_step_state and _forecast_step_forcing were set in __init__ but never read; only _forecast_step_boundary is used. Drop both. Co-Authored-By: Claude Opus 4.8 --- neural_lam/weather_dataset.py | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index ef5ca6ff6..7d2a7ec8d 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -125,23 +125,8 @@ def __init__( else: self.da_boundary_forcing = None - # Within-sample time step for the state series: this is the step - # between consecutive state times that __getitem__ exposes, used - # below to advance the forcing/boundary window across AR steps. - if self.datastore.is_forecast: - self._time_step_state = get_time_step( - self.da_state.elapsed_forecast_duration.values - ) - else: - self._time_step_state = get_time_step(self.da_state.time.values) - - # Forecast lead-time step for forcing/boundary, only meaningful when - # the corresponding datastore is in forecast mode. - self._forecast_step_forcing = None - if self.da_forcing is not None and self.datastore.is_forecast: - self._forecast_step_forcing = get_time_step( - self.da_forcing.elapsed_forecast_duration.values - ) + # Forecast lead-time step for the boundary, only meaningful when the + # boundary datastore is in forecast mode. self._forecast_step_boundary = None if self.datastore_boundary is not None: datastore_boundary = self.datastore_boundary From 8a2dc128f7a1f3ca3398d7b77ffe5a1e6a870786 Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 11:05:22 +0200 Subject: [PATCH 33/35] Make crop_time_if_needed the single coverage source of truth crop_time_if_needed now raises whenever the valid window is empty, so a successful crop always leaves da_state within boundary coverage. Drop the redundant post-crop check_time_overlap call (and its import) from WeatherDataset.__init__, and rename the now-misnamed coverage test. Co-Authored-By: Claude Opus 4.8 --- neural_lam/utils.py | 10 +++++----- neural_lam/weather_dataset.py | 18 ++---------------- tests/test_time_slicing.py | 4 ++-- 3 files changed, 9 insertions(+), 23 deletions(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index f45968dc1..d66b4ca60 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -999,14 +999,14 @@ def crop_time_if_needed( first_valid_idx = int( np.searchsorted(requested_tvals, required_min, side="left") ) - if first_valid_idx == len(requested_tvals): - raise ValueError( - "`da_available` covers no `da_requested` time at or after " - f"required_min={required_min}; cannot align." - ) last_valid_idx_plus_one = int( np.searchsorted(requested_tvals, required_max, side="right") ) + if first_valid_idx >= last_valid_idx_plus_one: + raise ValueError( + "`da_available` covers no `da_requested` time in " + f"[{required_min}, {required_max}]; cannot align." + ) n_removed_begin = first_valid_idx n_removed_end = len(requested_tvals) - last_valid_idx_plus_one diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 7d2a7ec8d..4df53d06f 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -13,11 +13,7 @@ # First-party from neural_lam.datastore.base import BaseDatastore -from neural_lam.utils import ( - check_time_overlap, - crop_time_if_needed, - get_time_step, -) +from neural_lam.utils import crop_time_if_needed, get_time_step class WeatherDataset(torch.utils.data.Dataset): @@ -138,9 +134,7 @@ def __init__( self.da_boundary_forcing.elapsed_forecast_duration.values ) - # Validate that the boundary covers the windows we will request, - # and if necessary crop the analysis-mode interior so that the - # very first/last samples don't fall outside boundary coverage. + # Crop the interior so the first/last samples stay within boundary. if self.da_boundary_forcing is not None: self.da_state = crop_time_if_needed( self.da_state, @@ -150,14 +144,6 @@ def __init__( num_past_steps=self.num_past_boundary_steps, num_future_steps=self.num_future_boundary_steps, ) - check_time_overlap( - self.da_state, - self.da_boundary_forcing, - da_requested_is_forecast=self.datastore.is_forecast, - da_available_is_forecast=datastore_boundary.is_forecast, - num_past_steps=self.num_past_boundary_steps, - num_future_steps=self.num_future_boundary_steps, - ) if self.datastore.is_ensemble and self.load_single_member: warnings.warn( diff --git a/tests/test_time_slicing.py b/tests/test_time_slicing.py index ca05092b5..fe29154ae 100644 --- a/tests/test_time_slicing.py +++ b/tests/test_time_slicing.py @@ -569,9 +569,9 @@ def test_forecast_boundary_anchors_on_init_not_target(): ) -def test_check_time_overlap_insufficient_raises(): +def test_insufficient_boundary_coverage_raises(): """If the boundary cannot be cropped enough to cover the requested - past-window, ``check_time_overlap`` surfaces a clear error.""" + past-window, ``crop_time_if_needed`` surfaces a clear error.""" interior_datastore = SinglePointDummyDatastore( state_data=ANALYSIS_STATE_VALUES, forcing_data=FORCING_VALUES, From 812d25a04e58ed4e9a17b2e986e198a471eee3ba Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 11:23:37 +0200 Subject: [PATCH 34/35] Raise ValueError instead of assert in boundary windowing Convert the lone center_time <= target_time assert in _window_forcing_in_time to a ValueError so the windowing path raises uniformly. The remaining da_state assertions are type-narrowing guards. Co-Authored-By: Claude Opus 4.8 --- neural_lam/weather_dataset.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 4df53d06f..a4e5b26aa 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -422,10 +422,12 @@ def _window_forcing_in_time( np.floor((target_time - forcing_at.values) / forecast_step) ) center_time = forcing_at.values + lead * forecast_step - assert center_time <= target_time, ( - "Boundary forecast valid time runs ahead of the interior " - f"target time ({center_time} > {target_time})." - ) + if center_time > target_time: + raise ValueError( + "Boundary forecast valid time runs ahead of the " + f"interior target time ({center_time} > " + f"{target_time})." + ) window_start = lead - num_past_steps window_end = lead + num_future_steps + 1 From c096dc43c9bd64394d748389877631cfa752207f Mon Sep 17 00:00:00 2001 From: sadamov Date: Thu, 25 Jun 2026 11:36:58 +0200 Subject: [PATCH 35/35] added helper for empty dataarrays --- neural_lam/weather_dataset.py | 58 ++++++++++++++++++++++------------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index a4e5b26aa..3c3e8563a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -468,6 +468,38 @@ def _window_forcing_in_time( return xr.concat(da_list, dim="time") + def _empty_windowed_dataarray( + self, grid_index: xr.DataArray, target_times: xr.DataArray + ) -> xr.DataArray: + """Build an empty windowed forcing/boundary dataarray. + + Used when no forcing (or no boundary) is configured: the feature + dimension has size 0 so downstream code can unpack a stable 5-tuple. + + Parameters + ---------- + grid_index : xr.DataArray + The ``grid_index`` coordinate to use (interior or boundary grid). + target_times : xr.DataArray + The ``time`` coordinate spanning the autoregressive target steps. + + Returns + ------- + xr.DataArray + Empty array with dims + ``("time", "grid_index", "forcing_feature")`` and a zero-length + feature dimension. + """ + return xr.DataArray( + data=np.empty((self.ar_steps, grid_index.size, 0)), + dims=("time", "grid_index", "forcing_feature"), + coords={ + "time": target_times, + "grid_index": grid_index, + "forcing_feature": [], + }, + ) + def _build_item_dataarrays( self, idx: int ) -> tuple[ @@ -577,17 +609,8 @@ def _build_item_dataarrays( forcing_feature_windowed=("forcing_feature", "window") ) else: - # create an empty forcing tensor with the right shape - da_forcing_windowed = xr.DataArray( - data=np.empty( - (self.ar_steps, da_state.grid_index.size, 0), - ), - dims=("time", "grid_index", "forcing_feature"), - coords={ - "time": da_target_times, - "grid_index": da_state.grid_index, - "forcing_feature": [], - }, + da_forcing_windowed = self._empty_windowed_dataarray( + da_state.grid_index, da_target_times ) if da_boundary_windowed is not None: @@ -595,7 +618,6 @@ def _build_item_dataarrays( forcing_feature_windowed=("forcing_feature", "window") ) else: - # create an empty boundary tensor with the right shape # Use the boundary datastore's grid_index if available, otherwise # fall back to state grid_index (for the no-boundary case the # last dim is 0 anyway) @@ -610,16 +632,8 @@ def _build_item_dataarrays( ) else: boundary_grid_index = da_state.grid_index - da_boundary_windowed = xr.DataArray( - data=np.empty( - (self.ar_steps, boundary_grid_index.size, 0), - ), - dims=("time", "grid_index", "forcing_feature"), - coords={ - "time": da_target_times, - "grid_index": boundary_grid_index, - "forcing_feature": [], - }, + da_boundary_windowed = self._empty_windowed_dataarray( + boundary_grid_index, da_target_times ) return (