feat : a general probabilistic forecasting interface - #700
Open
Sir-Sloth-The-Lazy wants to merge 37 commits into
Open
feat : a general probabilistic forecasting interface#700Sir-Sloth-The-Lazy wants to merge 37 commits into
Sir-Sloth-The-Lazy wants to merge 37 commits into
Conversation
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.
Contributor
Author
|
@joeloskarsson @observingClouds , this is the implementation of the dicussion on issue #685 hope this is what you wanted ! 👀. |
Sir-Sloth-The-Lazy
marked this pull request as draft
July 5, 2026 04:20
observingClouds
left a comment
observingClouds
left a comment
Contributor
There was a problem hiding this comment.
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 |
Contributor
There was a problem hiding this comment.
"injected scoring rule". Can this be written more explicitly? So that it is clear to users where the scoring rule is set.
| 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 |
Contributor
There was a problem hiding this comment.
fn always reminds me as an abbreviation for filename. Maybe use score_func or score_metric?
Contributor
There was a problem hiding this comment.
Or even just score.
Contributor
Author
There was a problem hiding this comment.
renamed to score_metric as per the suggestion.
Contributor
Author
|
@observingClouds , Hope the latest commit solves the concerns 😁 |
joeloskarsson
requested changes
Jul 7, 2026
joeloskarsson
left a comment
joeloskarsson
left a comment
Collaborator
There was a problem hiding this comment.
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.
Contributor
Author
|
Thank you for taking out the time to review this @joeloskarsson means a lot. I will start working on the review now 😃 |
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 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. |
Contributor
There was a problem hiding this comment.
I'm not much of a fan of these inline comments and think the code should be clear enough to explain itself.
Contributor
Author
There was a problem hiding this comment.
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.
leifdenby
reviewed
Jul 8, 2026
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
requested review from
joeloskarsson,
leifdenby and
observingClouds
July 10, 2026 08:05
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
marked this pull request as ready for review
July 15, 2026 08:01
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. |
21 tasks
joeloskarsson
requested changes
Jul 27, 2026
joeloskarsson
left a comment
joeloskarsson
left a comment
Collaborator
There was a problem hiding this comment.
Went over the full PR again now. Things had changed quite a bit since my last overview, so had a few points that needs looking into. Most things are small though, so overall I think this should be easy fixes. The thing considering some more thought is the updated Forecaster class hierarchy.
_ensemble_step took the square root of each batch's MSE and let Lightning average those roots over the epoch, which does not give the RMSE: the root does not commute with the averaging, so the reported number was neither the RMSE nor a mean of RMSEs over a meaningful population. Accumulate only the per-variable squared errors per batch and reduce them once per epoch in a new _log_ensemble_rmse: gather across devices, average over every sample of the epoch and sum over variables, then take the root. Validation logs it from a new on_validation_epoch_end override (before the inherited implementation clears the metric lists) and testing from the existing on_test_epoch_end. This also drops the second metrics.mse call that recomputed, with sum_vars=False, the errors already computed at the top of _ensemble_step; sum_vars=True is exactly that result summed over the variable dimension, so the per-step values are now derived from the single computation.
How a forecast is produced and how a training objective is computed from it are orthogonal, but ARForecaster bundled both: auto-regressive unrolling and the deterministic single-forecast loss. That left no way to express an auto-regressive forecaster trained a different way other than inheriting the deterministic objective and overriding it, which is why ProbabilisticARForecaster had to re-declare compute_training_loss abstract after inheriting a concrete one, and why it took a loss argument naming a pointwise scoring rule that will not apply to it. Split the two axes: - ARForecaster now covers only auto-regressive unrolling (predictor, boundary masks, forward) and leaves compute_training_loss abstract. - DeterministicForecaster (new) supplies the objective half: score a single forecast with a configured scoring rule, plus the reporting score() and the pred_std fallback it needs. It makes no assumption about how the forecast is produced. - DeterministicARForecaster combines the two and is the deterministic model the CLI builds; this is the rename of the old ARForecaster. - ProbabilisticARForecaster is now ARForecaster + ProbabilisticForecaster, so it never inherits a concrete compute_training_loss to re-abstract, and takes neither config nor loss: a concrete probabilistic forecaster brings whatever configuration its own objective needs. The per-variable std computation moves to loss_weighting.get_per_var_std so both objective families can reuse it, and the abstract score() drops off Forecaster since it is specific to the deterministic objective (only DeterministicForecasterModule calls it; the probabilistic module scores ensembles with metrics.mse directly).
Moving the objective onto the Forecaster left --train_steps_to_log accepted
but inert, since compute_training_loss returns only a scalar. That is the
right general contract (an ELBO need not decompose over rollout steps), but
it dropped a breakdown every model in the repo can actually produce.
Give the decomposition a home on the class that knows the objective has one:
DeterministicForecaster.compute_step_losses returns the scoring rule per
predicted step, and compute_training_loss is now its mean, so the two cannot
drift. DeterministicForecasterModule overrides training_step to log the
breakdown through the existing _log_step_loss helper, producing the same
train_loss_unroll{i} keys as before. The general
BaseForecasterModule.training_step still logs only the scalar, so a
forecaster whose objective does not decompose is unaffected.
This deliberately sits on the deterministic objective rather than on the
auto-regressive rollout: a direct (non-AR) forecaster scored per lead time
decomposes just as well, so AR-ness is not what makes the breakdown
available.
Deleting this helper during the module split reintroduced the duplication mllam#675 had removed: validation_step and test_step each repeated the same common_step, score, mean-over-batch sequence. Restore it on DeterministicForecasterModule, which is the only module that scores a single prediction this way (the probabilistic one samples and scores ensembles instead, and nothing else calls common_step). training_step is deliberately not folded back in, so the helper now serves two callers rather than the original three. It goes through the forecaster's training objective rather than the reporting scoring rule used here; those two happen to coincide for the deterministic objective, but routing training through score() would put the choice of objective back in the module, which is what this PR moves onto the Forecaster.
…steps The abstract validation_step/test_step docstrings pointed at self.val_metrics/self.test_metrics as if this class owned them. It does consume both in its epoch-end hooks but never creates them, since which metrics are collected depends on the evaluation mode. Declare them as annotations here to state that contract explicitly, and reword the docstrings to refer to the ones the subclass creates. Keep the abstract declarations, and record why inline. LightningModule does define validation_step and test_step, but as no-op stubs rather than abstract methods (only training_step even warns), so dropping these would let a module that omits them instantiate happily and silently skip evaluation instead of failing at construction. Added a test covering that.
…mplied Apply the suggested module docstring wording, and carry the same correction through the surrounding docs: DeterministicForecasterModule and DeterministicForecaster describe a single forecast scored per predicted step, which need not have been produced by unrolling. "Rollout" is kept where it is accurate, i.e. on DeterministicARForecaster and the AR machinery itself.
… sampling Replace the *args/**kwargs passthrough on DeterministicForecasterModule and ProbabilisticForecasterModule with the full explicit signature and parameter docs, so each module's constructor is self-describing rather than pointing at the base class. eval_ensemble_size becomes keyword-only, matching that it is required and has no sensible default. Writing out the signatures exposed a latent fragility in hparam handling: save_hyperparameters() collects the arguments of whichever __init__ frame calls it, which was the base class's only because the subclasses forwarded opaquely. With explicit subclass signatures it instead captured their arguments as passed, i.e. before the base resolves mutable defaults and unpacks a legacy args namespace, leaving hparams.val_steps_to_log as None. Write the resolved values back after saving so hparams no longer depends on which frame happens to make the call. Also document that sample_ensemble draws members sequentially, and that this is an implementation choice rather than a constraint, since members are independent given the inputs.
joeloskarsson
requested changes
Aug 2, 2026
joeloskarsson
left a comment
joeloskarsson
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the work, had another look over. The big thing in this batch is the new multi-class inheritance with the ARForecaster, that I think causes us some issues and requires some work.
A couple general comments also, to make reviewing easier:
- To the extent possible, try to fix one thing in each commit, so that when commits are linked in the reviewer comments one does not have to pick apart what part of the change is relevant.
- Many comments (both in the code, but mainly in github responses) get a bit winding and long to read (perhaps a symptom of too unconstrained AI-assisted writing?). Please try to be concise and only make the points needed for the discussion.
save_hyperparameters() decides what to record by walking the constructor chain and letting the most derived __init__ win, on the assumption that its arguments are the authoritative record of how the object was built. BaseForecasterModule breaks that assumption: it unpacks a legacy args namespace and resolves mutable defaults after receiving them, so the values it runs on are not the values a subclass was called with. While the modules forwarded through *args/**kwargs they had no such variables to find and the base's frame was read instead, which is why this only surfaced once the constructors were written out and hparams.val_steps_to_log arrived as None. Passing a mapping makes save_hyperparameters use it verbatim and skip the inspection, so what is recorded equals what is used regardless of how a subclass writes its signature -- removing the undocumented requirement that subclasses not declare these parameters. This also replaces the write-back added in ade1929, which could not fully work: the snapshot into _hparams_initial is the final statement of save_hyperparameters, so correcting values after it returns reaches only the live hparams. Verified equal before and after on key sets, values and checkpoint contents, with _hparams_initial now consistent where it previously diverged on four parameters (six with a legacy args namespace). Skipping the inspection means a subclass's own hyperparameters are no longer collected, so ProbabilisticForecasterModule records eval_ensemble_size itself; a second call merges rather than replaces. Note this was only ever saved as a side effect of writing out its signature in ade1929, so keeping it is a deliberate choice -- it changes evaluation results and should round-trip through a checkpoint.
Contributor
Author
|
@joeloskarsson I aasure you, I do not use AI for responses on github however I agree that i use it to code recently, I will reduce that. I will , from now on, take a special interest in the responses, I wanted it to be as descriptive as possible and that would make them winding, I will work on that ! 👍🏻 Thank you ! |
DeterministicForecaster required subclasses to call _configure_scoring to set up the scoring rule and buffers, which are its own concern. Give it a real __init__ instead and have every forecaster mix-in consume the keyword arguments it owns and forward the rest along the MRO. datastore is taken by Forecaster, since both mix-ins need it. Arguments no mix-in claims reach nn.Module and raise rather than being silently dropped. Also drop the one cross-mix-in read at construction time: the per_var_std fallback was allocated based on predicts_std, which the mix-in supplying forward answers, so the two only composed with the objective listed first. Register it whenever a config is given and let _resolve_pred_std settle per call whether it is used, so the mix-ins compose in either order.
…core Drop the metric argument from Forecaster.score, leaving it to apply the forecaster's own scoring rule. Compute mse and mae in DeterministicForecasterModule directly from neural_lam.metrics, as ProbabilisticForecasterModule already does. Values are unchanged: both metrics replace the std argument with ones internally.
The per-step losses it returns are the ones compute_step_losses produces for this objective, gradients included, so all three steps can use it.
The class-level annotations declared a contract nothing checked. Subclasses already create both dicts, and the docstrings on validation_step and test_step say so.
Split the reduction and logging into their own method so the decorator can gate it. The all_gather stays in the caller, which every rank must reach.
| std_placeholder = torch.ones( | ||
| target_states.shape[-1], device=target_states.device | ||
| ) | ||
| entry_mses = metrics.mse( |
Collaborator
There was a problem hiding this comment.
Can we figure out a better solution for this? I would like to just pass it as None, but I don't think the metric allows it? Even though it is unused.
Collaborator
There was a problem hiding this comment.
But it points to a design problem we should resolve rather than working around.
Contributor
Author
There was a problem hiding this comment.
Metrics now declare it: metrics.requires_pred_std(metric). _resolve_pred_std uses that to return None when the configured loss ignores the std, so a forecaster with --loss mse no longer needs a config/per_var_std it would only discard, the ValueError is now specific to scoring rules that actually use one.
Collaborator
There was a problem hiding this comment.
The modules got a lot cleaner now I think, without the placeholder stds.
Why does the other metrics than mse and mae have to use pred_std: Optional[torch.Tensor]? Now this signals that it can be None, but that will always cause a crash. Does not seem like intuitive design. This does to some sense mean that the different metrics have different signature (in that pred_std is optional only for some), but I think this is still the better solution, as the signature where pred_std is explicitly given is still the consistent one.
Collaborator
|
@Sir-Sloth-The-Lazy I had another look over this, commented on and resolved the open threads. Thanks for explaining the parts where I am missing/forgetting things 🙏 I am maybe prioritizing speed/review-frequency over digging that deep into some things, but I think that might still be the most efficient way for us to iterate on this :) |
DeterministicForecaster takes a fixed signature again; ARForecaster alone forwards **kwargs and is listed first in the bases. Trim the docstrings and comments that narrated the mix-in mechanics.
The deterministic module gets its per-step breakdown from score() since 3ca990d, leaving compute_step_losses with no caller outside the class. Forecaster now exposes two loss entry points: compute_training_loss, which produces its own forecast and returns the finished objective, and score, which applies the same rule to a forecast the caller already has.
The all_gather_cat has to run on every rank, so the decorator could not gate the whole method and splitting it in two just to apply it was worse than the guard it replaced.
Summing the ensemble-mean MSE over variables reported a quantity with no
clear meaning, and the per-lead per-variable path through
aggregate_and_plot_metrics already covers the metric properly. Drop it and
log compute_training_loss as {phase}_mean_loss instead, which restores the
scalar ModelCheckpoint monitors and keeps it the model's own objective.
Nothing monitors a test-phase loss, so computing it there spent a forward pass per batch for a number no one reads.
mse and mae documented pred_std as unused but implemented themselves through wmse and wmae with torch.ones_like(pred_std), so the argument was load-bearing for its shape and callers holding no std had to fabricate one just to be ignored. Compute them directly instead and default pred_std to None across every metric, with the std-dependent ones raising a clear ValueError rather than failing inside the maths. Metrics now also declare the requirement through requires_pred_std, so DeterministicForecaster resolves the per_var_std fallback only for a scoring rule that uses it and an unweighted loss needs no config.
joeloskarsson
requested changes
Aug 9, 2026
joeloskarsson
left a comment
joeloskarsson
left a comment
Collaborator
There was a problem hiding this comment.
Looking good! This is pretty close to being able to merge now, just some small polishing :) I think the overall design should be good.
| ) | ||
| return self.per_var_std | ||
|
|
||
| def score( |
Collaborator
There was a problem hiding this comment.
As discussed in #700 (comment), let's rename this to something like compute_loss_from_forecast.
Comment on lines
+221
to
+222
| self._log_objective(batch) | ||
| entry_mses = self._ensemble_step(batch) |
Collaborator
There was a problem hiding this comment.
Just some extra documentation of this:
Suggested change
| self._log_objective(batch) | |
| entry_mses = self._ensemble_step(batch) | |
| # Note that we here do two forward passes: One for computing loss | |
| # and one for computing ensemble metrics. Required as computing loss | |
| # might not involve making a forecast the same way as during inference. | |
| self._log_objective(batch) | |
| entry_mses = self._ensemble_step(batch) |
Comment on lines
+240
to
+241
| entry_mses = self._ensemble_step(batch) | ||
| self.test_metrics["ens_mse"].append(entry_mses) |
Collaborator
There was a problem hiding this comment.
Suggested change
| entry_mses = self._ensemble_step(batch) | |
| self.test_metrics["ens_mse"].append(entry_mses) | |
| # Note that we here do two forward passes: One for computing loss | |
| # and one for computing ensemble metrics. Required as computing loss | |
| # might not involve making a forecast the same way as during inference. | |
| entry_mses = self._ensemble_step(batch) | |
| self.test_metrics["ens_mse"].append(entry_mses) |
| std_placeholder = torch.ones( | ||
| target_states.shape[-1], device=target_states.device | ||
| ) | ||
| entry_mses = metrics.mse( |
Collaborator
There was a problem hiding this comment.
The modules got a lot cleaner now I think, without the placeholder stds.
Why does the other metrics than mse and mae have to use pred_std: Optional[torch.Tensor]? Now this signals that it can be None, but that will always cause a crash. Does not seem like intuitive design. This does to some sense mean that the different metrics have different signature (in that pred_std is optional only for some), but I think this is still the better solution, as the signature where pred_std is explicitly given is still the consistent one.
| "train_steps_to_log": train_steps_to_log, | ||
| "metrics_watch": metrics_watch, | ||
| "var_leads_metrics_watch": var_leads_metrics_watch, | ||
| "args": args, |
Collaborator
There was a problem hiding this comment.
Should args be saved here? Thinking that if we had a legacy args input, would the other hyperparameters then had been updated from it, so no need to keep storing it?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
ForecasterModuleonto theForecaster. A new abstractForecaster.compute_training_lossreturns a finished(loss, loss_components)pair, so each forecaster owns its complete training objective (including assembling it from any internal terms).ForecasterModule.training_stepnow only injects the configured scoring rule (--loss), the interior mask andper_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
ARForecastertraining 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 onForecaster, any out-of-treeForecastersubclass now has to implement it; all in-repo forecasters go throughARForecaster.Add the probabilistic side, built on top of the deterministic one (no sample dimension leaks into deterministic components):
ProbabilisticForecaster(abstract): declaressample_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 overridecompute_training_loss.ProbabilisticForecasterModule: training inherited unchanged; validation samples aneval_ensemble_sizeensemble 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 theval_mean_loss/val_loss_unroll{i}log keys so existing checkpoint callbacks work unchanged.test_stepraisesNotImplementedErrorrather 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 ownkl_beta/crps_weightconfig), ensemble test evaluation + plotting, andtrain_model.pywiring forprobabilistic models.
Issue Link
Graph-EFM instantiation and ensemble evaluation follow in later PRs).
prob_model_lamonmain, see issue Merge Graph-EFM model fromprob_model_lambranch #62Type 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