Refactor model class hierarchy into composable Forecaster/StepPredictor layers by Sir-Sloth-The-Lazy · Pull Request #208 · mllam/neural-lam

Sir-Sloth-The-Lazy · 2026-02-21T13:19:00Z

Describe your changes

Refactors the monolithic ARModel class into a composable hierarchy of smaller, focused components:

ForecasterModule (pl.LightningModule): Training loop, metrics, plotting, optimizer config
ARForecaster (nn.Module): Auto-regressive unrolling with boundary masking
StepPredictor (nn.Module): Single-step prediction, normalization, clamping

This separation makes it straightforward to add non-autoregressive forecasters, new step predictor architectures (e.g. Vision Transformers), or ensemble strategies without modifying the training infrastructure.

Also fixes two pre-existing bugs:

interior_mask_bool shape (1,) → (N,) for correct loss masking
all_gather_cat dimension collapse on single-device runs

Issue Link

Closes #49

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
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
I have requested a reviewer and an assignee

- Update test_datasets.py to use ForecasterModule instead of GraphLAM - Update test_plotting.py to use ForecasterModule instead of GraphLAM - Fix interior_mask_bool property shape (1,) -> (N,) for correct loss masking - Fix all_gather_cat to handle single-device runs without incorrect dim collapse

Sir-Sloth-The-Lazy · 2026-02-21T13:25:03Z

@joeloskarsson I have made the changes to the issue #49 please check if this is up to mark. I would love to hear your feedback.

Sir-Sloth-The-Lazy · 2026-02-22T04:47:41Z

@observingClouds and @leifdenby I have made the changes explained in the issue #49 if there are some mistakes please tell me about it would be helpful to know what i have done wrong thank you sorry for pestering you

leifdenby

Thank you for working on this @Sir-Sloth-The-Lazy!

This PR changes quite a lot (as expected of course because you are redesigning the class hierarchy for the models). I would like to give this a more thorough review, but first I need to get a clearer overview.

Would you be able to make a diagram which explains what each of the classes (forecast, step-predictor) take as input and return? With this I mean details of the shapes of the inputs and outputs too. I think we should put that in some documentation so that going forward the class hierarchy of the model architectures is clearer.

Thanks!

Sir-Sloth-The-Lazy

I have provided the clarification for the design choices. If some feel unsatisfactory please let me know , I would be happy to find some other way ;) @leifdenby

Sir-Sloth-The-Lazy · 2026-02-27T10:09:57Z

@leifdenby sorry for disturbing you, just a reminder , wanted to know if this is the write way or should i find another way to redesign the class

joeloskarsson

Thanks for starting all the work with this! I think this is a huge improvement to the codebase, and many things clearly become cleaner 😄 I added a number of comments and some points for discussion.

This is quite a large change to the codebase, that touches much of the code. There are many PRs open that will cause changes to code that is being refactored here. Mostly these PRs will make this change simpler I think, which is good. But because of this there will be some planning needed about when to merge this, and we will need to discuss this in our monthly dev meetings. But should be no issue to get this to a merge:able state and then making a plan for when to merge it!

Sir-Sloth-The-Lazy · 2026-03-01T14:22:48Z

@joeloskarsson thank you for this detailed review. Honestly , I was not feeling motivated enough to code today but the sheer amount of effort you have put in giving this feedback just made me instantly wanting to work again. I feel the time you would have taken out to read this code change would be huge! Thanks again to all of you guys @leifdenby @joeloskarsson will make the recommended changes and come again cheers ! ;)

AdMub · 2026-03-01T23:06:54Z

Hi @Sir-Sloth-The-Lazy, fantastic work getting this foundational PR started! I've been following the discussions in #49 closely, as I am planning to focus my GSoC proposal on extending this exact hierarchy to support the probabilistic and ensemble models (ARProbModel, Diffusion-LAM).

@joeloskarsson, regarding your comment on the pred_std handling: If the StepPredictor returns both next_state and pred_std, would you prefer the Forecaster to blindly pass both up to the ForecasterModule for loss computation, or should the Forecaster (specifically the ARForecaster subclass) actively use that pred_std during its unrolling loop to sample the next state for ensemble generation?

