Add boundary data overlay to evaluation plots#636
Conversation
Add support for loading boundary forcing from a separate datastore, enabling LAM models to ingest boundary conditions from a different domain (e.g. ERA5 boundaries for a COSMO/DANRA interior). - NeuralLAMConfig accepts optional `datastore_boundary` field - load_config_and_datastore returns 3-tuple (config, datastore, datastore_boundary) - WeatherDataset loads, windows, and standardizes boundary forcing - __getitem__ returns 5-tuple (init_states, target_states, forcing, boundary, target_times) - New CLI args --num_past_boundary_steps / --num_future_boundary_steps - ForecasterModule.common_step unpacks boundary (not yet wired to forward) - 4 new boundary-specific tests, all 157 tests pass refs mllam#108 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
d5a4114 to
6d67242
Compare
Add MDP-based ERA5 boundary example at tests/datastore_examples/mdp/era5_1000hPa_danra_100m_winds/ with config.yaml, era5.datastore.yaml (WeatherBench2 64x32 equiangular), and danra.datastore.yaml (DANRA 100m winds interior). Add DATASTORES_BOUNDARY_EXAMPLES dict and init_datastore_boundary_example() to conftest.py for use in boundary integration tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
MDPDatastore.__init__ crashed with KeyError when loading a datastore that has only forcing+static (no state), e.g. ERA5 boundary data. Fix is_ensemble check to guard against missing state, and grid_shape_state to fall back to forcing/static categories. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Make _get_analysis_times fall back to forcing file patterns when no
state files exist, guard get_dataarray("state") against empty var_names,
and prevent empty feature list from matching state loading path.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Guard against missing state/static feature keys in the zarr, not just forcing. Boundary-only datastores (e.g. ERA5) may lack state_feature entirely. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Return len(grid_index) directly instead of computing from grid_shape_state, which is more robust for boundary-only datastores. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
| target_time = np.array( | ||
| time_slice[t_i - 1].cpu(), dtype="datetime64[ns]" | ||
| ) | ||
| boundary_da_t = da_boundary_forcing.sel( |
There was a problem hiding this comment.
method="nearest" can silently overlay the wrong boundary timestep when the boundary datastore has missing times or a different cadence. For verification plots this should be exact, or at least use an explicit tolerance tied to the datastore step length and fail/skip when no valid boundary time exists.
There was a problem hiding this comment.
I actually like this to be nearest. Often the boundary will have coarser temporal resolution (e.g. 6h) and than the interior e.g. 1h. Now, neural-lam allows to select any number of future and past timesteps from the buondary that are nearest to that interior timestep creating a window. And I think it makes sense to display the nearest boundary in the verification plot because that is the boundary forcing closest to the truth that neural-lam can use in the inference. So unless you have strong objections, i'd rather leave this as is.
There was a problem hiding this comment.
I added a comment in the code about this rationale: 7b82b07
…eat/boundary-datastore
- Shrink the ERA5 boundary test dataset to 2022-03-30..2022-04-12 (was 1990-2022), add per-input lat/lon coord_ranges and enable mllam's convex-hull domain_cropping with include_interior_points=true. Stats computation drops from minutes to seconds and the cached zarr stays under 1 MB. - Stack [longitude, latitude] directly into grid_index in the era5 dim_mapping (per mllam's example.era5_cropped.yaml) so the original coord names survive for the convex-hull crop -- removes the need for any rename-preserve workaround in neural-lam. - Generalise the MDPDatastore units loop over self.spatial_coordinates with sensible defaults for x/y (m) and longitude/latitude/lon/lat (degrees_*) so ERA5-style geographic datastores work. - Register a pytest `slow` marker and a `--run-slow` CLI flag so the ERA5 boundary integration test (added in this PR) is skipped by default and can be run on demand. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
72c3197 to
0d6fc57
Compare
|
thanks @kshirajahere for the review, I'll have a look soon. PS: Tests fail also on main right now because of missing torch wheel for 2.12.0. |
|
The plots look amazing @sadamov :D |
|
I see the extent of the plot is now determined based on the full boundary size (please clarify if I missunderstood). I wonder if we want to give the user some more options here, for example if the boundary is very large one might want to only zoom in on the interior, with this covering most of the plot. Not sure where such an option best sits though. |
not quite. first ,there is a simple check (boundary data present yes/no), if no -> plot interior only. If yes -> then plot a 1° margin (in projection) around the interior. this 1° is independent of the full boundary size and can be set via |
Drop state metadata after init and raise KeyError on `state` lookups so plotting/model code that accidentally queries state on a boundary fails loudly. Add a focused test asserting plot_prediction works with such a datastore via the forcing grid. Addresses review comment on PR mllam#636 about testing the no-state path.
Drop state metadata after init and raise KeyError on `state` lookups so plotting/model code that accidentally queries state on a boundary fails loudly. Real ERA5-style boundary datastores expose only forcing fields, and the existing boundary tests (test_datasets.py) only ever access forcing on the boundary, so making the dummy state-less brings it closer to real boundary semantics without changing test behaviour.
efa0669 to
8efcb37
Compare
|
Aha, I see 😄 Yea I guess wiring through |
Boundary forcing was previously passed through unchanged in
`on_after_batch_transfer`, leaving the only normalization step on a
separate code path and inconsistent with how interior state/forcing are
handled. Wire it through the same on-device hook.
- `ForecasterModule.__init__` takes a new optional `datastore_boundary`
arg (excluded from `save_hyperparameters` alongside `datastore` and
`forecaster`, so it must be passed at `load_from_checkpoint` time).
- Register `boundary_mean` / `boundary_std` from
`datastore_boundary.get_standardization_dataarray("forcing")` when a
boundary datastore is provided; otherwise leave both as None.
- `on_after_batch_transfer` standardizes the boundary tensor the same
way it standardizes forcing: feature-major `(feature, window)`
per-feature stats are tiled once on the first batch and cached.
- Update the NOTE in `common_step` to reflect that boundary is now
standardized but still not consumed by the forecaster (mllam#108).
- Pass `datastore_boundary` from `train_model.main` to the module.
- New tests: `test_boundary_standardized_when_datastore_provided` and
`test_boundary_passthrough_when_no_boundary_datastore`.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
kshirajahere
left a comment
There was a problem hiding this comment.
Fascinating work. Thanks for giving me the opportunity to review it @sadamov really grateful :) left some 3 comments which I think might be worth addressing
The forecast-boundary path selected its analysis_time by pad-matching the first target time, which can pick a boundary forecast launched after model init - unavailable operationally. Anchor on the model init time (state_times[init_steps - 1], strictly before) instead, matching the research branch, and assert the per-step boundary valid time never runs ahead of the target. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Resolved conflicts: - neural_lam/datastore/mdp.py: kept the state-or-forcing-required validation while taking main's local _ds variable pattern (refactored upstream). - neural_lam/train_model.py: moved --num_past_boundary_steps and --num_future_boundary_steps to use data_group.add_argument, matching the argument-groups refactor from mllam#641. - neural_lam/weather_dataset.py: kept mllam#635's renamed _window_same_forecast_by_idx and shared_kwargs setup pattern; added type hints to match the post-mllam#631 type-hint sweep. Updated __getitem__ and __iter__ return types from 4-tuple to 5-tuple to reflect the new boundary tensor. Typed shared_kwargs as dict[str, Any] to satisfy mypy on the **kwargs splat. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Resolved conflicts: - neural_lam/train_model.py: moved --num_past_boundary_steps and --num_future_boundary_steps to use data_group.add_argument, matching the argument-groups refactor from mllam#641 (same as mllam#635 resolution). - neural_lam/weather_dataset.py: kept mllam#636's _slice_forcing_time signature with the extra num_past_steps/num_future_steps params, added type hints from main. Typed shared_kwargs as dict[str, Any] to satisfy mypy on the **kwargs splat. Updated _build_item_dataarrays, __getitem__, and __iter__ return types from 4-tuple to 5-tuple to reflect the new boundary tensor. - neural_lam/vis.py: kept mllam#636's boundary_da / boundary_datastore / boundary_margin_degrees args on plot_on_axis, plot_prediction, and plot_spatial_error. crop_to_interior is preserved as a deprecated parameter on plot_prediction and plot_spatial_error (already handled with a deprecation warning earlier on this branch). Type hints from main's post-mllam#631 sweep are applied throughout. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Matches the CI invocation introduced in mllam#651 so this branch's eventual ERA5 boundary integration test (and any other @pytest.mark.slow on this branch) is exercised by CI, not just by local --run-slow invocations. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Per the simplification in mllam#651, the custom --run-slow flag is being dropped in favour of relying on pytest's native -m marker selection. Remove the corresponding pytest_addoption + pytest_collection_modifyitems from conftest and the --run-slow flag from CI on this branch too, so this PR doesn't re-introduce conflicts when mllam#651 lands. The slow marker registration in pyproject.toml stays, ready for use on any future @pytest.mark.slow test on this branch. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ategory Replace the `datastore` + optional `datastore_boundary` config keys with a single `datastores` mapping keyed by a user-chosen name. A datastore's role is now implied by the categories it provides rather than by a dedicated key: a datastore with `state` data is the interior (model input and output), one without `state` data is used for input only (e.g. boundary forcing). At least one datastore must provide `state`; at most one boundary datastore is currently supported. `load_config_and_datastore` classifies the datastores and keeps returning the same `(config, interior, boundary)` tuple, so downstream callers are unchanged. Implements the design agreed in the discussion on mllam#652 (leifdenby's named `datastores` proposal, endorsed by joeloskarsson). refs mllam#108 refs mllam#652 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # neural_lam/config.py # neural_lam/models/module.py # neural_lam/weather_dataset.py # pyproject.toml
03da1dd to
65ae26e
Compare
The `slow` marker registration still described the markers as "deselected by default" (a leftover from the removed `--run-slow` flag). Slow tests actually run by default and are skipped with `-m "not slow"`, matching the behaviour agreed in mllam#651. Description-only change; no behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add optional boundary_da and boundary_datastore parameters to plot_on_axis and plot_prediction so that boundary data from a separate datastore can be overlaid on verification plots. Closes mllam#108 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Pass datastore_boundary through train_model.py into ForecasterModule. During --eval, plot_examples loads raw boundary forcing and overlays it underneath prediction/target panels via vis.plot_prediction. Add four boundary plotting tests using BoundaryDummyDatastore from PR mllam#635. Update README to document boundary plotting during evaluation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- `vis.plot_on_axis` and `vis.plot_prediction` no longer accept the `crop_to_interior` flag; the axes extent is always cropped to the interior bbox in the datastore's own projection. When `boundary_da` is provided the bbox is expanded by `boundary_margin_degrees` (new parameter, default 1.0), converted to projection units at the interior center so the margin is geometrically symmetric instead of lopsided in projections where lon-distance shrinks with latitude. - Cropped boundary datastores (e.g. ERA5 after convex-hull cropping) are now supported in `plot_on_axis`: when the boundary's grid is irregular the per-point `latitude`/`longitude` aux coords are used to reconstruct a regular grid via `set_index(...).unstack(...)`, with NaN for missing cells, yielding the expected "donut" of boundary forcing around the interior. - `plot_prediction` now sizes the figure from the interior's projected aspect ratio so portrait-oriented domains aren't squashed into a fixed landscape figure. - `plot_spatial_error` follows the same crop-to-interior default and drops `crop_to_interior` from its signature. - Tests updated for the new API and extent semantics; the obsolete `plot_boundary_context` flag concept is removed entirely -- the boundary is always rendered when a boundary datastore is configured. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
During evaluation the boundary overlay for a state variable is now drawn only when the boundary datastore exposes a forcing feature with the same name (previously paired by feature index, which silently mismatched e.g. DANRA `v100m` against ERA5 `mean_sea_level_pressure`). State variables with no matching boundary feature plot without an overlay -- no error, no warning. To pair fields across datastores that use different native variable names, align the names in the two mllam-data-prep configs (documented in README). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Asserts plot_prediction works against a boundary datastore with no state
category (forcing-only), exercising the get_xy("forcing") / get_lat_lon(
"forcing") path in vis.plot_on_axis. Pairs with the BoundaryDummyDatastore
state-less change in mllam#635.
Cadence between interior and boundary often differs (e.g. ERA5 6h vs DANRA 1h); we want every interior step to show the closest available boundary state, not skip frames. Document this on the sel() call so a future reviewer doesn't second-guess the missing tolerance.
Code defaults to 1.0 in vis.plot_on_axis and vis.plot_prediction, but the docstrings and CHANGELOG entry said 2.0. Align the docs with the actual value.
- tests/conftest.py: chdir to the boundary config's directory so mllam-data-prep can resolve the relative `interior_dataset_config_path` reference regardless of where pytest is invoked from; resolve the config path to absolute first so the chdir does not break the lookup. - neural_lam/vis.py: restore `crop_to_interior` keyword on `plot_prediction` and `plot_spatial_error` for backward compatibility, emit DeprecationWarning when passed. - README.md: clarify that boundary data is loaded for visualization only and that model-side consumption is a follow-up (mllam#108). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add an optional `boundary_datastore` field to `PlottingConfig` naming which entry in `datastores` supplies the boundary forcing for evaluation overlays. When unset it falls back to the single datastore without `state` data; `load_config_and_datastore` validates that a named datastore exists and is stateless. Implements the plotting-config selection agreed in mllam#652 (option 3). refs mllam#108 refs mllam#652 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
65ae26e to
aa04b59
Compare
Loading a config that still uses the old top-level `datastore:` key now raises a clear InvalidConfigError pointing to the named `datastores:` mapping. Add a regression test and note the behaviour in the CHANGELOG. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add an input/output table for the state/forcing/static categories in the Data section, as requested in mllam#652, so the datastore role-by-category explanation has an explicit definition to reference. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>



Describe your changes
Wire the boundary datastore from
config.yamlintoForecasterModuleand the plotting pipeline so that evaluation plots (--eval) overlay boundary forcing data from a separate domain underneath the interior prediction and ground truth panels.Changes:
vis.plot_on_axisandvis.plot_predictionaccept optionalboundary_daandboundary_datastoreparameters. When provided, boundary data is plotted underneath the interior field at alpha=0.5 and the map extent is expanded to cover the boundary domain.get_xy("forcing")/get_lat_lon("forcing")so that boundary datastores without state variables (e.g. ERA5 forcing-only) work correctly.ForecasterModule.__init__accepts an optionaldatastore_boundary. Duringplot_examples, raw (unstandardized) boundary forcing is loaded from the boundary datastore and passed through tovis.plot_predictionfor each timestep and variable (matched by index).train_model.pypassesdatastore_boundaryfrom config intoForecasterModule.@pytest.mark.slow).Depends on #635 (boundary datastore data loading). Please merge #635 first; after that the base branch of this PR can be updated to
main.Issue Link
refs #108
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee
closes #108