Skip to content

Add boundary data overlay to evaluation plots#636

Open
sadamov wants to merge 43 commits into
mllam:mainfrom
sadamov:feat/boundary-plotting
Open

Add boundary data overlay to evaluation plots#636
sadamov wants to merge 43 commits into
mllam:mainfrom
sadamov:feat/boundary-plotting

Conversation

@sadamov

@sadamov sadamov commented May 11, 2026

Copy link
Copy Markdown
Collaborator

Describe your changes

Wire the boundary datastore from config.yaml into ForecasterModule and 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_axis and vis.plot_prediction accept optional boundary_da and boundary_datastore parameters. 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.
  • Boundary grid coordinates are read via 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 optional datastore_boundary. During plot_examples, raw (unstandardized) boundary forcing is loaded from the boundary datastore and passed through to vis.plot_prediction for each timestep and variable (matched by index).
  • train_model.py passes datastore_boundary from config into ForecasterModule.
  • Integration test with real DANRA + ERA5 boundary data (marked @pytest.mark.slow).
  • README updated to document boundary overlay behavior during evaluation.

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

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section
    reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

closes #108

sadamov and others added 2 commits May 11, 2026 16:00
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>
@sadamov sadamov self-assigned this May 11, 2026
@sadamov sadamov added the enhancement New feature or request label May 11, 2026
@sadamov sadamov added this to the v0.8.0 milestone May 11, 2026
@sadamov sadamov force-pushed the feat/boundary-plotting branch from d5a4114 to 6d67242 Compare May 11, 2026 15:02
@sadamov sadamov requested a review from joeloskarsson May 11, 2026 15:02
@sadamov sadamov changed the title Add boundary datastore plotting support to vis.py Add boundary data overlay to evaluation plots May 11, 2026
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>
@sadamov

sadamov commented May 11, 2026

Copy link
Copy Markdown
Collaborator Author

here are some examples for dummy data with and without boundary each:

boundary_plot_interior_only boundary_plot_with_overlay

sadamov and others added 4 commits May 11, 2026 17:42
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>
Comment thread neural_lam/vis.py Outdated
Comment thread neural_lam/vis.py Outdated
target_time = np.array(
time_slice[t_i - 1].cpu(), dtype="datetime64[ns]"
)
boundary_da_t = da_boundary_forcing.sel(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@sadamov sadamov May 15, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@sadamov sadamov May 15, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a comment in the code about this rationale: 7b82b07

sadamov and others added 3 commits May 14, 2026 13:52
- 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>
@sadamov sadamov force-pushed the feat/boundary-plotting branch from 72c3197 to 0d6fc57 Compare May 14, 2026 12:53
@sadamov

sadamov commented May 14, 2026

Copy link
Copy Markdown
Collaborator Author

thanks @kshirajahere for the review, I'll have a look soon.
In the meantime I also implemented a "real-world" test-case using DANRA and ERA5 data across #635 and #636. The resulting plot looks like this:
image

PS: Tests fail also on main right now because of missing torch wheel for 2.12.0.

@kshirajahere

Copy link
Copy Markdown
Contributor

The plots look amazing @sadamov :D

@joeloskarsson

Copy link
Copy Markdown
Collaborator

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.

@sadamov

sadamov commented May 15, 2026

Copy link
Copy Markdown
Collaborator Author

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 boundary_margin_degrees. I did however not expose this to the user yet and it would have to be wired through ForecasterModule.plot_examples, if you wanted that.

sadamov added a commit to sadamov/neural-lam that referenced this pull request May 15, 2026
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.
@sadamov sadamov force-pushed the feat/boundary-plotting branch from efa0669 to 8efcb37 Compare May 15, 2026 08:32
@joeloskarsson

Copy link
Copy Markdown
Collaborator

Aha, I see 😄 Yea I guess wiring throughboundary_margin_degrees to the user is exactly what I was thinking about then. To follow all other plotting config, this should be an argparse flag. But this also feels like we are getting a ridiculous amount of flags 😮‍💨 Not sure if it is worth it. Maybe an option is to group this with the configuration of interior-boundary variable-mappings somehow? That is also plotting config really.

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 kshirajahere left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread neural_lam/vis.py
Comment thread README.md Outdated
sadamov and others added 2 commits May 28, 2026 21:22
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>
sadamov added a commit to sadamov/neural-lam that referenced this pull request Jun 7, 2026
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>
sadamov and others added 2 commits June 7, 2026 16:48
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>
sadamov and others added 2 commits June 15, 2026 15:12
…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
@sadamov sadamov force-pushed the feat/boundary-plotting branch from 03da1dd to 65ae26e Compare June 16, 2026 10:27
sadamov and others added 11 commits June 16, 2026 13:03
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>
@sadamov sadamov force-pushed the feat/boundary-plotting branch from 65ae26e to aa04b59 Compare June 16, 2026 11:11
sadamov and others added 6 commits June 16, 2026 20:59
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Plotting Boundaries in train_model.py --eval not just interior

3 participants