I want to make sure I accurately map out the data flow for the generative architectures in my proposal draft!

…r hierarchy - Replace opaque argparse.Namespace with explicit keyword arguments in StepPredictor, BaseGraphModel, BaseHiGraphModel, GraphLAM, HiLAM, and HiLAMParallel __init__ methods - Reorder methods in step_predictor.py: forward/expand_to_batch now appear before clamping methods - Update all instantiation sites (train_model.py, test_training.py, test_prediction_model_classes.py) to pass explicit kwargs - HiLAM helper methods (make_same/up/down_gnns) now use self.hidden_dim and self.hidden_layers instead of args parameter Addresses review comments on PR mllam#208.

- Rename border to boundary in Forecaster - Pass Forecaster object to ForecasterModule init instead of Predictor - Remove inline imports in ForecasterModule - Move loss-related pred_std logic fully into ForecasterModule - Delete obsolete test_refactored_hierarchy.py

kshirajahere · 2026-03-04T16:18:28Z

Following the review thread and thinking ahead to probabilistic support (#62): would it make sense to formalize the interface as StepPredictor -> (pred_state, pred_std | None), then keep loss-weighting and mask handling in ForecasterModule (with boundary mask read directly from datastore)? @joeloskarsson That seems to reduce coupling and may make ensemble/probabilistic evaluation plumbing cleaner later.

joeloskarsson · 2026-03-05T20:09:42Z

@AdMub regarding using pred_std during unrolling: While sampling using pred_std (just adding Gaussian noise with this std) would be the theoretically sound think to do, this is a very bad idea in practice and just makes us end up with noisy forecasts. So I see no need for this, even optionally.

joeloskarsson · 2026-03-05T20:11:37Z

@kshirajahere yes, I agree with you on this.

@Sir-Sloth-The-Lazy I interpret your review request that you have made changes and I should have a look again? Could you leave a short comment on each of my requested changes, just telling how you fixed each? (if you want link to a commit) That makes it a lot easier for me to review again 🫶 Thanks!

@GiGiKoneti

Squashed rebase onto current main of @GiGiKoneti's contribution from PR mllam#631, after the WeatherDataset changes from mllam#208 / mllam#239. Co-Authored-By: GiGi Koneti <GiGiKoneti@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Rebuilt onto current main. The original PR mllam#631 branch predated mllam#239 (GPU-side standardization) and mllam#208 (ForecasterModule refactor), so restoring its file wholesale would have reverted the standardize=True path back to the CPU. Instead, applied type hints directly to the current signatures and added a couple of `assert ... is not None` guards so the now-typed __len__ and _build_item_dataarrays bodies stay mypy-clean. Co-Authored-By: GiGi Koneti <GiGiKoneti@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Shristi-Goel

Re-applies @Shristi-Goel's PR mllam#205 onto current main. The target file `models/ar_model.py` was removed by mllam#208, so the change is now applied to `plot_examples` in `models/module.py` per the suggestion in the PR review (swap the two branches of the existing isinstance check rather than removing it - Wandb's carousel benefits from the bare key). Co-Authored-By: Shristi Goel <Shristi-Goel@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@RajdeepKushwaha5

…mllam#437) Re-applies @RajdeepKushwaha5's PR mllam#437 onto current main. The target file `models/ar_model.py` was removed by mllam#208, so the fix is applied to `on_test_epoch_end` in `models/module.py` next to the existing `matched_metrics`/`spatial_loss_maps` resets. Skipped the integration test from the PR diff: it referenced the pre-mllam#208 `GraphLAM(args=, ...)` constructor and the pre-mllam#239 `standardize=True` kwarg, so rebasing it would mean rewriting from scratch. Co-Authored-By: Rajdeep Singh <RajdeepKushwaha5@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@RajdeepKushwaha5

…mllam#437) Re-applies @RajdeepKushwaha5's PR mllam#437 onto current main. The target file `models/ar_model.py` was removed by mllam#208, so the fix is applied to `on_test_epoch_end` in `models/module.py` next to the existing `matched_metrics`/`spatial_loss_maps` resets. Skipped the integration test from the PR diff: it referenced the pre-mllam#208 `GraphLAM(args=, ...)` constructor and the pre-mllam#239 `standardize=True` kwarg, so rebasing it would mean rewriting from scratch. Co-Authored-By: Rajdeep Singh <RajdeepKushwaha5@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Debadri-das

Re-applies @Debadri-das's PR mllam#523 onto current main. The target file `base_graph_model.py` was moved to `models/step_predictors/graph/base.py` by mllam#208; the surrounding code is otherwise unchanged. Expanded the NOTE-comment to explain WHY the scaling is needed (initial pred_std otherwise sits at softplus(0) = ln(2), independent of variable scale, which blows up NLL/CRPS in early training for any variable whose physical step-diff std is much larger). Co-Authored-By: Debadri Das <Debadri-das@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Debadri-das

Re-applies @Debadri-das's PR mllam#523 onto current main. The target file `base_graph_model.py` was moved to `models/step_predictors/graph/base.py` by mllam#208; the surrounding code is otherwise unchanged. Expanded the NOTE-comment to explain WHY the scaling is needed (initial pred_std otherwise sits at softplus(0) = ln(2), independent of variable scale, which blows up NLL/CRPS in early training for any variable whose physical step-diff std is much larger). Co-Authored-By: Debadri Das <Debadri-das@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Ayushhgit

Re-applies @Ayushhgit's PR mllam#253 onto current main, scoped down to the clear-cut feature-weights fix and rebased onto the post-mllam#208 layout: - The original PR also added `+ 1e-6` to `softplus(pred_std_raw)` to guard against float32 underflow at very negative `pred_std_raw`. That overlapped with mllam#523's `* self.diff_std` scaling on the same line and is largely redundant once pred_std starts on the empirical scale, so dropped here to avoid coupling the two PRs. - Target file `models/ar_model.py` is gone since mllam#208; fix is applied at `models/module.py:120` next to the existing `per_var_std` registration. - Use `torch.finfo(torch.float32).eps` instead of a hard-coded `1e-8` for the sqrt epsilon, matching the eps choice used a few lines below for `state_std` / `forcing_std`. - Slimmed tests/test_numerical_stability.py to the 3 feature-weights assertions (finite output, eps does not perturb non-zero weights, warning fires with the zero index). Dropped the softplus / wmae / wmse integration tests since the softplus piece is no longer in scope. Refs mllam#526. Co-Authored-By: Ayush <Ayushhgit@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Mani212005

…llam#240) Re-applies @Mani212005's PR mllam#240 onto current main. The original PR was pre-mllam#208 and pre-mllam#641; this version applies the same Lightning-native pattern to current main with a much smaller surface area because the machinery is already in place via `load_forecaster_module_from_checkpoint` (introduced as part of the mllam#208 refactor). Source change: - Remove the `if not self.restore_opt: ...` fresh-optimizer hack from `ForecasterModule.on_load_checkpoint`. The hack only reset the optimizer state and silently left epoch / scheduler / callbacks restored from the checkpoint, which is not what `--restore_opt=False` is supposed to mean. - In `train_model.py`, branch on `args.restore_opt` for the `--load` path: weights-only goes through `load_forecaster_module_from_checkpoint` + `trainer.fit(model)` (no `ckpt_path`), full-resume keeps `trainer.fit(..., ckpt_path=args.load)`. Skipped the 239-line tests/test_checkpoint_restore.py from the original PR: it targeted the pre-mllam#208 API and would require a full rewrite. The existing `tests/test_training.py` smoke tests pass on the new code path. Verified: 18 training-related tests pass locally (test_training.py, test_cli.py, test_train_model_warnings.py). Co-Authored-By: Mani <Mani212005@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Mani212005

…llam#240) Re-applies @Mani212005's PR mllam#240 onto current main, picking up the naming/rename agreed with @joeloskarsson in the review thread. Source changes: - Rename CLI flag `--restore_opt` -> `--load_training_state` per @joeloskarsson's review comment (the old name was misleading because the flag controls epoch / scheduler / callbacks too, not just the optimizer). The internal `ForecasterModule(restore_opt=...)` keyword stays for now to keep legacy ARModel checkpoints loading via the `args=` shim. - Remove the `if not self.restore_opt: ...` fresh-optimizer hack from `ForecasterModule.on_load_checkpoint`. The hack only reset the optimizer state and silently left epoch / scheduler / callbacks restored from the checkpoint. - In `train_model.py`, branch on `args.load_training_state` for the `--load` path: weights-only goes through `load_forecaster_module_from_checkpoint` + `trainer.fit(model)` with no `ckpt_path`, full-resume keeps `trainer.fit(..., ckpt_path=args.load)`. - Add CLI assertion that `--load_training_state` requires `--load`. Test changes: - Drop the 239-line tests/test_checkpoint_restore.py from the original PR (targeted pre-mllam#208 API; @joeloskarsson also flagged it as adapting code to fit tests rather than testing real behaviour). - Add `test_load_training_state_requires_load` covering the new CLI assertion (parser.error -> SystemExit before any datastore loading). - Update tests/test_train_model_warnings.py mocks to set `load_training_state=False` explicitly (MagicMock returns truthy by default, which would trip the new assertion). Co-Authored-By: Mani <Mani212005@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Jayant-kernel

