diff --git a/AGENTS.md b/AGENTS.md index c1a9c72d6..a3315f0f8 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 ccf79fdb1..722fc977b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ 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` 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 `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`, @@ -25,6 +27,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 68f7faa8b..8229f4b54 100644 --- a/README.md +++ b/README.md @@ -154,9 +154,13 @@ data/ And the content of `config.yaml` could in this case look like: ```yaml -datastore: - kind: mdp - config_path: danra.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 @@ -175,7 +179,16 @@ training: For now the neural-lam config only defines few things: -1. The kind of datastore and the path to its config +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). Exactly one datastore + 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. @@ -215,6 +228,20 @@ 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). | + +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 @@ -380,9 +407,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 @@ -584,13 +612,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 1da43fff9..fd33073b3 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 ( @@ -127,15 +128,22 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): Attributes ---------- - datastore : DatastoreSelection - Configuration specifying which datastore backend to use and its - associated settings. + 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). 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()``. """ - datastore: DatastoreSelection + datastores: Dict[str, DatastoreSelection] training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig) class _(dataclass_wizard.JSONWizard.Meta): @@ -174,9 +182,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 @@ -186,9 +198,35 @@ 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 (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) + 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: @@ -196,12 +234,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 - ) - return config, datastore + # 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, + ) + 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/neural_lam/create_graph.py b/neural_lam/create_graph.py index c67e43975..c00100504 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -730,7 +730,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/datastore/mdp.py b/neural_lam/datastore/mdp.py index 7cad45d7b..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 @@ -89,9 +102,19 @@ def __init__( _ds = mdp.create_dataset(config=self._config) _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 _ds and "forcing" not in _ds: + raise ValueError( + f"Datastore at {self._config_path} contains neither 'state' " + "nor 'forcing' data. At least one is required." + ) + self._ds = _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 @@ -209,10 +232,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]: """ @@ -229,10 +254,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. @@ -300,10 +327,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) @@ -477,9 +517,13 @@ def grid_shape_state(self) -> CartesianGridShape: 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 (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_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) @@ -567,3 +611,14 @@ 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 the number of grid points in the dataset. + + Returns + ------- + int + The number of grid points in the dataset. + """ + return len(self._ds.grid_index) 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 e14bbc740..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 @@ -252,11 +252,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, ) @@ -400,7 +406,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 @@ -551,22 +558,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): """ diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 71ce79510..c93abb06f 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -3,7 +3,7 @@ # Standard library import os import warnings -from typing import Any +from typing import Any, Optional # Third-party import matplotlib.pyplot as plt @@ -38,6 +38,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, @@ -116,13 +117,16 @@ def __init__( if var_leads_metrics_watch is None: var_leads_metrics_watch = {} - # datastore and forecaster are excluded from saved hparams and must - # be provided explicitly when calling load_from_checkpoint. Saving - # args 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"]) + # datastore and forecaster are excluded from saved hparams and must be + # provided explicitly when calling load_from_checkpoint. Saving args + # 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", "datastore_boundary", "forecaster"] + ) self.datastore = datastore + self.datastore_boundary = datastore_boundary self.forecaster = forecaster self.matched_metrics: set = set() @@ -207,6 +211,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) @@ -308,10 +344,13 @@ 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. + 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, batch_times = batch + init_states, target_states, forcing, boundary, batch_times = batch init_states = (init_states - self.state_mean) / self.state_std target_states = (target_states - self.state_mean) / self.state_std @@ -334,7 +373,20 @@ def on_after_batch_transfer(self, batch, dataloader_idx): forcing - self.forcing_mean_tiled ) / self.forcing_std_tiled - return init_states, target_states, forcing, batch_times + 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): """ @@ -344,7 +396,7 @@ def common_step(self, batch): ---------- batch : tuple The batch of data containing initial states, target states, - forcing features, and batch times. + forcing features, boundary forcing features, and batch times. Returns ------- @@ -352,7 +404,16 @@ def common_step(self, batch): A tuple containing prediction, target states, predicted standard deviation, and batch times. """ - init_states, target_states, forcing_features, batch_times = batch + ( + init_states, + target_states, + forcing_features, + boundary_features, + batch_times, + ) = batch + # 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 ) @@ -614,7 +675,7 @@ 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 77c9b3cfd..4c4b7722d 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -262,7 +262,7 @@ def main() -> None: ) 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 f98065c49..02c116871 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -346,6 +346,20 @@ def main(input_args=None): default=1, help="Number of future time steps to use as input for forcing data", ) + data_group.add_argument( + "--num_past_boundary_steps", + type=int, + default=1, + 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, when present", + ) data_group.add_argument( "--load_single_member", action="store_true", @@ -390,8 +404,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 + ) # Check --var_leads_metrics_watch variable indices against the datastore # so users get an immediate error instead of an IndexError deep in the @@ -413,6 +429,9 @@ def main(input_args=None): ar_steps_eval=args.ar_steps_eval, num_past_forcing_steps=args.num_past_forcing_steps, num_future_forcing_steps=args.num_future_forcing_steps, + num_past_boundary_steps=args.num_past_boundary_steps, + num_future_boundary_steps=args.num_future_boundary_steps, + datastore_boundary=datastore_boundary, load_single_member=args.load_single_member, batch_size=args.batch_size, num_workers=args.num_workers, @@ -463,6 +482,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/neural_lam/utils.py b/neural_lam/utils.py index 942eb2068..d66b4ca60 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -12,8 +12,10 @@ from typing import Any, Iterator, Union, overload # 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 @@ -325,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) @@ -354,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) @@ -825,3 +827,197 @@ 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( + 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 ``da_available`` is wide enough to + support a windowed lookup driven by ``da_requested`` times with the + given past/future window sizes. + + Parameters + ---------- + 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 ``da_requested`` time, measured in + ``da_available`` steps. + + Raises + ------ + ValueError + If ``da_available`` does not cover the required time range. + """ + if da_requested_is_forecast: + times_requested = da_requested.analysis_time + else: + times_requested = da_requested.time + time_min_requested = times_requested.min().values + time_max_requested = times_requested.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_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 - ( + max_lead - num_future_steps * time_step_available + ) + else: + 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) + + 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_available > required_time_min: + raise ValueError( + "`da_available` starts too late to cover the requested window. " + f"Required start: {required_time_min}, " + f"but `da_available` starts at {time_min_available}." + ) + + if time_max_available < required_time_max: + raise ValueError( + "`da_available` ends too early to cover the requested window. " + f"Required end: {required_time_max}, " + f"but `da_available` ends at {time_max_available}." + ) + + +def crop_time_if_needed( + 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 ``da_requested`` so that + ``da_available`` covers every needed window. If ``check_time_overlap`` + 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`. + + Returns + ------- + xr.DataArray + Possibly cropped ``da_requested``. + """ + if da_requested is None or da_available is None: + return da_requested + + try: + check_time_overlap( + da_requested, + da_available, + da_requested_is_forecast, + da_available_is_forecast, + num_past_steps, + num_future_steps, + ) + return da_requested + except ValueError: + 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: + 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) + 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] + 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 + + first_valid_idx = int( + np.searchsorted(requested_tvals, required_min, side="left") + ) + 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 + + if n_removed_begin > 0 or n_removed_end > 0: + log_on_rank_zero( + 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", + ) + da_requested = da_requested.isel( + {crop_dim: slice(first_valid_idx, last_valid_idx_plus_one)} + ) + return da_requested 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 4168396a6..3c3e8563a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -3,7 +3,7 @@ # Standard library import datetime import warnings -from typing import Iterator, Optional, Union +from typing import Any, Iterator, Optional, Union # Third-party import numpy as np @@ -13,15 +13,55 @@ # First-party from neural_lam.datastore.base import BaseDatastore +from neural_lam.utils import crop_time_if_needed, get_time_step class WeatherDataset(torch.utils.data.Dataset): """Dataset class for weather data. - Loads and processes weather data from a given datastore. See - :meth:`__init__` for the full parameter list. + 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 + ---------- + datastore : BaseDatastore + The datastore to load the data from (e.g. mdp). + split : str, optional + The data split to use ("train", "val" or "test"). Default is "train". + ar_steps : int, optional + The number of autoregressive steps. Default is 3. + num_past_forcing_steps: int, optional + Number of past time steps to include in forcing input. If set to i, + forcing from times t-i, t-i+1, ..., t-1, t (and potentially beyond, + given num_future_forcing_steps) are included as forcing inputs at time t + Default is 1. + num_future_forcing_steps: int, optional + Number of future time steps to include in forcing input. If set to j, + 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 + sample. If `True`, only ensemble member 0 is used. Default is False, + so all members are used when available. """ + INIT_STEPS = 2 + def __init__( self, datastore: BaseDatastore, @@ -29,29 +69,14 @@ 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, ) -> 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 ------ @@ -66,8 +91,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( @@ -81,6 +109,42 @@ def __init__( "The datastore must provide state data for the WeatherDataset." ) + # 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 + ) + else: + self.da_boundary_forcing = None + + # 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 + 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 + ) + + # 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, + 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( "only using first ensemble member, so dataset size is " @@ -201,50 +265,28 @@ def __len__(self) -> int: def _slice_state_time( self, da_state: xr.DataArray, idx: int, n_steps: int ) -> xr.DataArray: - """ - 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 ) @@ -252,132 +294,220 @@ def _slice_state_time( {"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( - self, da_forcing: xr.DataArray, idx: int, n_steps: int + def _window_same_forecast_by_idx( + self, + da_forcing: xr.DataArray, + idx: int, + state_times: xr.DataArray, + num_past_steps: int, + num_future_steps: int, ) -> xr.DataArray: + """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 autoregressive time step. The two `init_steps` can also be - # used as past forcings. - init_steps = 2 + init_steps = self.INIT_STEPS da_list = [] - 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, self.num_past_forcing_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 - - 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 (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( + [model_init_time], method="pad" + )[0] + if forcing_at_idx < 0: + raise ValueError( + "Boundary/forcing analysis times start after the model " + f"init time ({model_init_time})." ) + forcing_at = da_forcing.analysis_time[forcing_at_idx] + if model_init_time == forcing_at.values: + if forcing_at_idx == 0: + raise ValueError( + "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] - 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) ) + center_time = forcing_at.values + lead * forecast_step + 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 - 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, self.num_past_forcing_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 - - # Slice the data over the desired time window - da_sliced = da_forcing.isel(time=slice(start_idx, end_idx + 1)) + 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}." + ) - da_sliced = da_sliced.rename({"time": "window"}) + window_start = forcing_time_idx - num_past_steps + window_end = forcing_time_idx + num_future_steps + 1 - # Assign the 'window' coordinate to be relative positions - da_sliced = da_sliced.assign_coords( - window=np.arange(len(da_sliced.window)) + 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) - # 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]} - ) + return xr.concat(da_list, dim="time") - da_list.append(da_sliced) + def _empty_windowed_dataarray( + self, grid_index: xr.DataArray, target_times: xr.DataArray + ) -> xr.DataArray: + """Build an empty windowed forcing/boundary dataarray. - # Concatenate the list of DataArrays along the 'time' dimension - da_concat = xr.concat(da_list, dim="time") + Used when no forcing (or no boundary) is configured: the feature + dimension has size 0 so downstream code can unpack a stable 5-tuple. - return da_concat + 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[xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray]: + ) -> tuple[ + xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray + ]: """ - 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 ---------- @@ -392,6 +522,9 @@ def _build_item_dataarrays( 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. """ @@ -417,20 +550,53 @@ def _build_item_dataarrays( 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, + ) + + if self.da_boundary_forcing is not None: + da_boundary_windowed = self._window_forcing_in_time( + da_forcing=self.da_boundary_forcing, + 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 # 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)) @@ -443,32 +609,49 @@ 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: + da_boundary_windowed = da_boundary_windowed.stack( + forcing_feature_windowed=("forcing_feature", "window") + ) + else: + # 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 = self._empty_windowed_dataarray( + boundary_grid_index, da_target_times ) return ( da_init_states, da_target_states, da_forcing_windowed, + da_boundary_windowed, da_target_times, ) def __getitem__( self, idx: int - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, np.ndarray]: + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ]: """ Return a single training sample, which consists of the initial states, - target states, forcing and batch times. + target states, forcing, boundary and batch times. The returned data is unstandardized; normalization is applied on-device in `ForecasterModule.on_after_batch_transfer`. @@ -485,13 +668,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 + ``(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) @@ -507,6 +696,7 @@ def __getitem__( da_init_states, da_target_states, da_forcing_windowed, + da_boundary_windowed, da_target_times, ) = self._build_item_dataarrays(idx=idx) @@ -523,17 +713,24 @@ 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, num_grid_nodes, num_state_vars) - # target_states: (ar_steps, num_grid_nodes, num_state_vars) - # forcing: (ar_steps, num_grid_nodes, num_forcing_vars * window) - # target_times: (ar_steps,) + # 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, target_times + return init_states, target_states, forcing, boundary, target_times def __iter__( self, - ) -> Iterator[tuple[torch.Tensor, torch.Tensor, torch.Tensor, np.ndarray]]: + ) -> Iterator[ + tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ] + ]: """ Convenience method to iterate over the dataset. @@ -648,6 +845,9 @@ def __init__( ar_steps_eval: int = 25, 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, @@ -679,8 +879,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.load_single_member = load_single_member @@ -707,22 +910,26 @@ 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] = { + "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, split="train", ar_steps=self.ar_steps_train, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) self.val_dataset = WeatherDataset( datastore=self._datastore, split="val", ar_steps=self.ar_steps_eval, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) if stage == "test" or stage is None: @@ -730,9 +937,7 @@ def setup(self, stage: Optional[str] = None) -> None: datastore=self._datastore, split=self.eval_split, ar_steps=self.ar_steps_eval, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, + **shared_kwargs, ) def train_dataloader(self) -> torch.utils.data.DataLoader: diff --git a/pyproject.toml b/pyproject.toml index 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] 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/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/.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..fa4bc8b80 --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/config.yaml @@ -0,0 +1,22 @@ +datastores: + danra: + kind: mdp + config_path: danra.datastore.yaml + era5: + kind: mdp + config_path: era5.datastore.yaml +training: + state_feature_weighting: + __config_class__: ManualStateFeatureWeighting + 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 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/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..9a03970bb --- /dev/null +++ b/tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/era5.datastore.yaml @@ -0,0 +1,108 @@ +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: 2022-03-30T00:00 + end: 2022-04-12T00:00 + step: PT6H + chunking: + time: 1 + splitting: + dim: time + splits: + train: + 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: 2022-03-30T00:00 + end: 2022-04-12T00:00 + test: + 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: + values: [1000,] + units: hPa + dim_mapping: + time: + method: rename + dim: time + forcing_feature: + method: stack_variables_by_var_name + dims: [level] + name_format: "{var_name}{level}hPa" + grid_index: + method: stack + 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 + forcing_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + 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: + static_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [longitude, latitude] + target_output_variable: static + +extra: + projection: + class_name: PlateCarree + kwargs: + central_longitude: 0.0 diff --git a/tests/dummy_datastore.py b/tests/dummy_datastore.py index b269bfb2c..6e30f423f 100644 --- a/tests/dummy_datastore.py +++ b/tests/dummy_datastore.py @@ -479,6 +479,58 @@ 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 with no state. + + 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=0, 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, + ) + 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): """Small offline datastore for ensemble WeatherDataset tests. 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..096c15928 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) @@ -70,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)) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 4b35840ec..94dbe5667 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -14,8 +14,15 @@ 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.dummy_datastore import DummyDatastore, EnsembleDummyDatastore +from tests.conftest import ( + init_datastore_boundary_example, + init_datastore_example, +) +from tests.dummy_datastore import ( + BoundaryDummyDatastore, + DummyDatastore, + EnsembleDummyDatastore, +) @pytest.mark.parametrize("datastore_name", DATASTORES.keys()) @@ -48,7 +55,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 +106,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]") @@ -211,9 +218,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) @@ -396,8 +405,8 @@ def test_ensemble_index_mapping_is_time_major(): load_single_member=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) @@ -420,8 +429,8 @@ def test_ensemble_forcing_uses_same_member_when_available(): load_single_member=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) @@ -443,8 +452,8 @@ def test_ensemble_forcing_without_member_dim_is_shared(): load_single_member=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) @@ -479,3 +488,159 @@ def test_forecast_ensemble_len_scales_with_default_all_members(): ) assert len(dataset_all) == len(dataset_single) * 3 + + +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, + ) + + 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 the boundary tensor has no NaN values + 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, + ) + + _, _, _, boundary, _ = dataset[0] + assert boundary.shape[-1] == 0 + + +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 + ) + + 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=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", + 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, + ) + + # 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_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 8d516bfbd..cf4772c8a 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,22 +20,27 @@ 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 - ) + datastores={ + "main": nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + } ) 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, ) def test_on_after_batch_transfer(): """The hook standardizes state and forcing as (x - mean) / std and - leaves shapes and target times untouched.""" + leaves shapes, boundary forcing and target times untouched.""" datastore = init_datastore_example("mdp") model = _build_module(datastore) @@ -49,17 +55,19 @@ def test_on_after_batch_transfer(): forcing = torch.randn( 1, ar_steps, num_grid_nodes, num_forcing_vars * window_size ) + boundary = torch.randn(1, ar_steps, num_grid_nodes, 3) target_times = torch.randint(0, 1000000, (1, ar_steps)) - norm_init, norm_target, norm_forcing, norm_times = ( + norm_init, norm_target, norm_forcing, norm_boundary, norm_times = ( model.on_after_batch_transfer( - (init_states, target_states, forcing, target_times), 0 + (init_states, target_states, forcing, boundary, target_times), 0 ) ) assert norm_init.shape == init_states.shape assert norm_target.shape == target_states.shape assert norm_forcing.shape == forcing.shape + assert torch.equal(norm_boundary, boundary) assert torch.equal(norm_times, target_times) expected_init = (init_states - model.state_mean) / model.state_std @@ -104,6 +112,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 diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 616d563de..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 @@ -526,7 +528,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 @@ -548,7 +550,7 @@ def test_plot_examples_integration_saves_figure( # Get first example. pred_slice = prediction_rescaled[0].detach() target_slice = target_rescaled[0].detach() - time_slice = batch[3][0] + time_slice = batch[4][0] # Create DataArrays. dataset = WeatherDataset(datastore=datastore, split="train") @@ -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 73e2f9054..0b82f17a8 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_time_slicing.py b/tests/test_time_slicing.py index a8b022eda..fe29154ae 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( @@ -115,7 +184,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 ] @@ -190,4 +259,398 @@ 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 + + +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. + # 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 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, 2:5] + ) + np.testing.assert_array_equal( + boundary[1, 0, :], boundary_values[expected_analysis_idx, 4:7] + ) + + +def test_insufficient_boundary_coverage_raises(): + """If the boundary cannot be cropped enough to cover the requested + past-window, ``crop_time_if_needed`` 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, + ) + + +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 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()}), diff --git a/tests/test_training.py b/tests/test_training.py index bf1a5884a..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, @@ -99,9 +106,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 @@ -129,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, @@ -161,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.