Skip to content

feat : a general probabilistic forecasting interface#700

Open
Sir-Sloth-The-Lazy wants to merge 17 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:feat/probabilistic-forecasting-interface
Open

feat : a general probabilistic forecasting interface#700
Sir-Sloth-The-Lazy wants to merge 17 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:feat/probabilistic-forecasting-interface

Conversation

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor

Describe your changes

This PR implements the abstract constructions agreed in the #685 discussion (RFC: a general probabilistic forecasting interface), deliberately without any concrete Graph-EFM instantiation, so the design can be reviewed on its own and this can server as the base for addition of any probabilistic model, not just graph_efm.

Move ownership of the training objective from ForecasterModule onto the Forecaster. A new abstract Forecaster.compute_training_loss returns a finished (loss, loss_components) pair, so each forecaster owns its complete training objective (including assembling it from any internal terms).ForecasterModule.training_step now only injects the configured scoring rule (--loss), the interior mask and per_var_std, then logs the returned loss and components (component names prefixed with the phase). Per the discussion,
this move is applied to the deterministic setup as well, so the same concept sits in the same place in both stacks: the deterministic ARForecaster training loss is unchanged in value (covered by an equality test), it is just computed by the forecaster itself. Note for reviewers: since the method is abstract on Forecaster, any out-of-tree Forecaster subclass now has to implement it; all in-repo forecasters go through ARForecaster.