…rams (mllam#232) Re-applies @Jayant-kernel's PR mllam#232 onto current main. The original test targeted the pre-mllam#208 API (`GraphLAM(args=..., config=..., datastore=...)` all on the top-level model). Rewritten for the post-mllam#208 layout: build a GraphLAM step predictor + ARForecaster + ForecasterModule explicitly and assert the ignore covers both heavy non-pickle-safe objects (`datastore` and the now-additional `forecaster`). Refs mllam#148. Co-Authored-By: Jayant <Jayant-kernel@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Jayant-kernel

…rams (mllam#232) Re-applies @Jayant-kernel's PR mllam#232 onto current main. The original test targeted the pre-mllam#208 API (`GraphLAM(args=..., config=..., datastore=...)` all on the top-level model). Rewritten for the post-mllam#208 layout: build a GraphLAM step predictor + ARForecaster + ForecasterModule explicitly and assert the ignore covers both heavy non-pickle-safe objects (`datastore` and the now-additional `forecaster`). Refs mllam#148. Co-Authored-By: Jayant <Jayant-kernel@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

CONTRIBUTING.md (this PR) now owns the general contributor workflow. Remove the duplicate sections from AGENTS.md (Issues, Pull Requests, Changelog) and add a pointer to CONTRIBUTING.md as the first thing agents should read. Kept and refined the AI-specific rules: - Search before creating issues/PRs (now consolidated to one section with the exact gh search commands). - Re-read the full thread before every action, reload after context gaps - the rule that matters most for agents. - Communication style (terse, no filler ban list). - AI attribution in commit trailers. Also updated the codebase reference to reflect the post-mllam#208 ForecasterModule / Forecaster / StepPredictor hierarchy, and added `pytest -m "not slow"` to the command list (per mllam#651). Net: 91 lines (down from 97), but the AI-specific content is denser. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Per request to make AGENTS.md easy to maintain by avoiding duplication with other docs: - Codebase section: dropped the snapshot module list (which already rots, was still showing pre-mllam#208 names in one version), replaced with a pointer at the README architecture overview and `git log --stat -- neural_lam/` for finding recently-moved files. - Commands section: kept only the create_graph / train_model / --eval CLI entry points (the AI-specific bit) and pointed at CONTRIBUTING.md > Before you push for the standard install / lint / test commands that already live there. - Re-read the thread rule: compressed three bullets to one paragraph, same content. Net: 68 lines (down from 91). Maintenance burden: one paragraph per section, mostly links. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…r hierarchy - Replace opaque argparse.Namespace with explicit keyword arguments in StepPredictor, BaseGraphModel, BaseHiGraphModel, GraphLAM, HiLAM, and HiLAMParallel __init__ methods - Reorder methods in step_predictor.py: forward/expand_to_batch now appear before clamping methods - Update all instantiation sites (train_model.py, test_training.py, test_prediction_model_classes.py) to pass explicit kwargs - HiLAM helper methods (make_same/up/down_gnns) now use self.hidden_dim and self.hidden_layers instead of args parameter Addresses review comments on PR mllam#208.

- Rename border to boundary in Forecaster - Pass Forecaster object to ForecasterModule init instead of Predictor - Remove inline imports in ForecasterModule - Move loss-related pred_std logic fully into ForecasterModule - Delete obsolete test_refactored_hierarchy.py

Sir-Sloth-The-Lazy added 8 commits February 21, 2026 17:42

refactor: add Forecaster and StepPredictor abstract base classes

e7eb060

refactor: add ARForecaster implementing auto-regressive unrolling

4aa3112

refactor: add ForecasterModule as pl.LightningModule wrapper

ea5f185

refactor: change BaseGraphModel base class to StepPredictor

6892997

refactor: update train_model.py and __init__.py for new hierarchy

0bd2f5e

test: add tests for refactored model hierarchy and checkpoint loading

1781aba

refactor: archive legacy monolithic ar_model.py

3e1011e

sadamov requested a review from joeloskarsson February 21, 2026 14:57

leifdenby requested changes Feb 23, 2026

View reviewed changes

Comment thread neural_lam/models/archive/ar_model.py Outdated

Comment thread neural_lam/models/forecaster_module.py Outdated

Sir-Sloth-The-Lazy commented Feb 24, 2026

View reviewed changes

Removed ar_model.py

5305a65

Sir-Sloth-The-Lazy requested a review from leifdenby February 24, 2026 12:04

Merge branch 'main' into refactor/model-class-hierarchy-issue-49

b013643

joeloskarsson mentioned this pull request Mar 1, 2026

Refactor model class hierarchy #49

Closed

joeloskarsson requested changes Mar 1, 2026

View reviewed changes

joeloskarsson mentioned this pull request Mar 1, 2026

add support for Step Predictor and Forecaster abstract Interfaces. #300

Closed

21 tasks

Sir-Sloth-The-Lazy requested a review from joeloskarsson March 3, 2026 16:10

Sir-Sloth-The-Lazy added 3 commits March 3, 2026 21:47

Merge branch 'refactor/model-class-hierarchy-issue-49' into main

7dc80f5

Merge branch 'main' of github.com:Sir-Sloth-The-Lazy/neural-lam

b59f559

sadamov mentioned this pull request Jun 4, 2026

Modernize linting stack: migrate to Ruff and streamline pre-commit hooks #614

Open

21 tasks

This was referenced Jun 5, 2026

fix pred_std scaling #523

Closed

Mask boundary nodes in spatial loss maps during test_step #568

Open

sadamov mentioned this pull request Jun 9, 2026

Takeover: align #252 with README glossary, merge 25 main commits, address Joel review Mohit-Lakra/neural-lam#6

Merged

Sir-Sloth-The-Lazy mentioned this pull request Jul 1, 2026

RFC: a general probabilistic forecasting interface #685

Open

aliviahossain mentioned this pull request Jul 14, 2026

AGENTS.md "Key modules" section is out of date and points agents at the wrong tree #704

Closed

Uh oh!

Conversation

Sir-Sloth-The-Lazy commented Feb 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

Issue Link

Type of change

Checklist before requesting a review

Uh oh!

Sir-Sloth-The-Lazy commented Feb 21, 2026

Uh oh!

Sir-Sloth-The-Lazy commented Feb 22, 2026

Uh oh!

leifdenby left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy left a comment

Choose a reason for hiding this comment

Uh oh!

Sir-Sloth-The-Lazy commented Feb 27, 2026

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy commented Mar 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AdMub commented Mar 1, 2026

Uh oh!

kshirajahere commented Mar 4, 2026

Uh oh!

joeloskarsson commented Mar 5, 2026

Uh oh!

joeloskarsson commented Mar 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

Sir-Sloth-The-Lazy commented Feb 21, 2026 •

edited

Loading

Sir-Sloth-The-Lazy commented Mar 1, 2026 •

edited

Loading