Add the probabilistic side, built on top of the deterministic one (no sample dimension leaks into deterministic components):

  • ProbabilisticForecaster (abstract): declares sample_ensemble, producing members stacked along a new dimension after batch, (B, S, pred_steps, num_grid_nodes, d_state). This encodes the module's only assumption "the forecaster can create ensemble forecasts of the correct shape" as a type contract, and leaves room for non-AR implementations (diffusion/flow) later.

  • ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): unrolls independent trajectories through a step predictor that samples its output (sequential loop as the safe first default), and by default trains on the configured scoring rule applied to the ensemble mean. Model-specific objectives (e.g. Graph-EFM's ELBO + CRPS) will live in subclasses that override compute_training_loss.

  • ProbabilisticForecasterModule: training inherited unchanged; validation samples an eval_ensemble_size ensemble and logs the RMSE of the ensemble mean (the example ensemble metric suggested in RFC: a general probabilistic forecasting interface #685, standing in until the ensemble-metrics PR), reusing the val_mean_loss / val_loss_unroll{i} log keys so existing checkpoint callbacks work unchanged. test_step raises NotImplementedError rather than silently evaluating a single member deterministically; ensemble test evaluation and plotting are a follow-up.

Dependencies: none. In particular this PR does not depend on the planned ensemble metrics (crps_ens, spread_squared). Follow-ups: ensemble metrics, the Graph-EFM forecaster (ELBO + CRPS, its own kl_beta/crps_weight config), ensemble test evaluation + plotting, and train_model.py wiring for
probabilistic models.

Issue Link

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 (no README changes needed — no user-facing CLI/config changes in this PR)
  • 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.

Add abstract Forecaster.compute_training_loss returning a finished
(loss, loss_components) pair, so each forecaster owns its complete
training objective. ForecasterModule.training_step now only injects the
configured scoring rule, interior mask and per_var_std, and logs the
result. The deterministic ARForecaster loss is unchanged in value.

Add the abstract ProbabilisticForecaster (sample_ensemble capability),
ProbabilisticARForecaster (sequential sampled rollouts, trains on the
configured score of the ensemble mean) and a minimal
ProbabilisticForecasterModule whose validation samples an ensemble and
logs the RMSE of the ensemble mean. Interface design from mllam#685.
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@joeloskarsson @observingClouds , this is the implementation of the dicussion on issue #685 hope this is what you wanted ! 👀.

@Sir-Sloth-The-Lazy
Sir-Sloth-The-Lazy marked this pull request as draft July 5, 2026 04:20

@observingClouds observingClouds 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.

Hi @Sir-Sloth-The-Lazy thanks for drafting this. I had a look at the discussion around the proposal in #685 and this PR. It looks well aligned. So far, I only have a few minor comments.

trajectory. This class adds ensemble forecasting on top: unrolling
several trajectories and stacking them along an ensemble dimension.
The default training objective scores the ensemble mean with the
injected scoring rule; forecasters with model-specific objectives

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.

"injected scoring rule". Can this be written more explicitly? So that it is clear to users where the scoring rule is set.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

solved in commit 3f5402d

states at each predicted step, used both as the prediction
targets and to overwrite boundary nodes during the rollouts.
Dims: same as one ensemble member.
score_fn : Callable

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.

fn always reminds me as an abbreviation for filename. Maybe use score_func or score_metric?

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.

Or even just score.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Solved in commit 3f5402d

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

renamed to score_metric as per the suggestion.

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@observingClouds , Hope the latest commit solves the concerns 😁

@joeloskarsson joeloskarsson left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

This looks quite nice I think, and is still pretty easy to follow. I had some outstanding thoughts regarding the connection to the deterministic version + some small things.

Comment thread neural_lam/models/module.py Outdated
init_states,
forcing_features,
target_states,
score_metric=self.loss,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

If the loss calculation is now entirely handled by the Forecaster, I think it could make more sense that self.loss is instantiated and sits on the Forecaster, rather than being passed to it? What do you think, is there any problem with that?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Good question, I don't think there's a hard problem with it, but it's a bit bigger than a one-liner so wanted to lay out the scope before doing it.

Right now self.loss isn't only feeding compute_training_loss ForecasterModule also calls it directly in validation_step/test_step (and for the spatial loss map) to report val/test loss, independent of whatever the forecaster's training objective ends up being. So "loss calculation entirely handled by the Forecaster" isn't quite true yet for the deterministic path the Module still owns a second, separate use of the same metric for reporting.

It also isn't uniform across the two modules we have: ProbabilisticForecasterModule.validation_step doesn't touch self.loss at all, it scores the ensemble mean with raw metrics.mse, and test_step isn't implemented there yet.

If we move it, the plan would be:

loss: str becomes a ctor arg on Forecaster (or ARForecaster), instantiated there as self.loss = metrics.get_metric(loss)
compute_training_loss drops the score_metric param entirely and reads self.loss internally
ForecasterModule.validation_step/test_step read self.forecaster.loss instead of self.loss
loss gets threaded through the forecaster constructors in train_model.py (currently only passed to ForecasterModule) and the ~15 test call sites that build ARForecaster/ProbabilisticARForecaster directly
None of that is a blocker, just confirming that's the scope before I make the change and whether ForecasterModule should keep a loss kwarg at all (for backward-compat / overriding the reporting metric independently of the training metric) or drop it entirely in favor of always reading it off the forecaster.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I am not a fan of having the ForecasterModule get the loss as self.forecaster.loss and then using that. A correct division of responsibilities would mean that the ForecasterModule does not know how to compute the loss for a batch, and that should be delegated to the Forecaster. So instead of getting the loss and using it, I think any loss computation that is needed in the validation/test step should be delegated to the Forecaster. If that means that we need additional functionality in the Forecaster (e.g. other types of loss reduction, not all dims) than we should implement that there.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

f0e01b1 committed the fix

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

f0e01b1 committed the fix

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think we need to differentiate taking (prediction, target, (optionally) predicted std) and computing:

  1. The loss. In the new setup here only the Forecaster knows how to compute its loss, and should be responsible for this.
  2. Other metrics. Computing a specific metric only requires knowledge about the shape of these 3 inputs. This is always known, and should not vary between different Forecasters. So I don't think it is required to have the Forecaster handle this.

Right now both computing the loss and all other metrics goes through the Forecaster using the score function. This makes sense for the loss, but seems unnecessary for other metrics (this adds a layer of obscurity that is not needed). I think we can keep direct metric computations in the Module, just call the Forecaster for the loss specifically.

Note that my precious comment was specific to the loss function (to not extract this from the Forecaster, but call the forecaster for this), not about other metrics such as entry-wise mse. Sorry if this was not very clear.

Comment thread neural_lam/models/module.py Outdated
Comment thread neural_lam/models/module.py Outdated
Comment thread neural_lam/models/forecasters/probabilistic.py Outdated
Comment thread neural_lam/models/forecasters/probabilistic.py
Comment thread neural_lam/models/forecasters/probabilistic.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread CHANGELOG.md Outdated
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

Thank you for taking out the time to review this @joeloskarsson means a lot. I will start working on the review now 😃

Sir-Sloth-The-Lazy and others added 7 commits July 8, 2026 09:10
score_metric/per_var_std were injected into compute_training_loss by
ForecasterModule and also used directly for val/test loss reporting,
duplicating config the forecaster already needs for its own objective.
ARForecaster/ProbabilisticARForecaster now own self.loss and
self.per_var_std (computed from an optional config ctor arg), and
ForecasterModule reads them off self.forecaster instead. Also trims
the CHANGELOG entry for mllam#685 down to one sentence per review feedback.
A Forecaster built without config now silently has per_var_std=None
when its predictor doesn't output its own std. Previously per_var_std
was always computed by ForecasterModule itself, so this gap didn't
exist; now that construction is split across two calls, catch it at
ForecasterModule init instead of crashing at the first val/test step.
Each member's predicted std is its own, not a spread computed across
the ensemble, so ensemble_std was a misleading name. Document on
ProbabilisticForecaster that a per-member std makes the predictive
distribution a mixture of Gaussians, and note in
ProbabilisticARForecaster.compute_training_loss that averaging the
per-member stds is a simplification of the true mixture variance
(which also includes the spread between member means).
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Scoring the ensemble mean with a pointwise metric only rewards the
mean being right, giving the model no incentive to keep a calibrated
spread, and risks training it to collapse to a point estimate. Redeclare
compute_training_loss as abstract on ProbabilisticARForecaster instead
of providing that as a default (it would otherwise silently fall back
to ARForecaster's single-rollout objective via MRO, not even the
ensemble mean). Concrete subclasses must define their own objective.

Tests that only need an instantiable forecaster now use a local
ConcreteProbabilisticARForecaster example (ensemble-mean scoring,
moved out of the library code); a new test locks in that the base
class itself cannot be instantiated.
Drop ProbabilisticARForecaster's ensemble_size constructor arg and the
implicit num_members=None -> self.ensemble_size fallback in
sample_ensemble; num_members is now always required. Baking a default
member count into the forecaster's state was unnecessary now that
compute_training_loss is abstract too (nothing in the shared base
class path used it) and just adds an implicit default callers could
silently rely on instead of deciding explicitly. The num_members < 1
validation moves from __init__ to sample_ensemble accordingly.
ProbabilisticForecasterModule.eval_ensemble_size follows suit: it no
longer defaults to None with a forecaster fallback, it's required.

Test-only ConcreteProbabilisticARForecaster (used wherever a concrete
probabilistic forecaster is needed for testing) gains its own
train_num_members for the training objective, since deciding how many
members to sample during training is now the concrete subclass's call.
Mirrors validation_step: samples eval_ensemble_size members and scores
the ensemble mean, same as validation. Factored the shared sampling +
scoring + logging into _ensemble_step(batch, phase) rather than
duplicating the block, since validation_step and test_step differ only
in their log-key prefix and which metrics dict collects the result.

Overrides on_test_epoch_end (rather than inheriting
ForecasterModule's) since this module's test_step doesn't populate
spatial_loss_maps or plot examples - the inherited version would crash
on torch.cat of an empty list.
Comment thread tests/test_probabilistic_forecaster.py Outdated
Comment on lines +64 to +65
# per_var_std is normally computed from config; override directly since
# this test only cares about the loss computation, not standardization.

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.

I'm not much of a fan of these inline comments and think the code should be clear enough to explain itself.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fair point, removed the assignment is clear on its own. Committed as 45ffdeb. Is there something more you want out of this ? I would really like to see, if this is enough or I can do more to make the code more readable.

Comment thread neural_lam/models/probabilistic_module.py Outdated
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

Thank you for your review ! I am on it 🤩

Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss
to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the
training loss, per review feedback.
Remove explanatory comments around the per_var_std overrides; the
assignments are clear on their own.
…crete deterministic/probabilistic modules

Introduce BaseForecasterModule (abstract) under models/forecasters/ holding
shared plumbing (training_step, common_step, batch standardization,
checkpoint compatibility, plotting/aggregation helpers), with
validation_step, test_step and on_test_epoch_end left abstract since they
differ meaningfully between evaluation modes. Rename ForecasterModule to
DeterministicForecasterModule and move it, alongside
ProbabilisticForecasterModule, into forecasters/ as siblings implementing
the shared contract, rather than one subclassing the other.
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@leifdenby @joeloskarsson @observingClouds, I hope these commits solve most of the issue, where i need a little more direction, I have left the comments asking for more help ! Thank you hope to here from you all soon. 🤩

Sir-Sloth-The-Lazy added a commit to Sir-Sloth-The-Lazy/neural-lam that referenced this pull request Jul 14, 2026
Removes estimate_likelihood/compute_step_loss and the per_var_std
buffer they existed to feed, plus the now-unused config constructor
arg (its only use was building per_var_std). Per the objective now
living on the Forecaster (mllam#700), the predictor stays a pure network
construct: encoder/prior/decoder plus the forward sampling path.
@Sir-Sloth-The-Lazy
Sir-Sloth-The-Lazy marked this pull request as ready for review July 15, 2026 08:01
@joeloskarsson

Copy link
Copy Markdown
Collaborator

Read through all of the fixes and discussion here now. Some small follow-up comments on my earlier points, but in general I think this is looking quite good. The refactor into a base module class was very good, but also moved some things around a bit. I would be happy to give this another full read through before approving once I am back from vacation. I think that should be timely if this goes in in v0.8.0.

Add JetBrains IDE project directory to the ignore list alongside
the existing .vim/.vscode entries.
…aster

DeterministicForecasterModule.validation_step/test_step read
self.forecaster.loss and self.forecaster.per_var_std directly, so the
Module still knew how to compute a loss from a prediction. Add an
abstract Forecaster.score (implemented on ARForecaster) that resolves
the pred_std fallback and applies a scoring rule internally; the Module
now only calls forecaster.score(...) and never touches loss/per_var_std
itself.
…ster

BaseForecasterModule.__init__ raised if forecaster.per_var_std was None
and the forecaster didn't predict its own std -- a check on the
Forecaster's configuration living in the wrong class. Move it into
ARForecaster via a new _resolve_pred_std helper, shared by score() and
compute_training_loss(): construction with config=None now always
succeeds (a valid state for forecasters only ever used for inference),
and the ValueError instead fires from the Forecaster itself, only once
scoring is attempted without any std to use.
…dules/

BaseForecasterModule, DeterministicForecasterModule and
ProbabilisticForecasterModule (Lightning wrappers around a Forecaster)
were mixed in with the Forecaster classes themselves under
neural_lam/models/forecasters/. Move them into their own
neural_lam/models/modules/ package (base.py, deterministic.py,
probabilistic.py) so the two concerns: what a forecaster is vs how
it's trained/evaluated by Lightning live in separate directories.
Pure move: internal imports and neural_lam/models/__init__.py updated
accordingly, no behavioural change.
Resolves the conflict on neural_lam/models/module.py: main's PR mllam#675
added --train_steps_to_log and deduplicated the val/test/train
prediction+loss+logging pattern into _compute_prediction_and_loss /
_log_step_loss, on top of the pre-mllam#700 module.py where the Module
still owned self.loss/self.per_var_std directly. This branch had since
deleted that file (split into modules/base.py, deterministic.py,
probabilistic.py, with loss ownership moved onto the Forecaster).

Ports mllam#675's dedup/logging helpers (_warn_skipped_steps, _log_step_loss,
the train_steps_to_log hparam) into modules/base.py, and rewires
DeterministicForecasterModule's validation_step/test_step to use them
via forecaster.score() instead of the removed self.loss/self.per_var_std.
_compute_prediction_and_loss is dropped rather than ported: it assumes a
single per-step-decomposable loss, which conflicts with
compute_training_loss's contract that a forecaster's training objective
may not decompose per step (e.g. an ELBO). training_step therefore still
only logs the aggregate train_loss; train_steps_to_log is accepted for
CLI/checkpoint compatibility but does not yet produce a
train_loss_unroll{i} breakdown -- documented inline as a deliberate
scope boundary, not an oversight.

Also fixes two call sites (deterministic.py, probabilistic.py) that
referenced the pre-rename _warn_skipped_val_steps, and absorbs the
unrelated utils.py -> utils/ package split from mllam#682 (auto-merged
cleanly, this branch never touched utils.py).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants