From 1713911ea38c1cc121deed3948fa16798b1c8915 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sun, 5 Jul 2026 09:31:42 +0530 Subject: [PATCH 1/9] feat: move training loss onto Forecaster, add probabilistic interface 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 #685. --- CHANGELOG.md | 17 + neural_lam/models/__init__.py | 5 + neural_lam/models/forecasters/__init__.py | 1 + .../models/forecasters/autoregressive.py | 73 +++++ neural_lam/models/forecasters/base.py | 64 ++++ .../models/forecasters/probabilistic.py | 257 +++++++++++++++ neural_lam/models/module.py | 30 +- neural_lam/models/probabilistic_module.py | 140 ++++++++ tests/test_probabilistic_forecaster.py | 304 ++++++++++++++++++ 9 files changed, 878 insertions(+), 13 deletions(-) create mode 100644 neural_lam/models/forecasters/probabilistic.py create mode 100644 neural_lam/models/probabilistic_module.py create mode 100644 tests/test_probabilistic_forecaster.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fd123dee..d5ab6ee8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add a general probabilistic forecasting interface: an abstract + `ProbabilisticForecaster` capable of sampling ensemble forecasts + (`sample_ensemble`, members stacked along a new dimension after batch), + its auto-regressive implementation `ProbabilisticARForecaster` (samples + independent trajectories through a stochastic step predictor and by + default trains on the configured scoring rule applied to the ensemble + mean) and a `ProbabilisticForecasterModule` whose validation samples an + ensemble and logs the RMSE of the ensemble mean. Move ownership of the + training objective from `ForecasterModule` onto the `Forecaster`: the + new abstract `Forecaster.compute_training_loss` returns a finished + `(loss, loss_components)` pair and `ForecasterModule.training_step` only + injects the configured scoring rule and interior mask and logs the + result. The deterministic `ARForecaster` training loss is unchanged in + value, only computed by the forecaster itself. + [\#685](https://github.com/mllam/neural-lam/issues/685) + @Sir-Sloth-The-Lazy + - Add `PropagationNet` GNN layer that incentivises directional message propagation from sender to receiver nodes, and expose it alongside `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`, diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..1bfe9eb5 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,7 +3,12 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.probabilistic import ( + ProbabilisticARForecaster, + ProbabilisticForecaster, +) from .module import ForecasterModule +from .probabilistic_module import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_lam import GraphLAM diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 7ea9f6fd..254c4ba0 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -5,3 +5,4 @@ # Local from .autoregressive import ARForecaster from .base import Forecaster +from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index a8135f62..92a4e938 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,6 +1,7 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" # Standard library +from typing import Callable # Third-party import torch @@ -144,3 +145,75 @@ def forward( pred_std = None return prediction, pred_std + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score the deterministic rollout with the injected scoring rule. + + Unrolls a single forecast over the full rollout, scores it against + the target states on interior nodes and averages over batch and + time. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during the rollout. + Dims: same as the prediction. + score_fn : Callable + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the wrapped predictor does not + output an std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the rollout, averaged over + batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; the deterministic objective has no separate components. + """ + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + if pred_std is None: + pred_std = per_var_std + + batch_loss = torch.mean( + score_fn( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4d957916..4b50c030 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,6 +2,7 @@ # Standard library from abc import ABC, abstractmethod +from typing import Callable # Third-party import torch @@ -79,3 +80,66 @@ def forward( per-variable std is substituted upstream by ``ForecasterModule``. Dims: same as ``prediction``. """ + + @abstractmethod + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the training objective for one batch. + + The forecaster owns its complete training objective: which forecasts + to produce from the batch, which loss terms to compute from them and + how to combine those terms into a single scalar. The wrapping + ``ForecasterModule`` only injects the configured scoring rule and + mask, logs the returned components and optimizes the returned loss. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during forecasting. + Dims: same as the prediction. + score_fn : Callable + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the forecaster does not predict its + own std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The full training loss for the batch, to take gradients + of. + loss_components : dict of {str: torch.Tensor} + Scalar loss-related quantities to log alongside the loss, keyed + by component name. The wrapping module prefixes the names with + the training phase. Empty when the objective has no separate + components worth logging. + """ diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..c6609b7d --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,257 @@ +"""Forecasters producing probabilistic (ensemble) forecasts.""" + +# Standard library +from abc import abstractmethod +from typing import Callable + +# Third-party +import torch + +# Local +from ...datastore import BaseDatastore +from ..step_predictors.base import StepPredictor +from .autoregressive import ARForecaster +from .base import Forecaster + + +class ProbabilisticForecaster(Forecaster): + """ + Forecaster whose forecasts are samples from a predictive distribution. + + Adds the capability that probabilistic evaluation and ensemble-based + objectives build on: sampling an ensemble of forecasts. How the + members are produced (auto-regressive sampling, diffusion, ...) is + left to subclasses; consumers only rely on the shape of the returned + ensemble. + """ + + @abstractmethod + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each + predicted step, identically in every member. Dims: same as one + member. + num_members : int or None + Number of ensemble members ``S`` to sample. When ``None``, the + forecaster's configured ensemble size is used. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + ensemble_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` + when the forecaster predicts an std, otherwise ``None``. Dims: + same as ``ensemble``. + """ + + +class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): + """ + Auto-regressive forecaster for step predictors that sample their output. + + Each call to the wrapped predictor draws a fresh sample of the next + state, so the inherited ``ARForecaster.forward`` unrolls one sampled + 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 + (ensemble scoring rules, variational objectives) override + ``compute_training_loss``. + """ + + def __init__( + self, + predictor: StepPredictor, + datastore: BaseDatastore, + ensemble_size: int, + ) -> None: + """ + Initialize the ProbabilisticARForecaster. + + Parameters + ---------- + predictor : StepPredictor + The predictor to use for each step. Each call should draw a + fresh sample of the next state. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + ensemble_size : int + Number of ensemble members to sample when no explicit member + count is given, in particular for the training objective. + """ + super().__init__(predictor, datastore) + if ensemble_size < 1: + raise ValueError( + f"ensemble_size must be at least 1, got {ensemble_size}" + ) + self.ensemble_size = ensemble_size + + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Unrolls ``num_members`` independent forecasts, each sampling fresh + randomness at every step, and stacks them along a new ensemble + dimension after the batch dimension. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each AR + step, identically in every member. Dims: same as one member. + num_members : int or None + Number of ensemble members ``S`` to sample. When ``None``, + ``self.ensemble_size`` is used. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + ensemble_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` + when the wrapped predictor outputs an std, otherwise ``None``. + Dims: same as ``ensemble``. + """ + if num_members is None: + num_members = self.ensemble_size + + member_list = [] + member_std_list = [] + for _ in range(num_members): + prediction, pred_std = self( + init_states, forcing_features, boundary_states + ) + member_list.append(prediction) + if pred_std is not None: + member_std_list.append(pred_std) + + ensemble = torch.stack(member_list, dim=1) + ensemble_std = ( + torch.stack(member_std_list, dim=1) if member_std_list else None + ) + return ensemble, ensemble_std + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score the ensemble mean with the injected scoring rule. + + Samples an ensemble of ``self.ensemble_size`` forecasts, averages + the members into an ensemble mean forecast, scores it against the + target states on interior nodes and averages over batch and time. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + 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 + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the wrapped predictor does not + output an std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the ensemble mean, averaged + over batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; this objective has no separate components. + """ + ensemble, ensemble_std = self.sample_ensemble( + init_states, forcing_features, target_states + ) + ensemble_mean = ensemble.mean(dim=1) + if ensemble_std is not None: + pred_std = ensemble_std.mean(dim=1) + else: + pred_std = per_var_std + + batch_loss = torch.mean( + score_fn( + ensemble_mean, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 71ce7951..62791733 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -338,7 +338,7 @@ def on_after_batch_transfer(self, batch, dataloader_idx): def common_step(self, batch): """ - Perform a common prediction step for training, validation, and testing. + Perform a common prediction step for validation and testing. Parameters ---------- @@ -362,6 +362,10 @@ def training_step(self, batch): """ Perform a single training step. + The training objective is fully assembled by the wrapped forecaster; + this method injects the configured scoring rule and interior mask, + then logs the loss and any loss components the forecaster returns. + Parameters ---------- batch : tuple @@ -372,20 +376,20 @@ def training_step(self, batch): torch.Tensor The computed loss for the training step. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.per_var_std - - batch_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ) + init_states, target_states, forcing_features, _ = batch + batch_loss, loss_components = self.forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=self.loss, + interior_mask_bool=self.interior_mask_bool, + per_var_std=self.per_var_std, ) - log_dict = {"train_loss": batch_loss} + log_dict = { + f"train_{name}": value for name, value in loss_components.items() + } + log_dict["train_loss"] = batch_loss self.log_dict( log_dict, prog_bar=True, diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py new file mode 100644 index 00000000..24f83049 --- /dev/null +++ b/neural_lam/models/probabilistic_module.py @@ -0,0 +1,140 @@ +"""Lightning module evaluating probabilistic forecasters as ensembles.""" + +# Third-party +import torch + +# Local +from .. import metrics +from .forecasters.probabilistic import ProbabilisticForecaster +from .module import ForecasterModule + + +class ProbabilisticForecasterModule(ForecasterModule): + """ + Lightning module for forecasters that sample ensemble forecasts. + + Training is inherited unchanged from ``ForecasterModule``: the wrapped + forecaster assembles its own training loss. Validation is ensemble + based instead of deterministic: an ensemble is sampled from the + forecaster and scored through its ensemble mean (root-mean-squared + error of the ensemble mean). The module only assumes that the + forecaster can sample ensemble forecasts of the correct shape; it makes + no assumption on how the members are produced. + """ + + # The wrapped forecaster must be able to sample ensemble forecasts + forecaster: ProbabilisticForecaster + + def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): + """ + Initialize the module and store the evaluation ensemble size. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``ForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + eval_ensemble_size : int or None + Number of ensemble members sampled during validation. ``None`` + uses the forecaster's configured ensemble size. + **kwargs + Keyword arguments forwarded to ``ForecasterModule.__init__`` + (``loss``, ``lr``, ...). + """ + super().__init__(*args, **kwargs) + if eval_ensemble_size is not None and eval_ensemble_size < 1: + raise ValueError( + "eval_ensemble_size must be at least 1, " + f"got {eval_ensemble_size}" + ) + self.eval_ensemble_size = eval_ensemble_size + self.val_metrics = {"ens_mse": []} + + def validation_step(self, batch, batch_idx): + """ + Perform a single ensemble validation step. + + Samples an ensemble from the forecaster and scores its ensemble + mean against the target states on interior nodes. Logs the + root-mean-squared error of the ensemble mean per configured rollout + step and averaged over the rollout, and collects per-variable + ensemble-mean MSE for epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + init_states, target_states, forcing_features, _ = batch + ensemble, _ = self.forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=self.eval_ensemble_size, + ) + ensemble_mean = ensemble.mean(dim=1) + # metrics.mse ignores the std argument, but requires one + std_placeholder = torch.ones( + target_states.shape[-1], device=target_states.device + ) + + time_step_mse = torch.mean( + metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + ), + dim=0, + ) + time_step_rmse = torch.sqrt(time_step_mse) + mean_rmse = torch.mean(time_step_rmse) + self._warn_skipped_val_steps(len(time_step_rmse), "val") + + val_log_dict = { + f"val_loss_unroll{step}": time_step_rmse[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_rmse) + } + val_log_dict["val_mean_loss"] = mean_rmse + self.log_dict( + val_log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + entry_mses = metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.val_metrics["ens_mse"].append(entry_mses) + + def test_step(self, batch, batch_idx): + """ + Not supported: ensemble test evaluation is not implemented. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + + Raises + ------ + NotImplementedError + Always; only training and ensemble validation are implemented + for probabilistic forecasters. + """ + raise NotImplementedError( + "Ensemble test evaluation is not implemented for " + "probabilistic forecasters." + ) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py new file mode 100644 index 00000000..34458c13 --- /dev/null +++ b/tests/test_probabilistic_forecaster.py @@ -0,0 +1,304 @@ +# Third-party +import pytest +import torch +from torch import nn + +# First-party +from neural_lam import config as nlconfig +from neural_lam import metrics +from neural_lam.models import ( + ARForecaster, + ForecasterModule, + ProbabilisticARForecaster, + ProbabilisticForecasterModule, + StepPredictor, +) +from tests.conftest import init_datastore_example + + +class ZeroStepPredictor(StepPredictor): + """Deterministic predictor always predicting the zero state.""" + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = torch.zeros_like(prev_state) + pred_std = torch.zeros_like(prev_state) if self.output_std else None + return pred_state, pred_std + + +class NoisyStepPredictor(StepPredictor): + """Stochastic predictor sampling a fresh state at every call.""" + + def __init__(self, datastore, **kwargs): + super().__init__(datastore, **kwargs) + self.noise_scale = nn.Parameter(torch.tensor(1.0)) + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = self.noise_scale * torch.randn_like(prev_state) + return pred_state, None + + +def _example_batch(datastore, B=2, pred_steps=3): + """Create constant example input tensors matching the datastore dims.""" + num_grid_nodes = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + num_past_forcing_steps = 1 + num_future_forcing_steps = 1 + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + target_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) * 5.0 + return init_states, forcing_features, target_states + + +def test_ar_forecaster_training_loss_matches_direct_score(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore) + + init_states, forcing_features, target_states = _example_batch(datastore) + score_fn = metrics.get_metric("mse") + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + per_var_std = torch.ones(d_state) + + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=score_fn, + interior_mask_bool=interior_mask_bool, + per_var_std=per_var_std, + ) + + prediction, _ = forecaster(init_states, forcing_features, target_states) + expected_loss = torch.mean( + score_fn( + prediction, + target_states, + per_var_std, + mask=interior_mask_bool, + ) + ) + + assert batch_loss.shape == () + assert loss_components == {} + torch.testing.assert_close(batch_loss, expected_loss) + + +def test_sample_ensemble_shapes_and_member_variability(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + + # Override masks to test boundary masking behaviour + forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) + forecaster.interior_mask[0, 0] = 1 # One node is interior + forecaster.boundary_mask = 1 - forecaster.interior_mask + + B, pred_steps, num_members = 2, 3, 4 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + num_grid_nodes = datastore.num_grid_points + d_state = target_states.shape[-1] + + torch.manual_seed(42) + ensemble, ensemble_std = forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=num_members, + ) + + assert ensemble.shape == ( + B, + num_members, + pred_steps, + num_grid_nodes, + d_state, + ) + assert ensemble_std is None + + # Members carry independent samples on the interior node + assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0]) + # Boundary nodes are overwritten with the true state in every member + assert torch.all(ensemble[:, :, :, 1:] == 5.0) + + # Without an explicit member count the configured ensemble_size is used + default_ensemble, _ = forecaster.sample_ensemble( + init_states, forcing_features, target_states + ) + assert default_ensemble.shape[1] == forecaster.ensemble_size + + +def test_probabilistic_training_loss_gradient_flow(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + + torch.manual_seed(42) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=metrics.get_metric("mse"), + interior_mask_bool=interior_mask_bool, + per_var_std=torch.ones(d_state), + ) + + assert batch_loss.shape == () + assert loss_components == {} + assert torch.isfinite(batch_loss) + + batch_loss.backward() + assert predictor.noise_scale.grad is not None + assert predictor.noise_scale.grad != 0.0 + + +def test_probabilistic_forecaster_rejects_empty_ensemble(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(ValueError, match="ensemble_size"): + ProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + + +def test_module_training_step_delegates_to_forecaster(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + batch = (init_states, target_states, forcing_features, batch_times) + + batch_loss = model.training_step(batch) + + expected_loss, _ = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=model.loss, + interior_mask_bool=model.interior_mask_bool, + per_var_std=model.per_var_std, + ) + + torch.testing.assert_close(batch_loss, expected_loss) + + +class MemberCountRecordingForecaster(ProbabilisticARForecaster): + """ProbabilisticARForecaster recording the requested member count.""" + + def sample_ensemble(self, *args, **kwargs): + self.last_num_members = kwargs.get("num_members") + return super().sample_ensemble(*args, **kwargs) + + +def test_probabilistic_module_validation_scores_ensemble_mean(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, ensemble_size=2 + ) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + eval_ensemble_size=3, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(42) + model.validation_step(batch, 0) + + # Validation samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.val_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) + + +def test_probabilistic_module_rejects_empty_eval_ensemble(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + + with pytest.raises(ValueError, match="eval_ensemble_size"): + ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + eval_ensemble_size=0, + ) + + +def test_probabilistic_module_test_step_not_implemented(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + batch = (init_states, target_states, forcing_features, batch_times) + + with pytest.raises(NotImplementedError): + model.test_step(batch, 0) From 3f5402d9b57fab41394522620cd8c15c53065b24 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Mon, 6 Jul 2026 13:32:10 +0530 Subject: [PATCH 2/9] Address PR review: clarify scoring-rule wording, rename score_fn to score_metric --- neural_lam/models/forecasters/autoregressive.py | 12 ++++++------ neural_lam/models/forecasters/base.py | 8 ++++---- neural_lam/models/forecasters/probabilistic.py | 17 +++++++++-------- neural_lam/models/module.py | 2 +- tests/test_probabilistic_forecaster.py | 10 +++++----- 5 files changed, 25 insertions(+), 24 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 92a4e938..9a2980c6 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -151,12 +151,12 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the deterministic rollout with the injected scoring rule. + Score the deterministic rollout with the given ``score_metric``. Unrolls a single forecast over the full rollout, scores it against the target states on interior nodes and averages over batch and @@ -182,12 +182,12 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollout. Dims: same as the prediction. - score_fn : Callable + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard @@ -209,7 +209,7 @@ def compute_training_loss( pred_std = per_var_std batch_loss = torch.mean( - score_fn( + score_metric( prediction, target_states, pred_std, diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4b50c030..ad32de00 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -87,7 +87,7 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: @@ -120,12 +120,12 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during forecasting. Dims: same as the prediction. - score_fn : Callable + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index c6609b7d..0965a4a4 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -81,8 +81,9 @@ class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): state, so the inherited ``ARForecaster.forward`` unrolls one sampled 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 + The default training objective scores the ensemble mean using the + scoring rule passed to ``compute_training_loss`` (from + ``neural_lam.metrics``); forecasters with model-specific objectives (ensemble scoring rules, variational objectives) override ``compute_training_loss``. """ @@ -186,12 +187,12 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with the injected scoring rule. + Score the ensemble mean with the given ``score_metric``. Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the @@ -217,12 +218,12 @@ def compute_training_loss( 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 + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard @@ -247,7 +248,7 @@ def compute_training_loss( pred_std = per_var_std batch_loss = torch.mean( - score_fn( + score_metric( ensemble_mean, target_states, pred_std, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 62791733..a4ca4d8d 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -381,7 +381,7 @@ def training_step(self, batch): init_states, forcing_features, target_states, - score_fn=self.loss, + score_metric=self.loss, interior_mask_bool=self.interior_mask_bool, per_var_std=self.per_var_std, ) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 34458c13..d22c68dc 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -58,7 +58,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): forecaster = ARForecaster(predictor, datastore) init_states, forcing_features, target_states = _example_batch(datastore) - score_fn = metrics.get_metric("mse") + score_metric = metrics.get_metric("mse") interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] per_var_std = torch.ones(d_state) @@ -67,14 +67,14 @@ def test_ar_forecaster_training_loss_matches_direct_score(): init_states, forcing_features, target_states, - score_fn=score_fn, + score_metric=score_metric, interior_mask_bool=interior_mask_bool, per_var_std=per_var_std, ) prediction, _ = forecaster(init_states, forcing_features, target_states) expected_loss = torch.mean( - score_fn( + score_metric( prediction, target_states, per_var_std, @@ -151,7 +151,7 @@ def test_probabilistic_training_loss_gradient_flow(): init_states, forcing_features, target_states, - score_fn=metrics.get_metric("mse"), + score_metric=metrics.get_metric("mse"), interior_mask_bool=interior_mask_bool, per_var_std=torch.ones(d_state), ) @@ -200,7 +200,7 @@ def test_module_training_step_delegates_to_forecaster(): init_states, forcing_features, target_states, - score_fn=model.loss, + score_metric=model.loss, interior_mask_bool=model.interior_mask_bool, per_var_std=model.per_var_std, ) From 987fecd7fc0bc6d5399754a8ffa3fbf0379f5301 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 09:10:48 +0530 Subject: [PATCH 3/9] Address PR review: move loss and per_var_std onto the Forecaster 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 #685 down to one sentence per review feedback. --- CHANGELOG.md | 18 +---- .../models/forecasters/autoregressive.py | 64 ++++++++++----- neural_lam/models/forecasters/base.py | 19 ++--- .../models/forecasters/probabilistic.py | 31 +++---- neural_lam/models/module.py | 80 +++++++------------ neural_lam/models/probabilistic_module.py | 2 +- neural_lam/train_model.py | 9 ++- tests/test_checkpoint.py | 3 +- tests/test_datasets.py | 5 +- tests/test_gnn_layers.py | 2 +- tests/test_gpu_normalization.py | 2 +- tests/test_plotting.py | 8 +- tests/test_prediction_model_classes.py | 15 ++-- tests/test_probabilistic_forecaster.py | 25 +++--- tests/test_training.py | 3 +- 15 files changed, 135 insertions(+), 151 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d5ab6ee8..2809daaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,20 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a general probabilistic forecasting interface: an abstract - `ProbabilisticForecaster` capable of sampling ensemble forecasts - (`sample_ensemble`, members stacked along a new dimension after batch), - its auto-regressive implementation `ProbabilisticARForecaster` (samples - independent trajectories through a stochastic step predictor and by - default trains on the configured scoring rule applied to the ensemble - mean) and a `ProbabilisticForecasterModule` whose validation samples an - ensemble and logs the RMSE of the ensemble mean. Move ownership of the - training objective from `ForecasterModule` onto the `Forecaster`: the - new abstract `Forecaster.compute_training_loss` returns a finished - `(loss, loss_components)` pair and `ForecasterModule.training_step` only - injects the configured scoring rule and interior mask and logs the - result. The deterministic `ARForecaster` training loss is unchanged in - value, only computed by the forecaster itself. +- Add a general probabilistic forecasting interface (`ProbabilisticForecaster`, + `ProbabilisticARForecaster`, `ProbabilisticForecasterModule`) and move + ownership of the training objective, scoring rule and per-variable std + from `ForecasterModule` onto the `Forecaster`. [\#685](https://github.com/mllam/neural-lam/issues/685) @Sir-Sloth-The-Lazy diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 9a2980c6..74f2dc77 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,13 +1,13 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" -# Standard library -from typing import Callable - # Third-party import torch # Local +from ... import metrics +from ...config import NeuralLAMConfig from ...datastore import BaseDatastore +from ...loss_weighting import get_state_feature_weighting from ..step_predictors.base import StepPredictor from .base import Forecaster @@ -19,7 +19,11 @@ class ARForecaster(Forecaster): """ def __init__( - self, predictor: StepPredictor, datastore: BaseDatastore + self, + predictor: StepPredictor, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", ) -> None: """ Initialize the ARForecaster. @@ -30,6 +34,14 @@ def __init__( The predictor to use for each step. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Only required for that case; + forecasters used purely for inference can omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. """ super().__init__() self.predictor = predictor @@ -45,6 +57,31 @@ def __init__( "interior_mask", 1.0 - self.boundary_mask, persistent=False ) + self.loss = metrics.get_metric(loss) + + # Store per_var_std here if the predictor does not output its own std + if not self.predicts_std and config is not None: + da_state_stats = datastore.get_standardization_dataarray( + category="state" + ) + state_feature_weights = get_state_feature_weighting( + config=config, datastore=datastore + ) + diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + feature_weights_t = torch.tensor( + state_feature_weights, dtype=torch.float32 + ) + self.register_buffer( + "per_var_std", + diff_std / torch.sqrt(feature_weights_t), + persistent=False, + ) + else: + self.per_var_std = None + @property def predicts_std(self) -> bool: """ @@ -138,7 +175,7 @@ def forward( prediction = torch.stack(prediction_list, dim=1) # If predictor outputs std, stack it; otherwise return None so - # ForecasterModule can substitute the constant per_var_std + # callers can substitute the constant per_var_std if pred_std_list: pred_std = torch.stack(pred_std_list, dim=1) else: @@ -151,12 +188,10 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the deterministic rollout with the given ``score_metric``. + Score the deterministic rollout with ``self.loss``. Unrolls a single forecast over the full rollout, scores it against the target states on interior nodes and averages over batch and @@ -182,17 +217,10 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollout. Dims: same as the prediction. - score_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the wrapped predictor does not - output an std, otherwise ``None``. Returns ------- @@ -206,10 +234,10 @@ def compute_training_loss( init_states, forcing_features, target_states ) if pred_std is None: - pred_std = per_var_std + pred_std = self.per_var_std batch_loss = torch.mean( - score_metric( + self.loss( prediction, target_states, pred_std, diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index ad32de00..8dfb0002 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,7 +2,6 @@ # Standard library from abc import ABC, abstractmethod -from typing import Callable # Third-party import torch @@ -87,18 +86,17 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ Compute the training objective for one batch. The forecaster owns its complete training objective: which forecasts to produce from the batch, which loss terms to compute from them and - how to combine those terms into a single scalar. The wrapping - ``ForecasterModule`` only injects the configured scoring rule and - mask, logs the returned components and optimizes the returned loss. + how to combine those terms into a single scalar, using its own + ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The + wrapping ``ForecasterModule`` only injects the interior mask, logs + the returned components and optimizes the returned loss. Parameters ---------- @@ -120,17 +118,10 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during forecasting. Dims: same as the prediction. - score_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the forecaster does not predict its - own std, otherwise ``None``. Returns ------- diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 0965a4a4..2bfd0458 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -2,12 +2,12 @@ # Standard library from abc import abstractmethod -from typing import Callable # Third-party import torch # Local +from ...config import NeuralLAMConfig from ...datastore import BaseDatastore from ..step_predictors.base import StepPredictor from .autoregressive import ARForecaster @@ -93,6 +93,8 @@ def __init__( predictor: StepPredictor, datastore: BaseDatastore, ensemble_size: int, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", ) -> None: """ Initialize the ProbabilisticARForecaster. @@ -107,8 +109,16 @@ def __init__( ensemble_size : int Number of ensemble members to sample when no explicit member count is given, in particular for the training objective. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Only required for that case; + forecasters used purely for inference can omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. """ - super().__init__(predictor, datastore) + super().__init__(predictor, datastore, config=config, loss=loss) if ensemble_size < 1: raise ValueError( f"ensemble_size must be at least 1, got {ensemble_size}" @@ -187,12 +197,10 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with the given ``score_metric``. + Score the ensemble mean with ``self.loss``. Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the @@ -218,17 +226,10 @@ def compute_training_loss( 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_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the wrapped predictor does not - output an std, otherwise ``None``. Returns ------- @@ -245,10 +246,10 @@ def compute_training_loss( if ensemble_std is not None: pred_std = ensemble_std.mean(dim=1) else: - pred_std = per_var_std + pred_std = self.per_var_std batch_loss = torch.mean( - score_metric( + self.loss( ensemble_mean, target_states, pred_std, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index a4ca4d8d..4cebe6ca 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -20,7 +20,6 @@ from .. import metrics, vis from ..config import NeuralLAMConfig from ..datastore import BaseDatastore -from ..loss_weighting import get_state_feature_weighting from ..weather_dataset import WeatherDataset from .forecasters.base import Forecaster @@ -38,7 +37,6 @@ def __init__( forecaster: Forecaster, config: NeuralLAMConfig, datastore: BaseDatastore, - loss: str = "wmse", lr: float = 1e-3, restore_opt: bool = False, n_example_pred: int = 1, @@ -54,13 +52,14 @@ def __init__( Parameters ---------- forecaster : Forecaster - The forecaster model to use for predictions. + The forecaster model to use for predictions. Owns the scoring + rule (``forecaster.loss``) and the constant per-variable std + fallback (``forecaster.per_var_std``) used for training and for + validation/test loss reporting here. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore Datastore providing grid metadata and data access. - loss : str, default "wmse" - The loss function to use. lr : float, default 1e-3 Learning rate for the optimizer. restore_opt : bool, default False @@ -79,7 +78,7 @@ def __init__( args : argparse.Namespace, optional Pre-refactor ``ARModel`` checkpoint hyperparameters. When provided, attributes on ``args`` take precedence over the - corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``, + corresponding explicit kwargs (``lr``, ``restore_opt``, ``n_example_pred``, ``create_gif``, ``val_steps_to_log``, ``metrics_watch``, ``var_leads_metrics_watch``) so legacy checkpoints round-trip through ``load_from_checkpoint`` @@ -90,10 +89,9 @@ def __init__( # inside an argparse Namespace under the single key 'args'. When # Lightning calls __init__ during load_from_checkpoint it would # otherwise drop 'args' (not in the new signature) and silently fall - # back to defaults for loss/lr/create_gif/etc. Unpack the namespace - # here so legacy checkpoints round-trip correctly. + # back to defaults for lr/create_gif/etc. Unpack the namespace here + # so legacy checkpoints round-trip correctly. if args is not None: - loss = getattr(args, "loss", loss) lr = getattr(args, "lr", lr) restore_opt = getattr(args, "restore_opt", restore_opt) n_example_pred = getattr(args, "n_example_pred", n_example_pred) @@ -139,29 +137,6 @@ def __init__( persistent=False, ) - # Store per_var_std here if predictor does not output std - if not self.forecaster.predicts_std: - da_state_stats = datastore.get_standardization_dataarray( - category="state" - ) - state_feature_weights = get_state_feature_weighting( - config=config, datastore=datastore - ) - diff_std = torch.tensor( - da_state_stats.state_diff_std_standardized.values, - dtype=torch.float32, - ) - feature_weights_t = torch.tensor( - state_feature_weights, dtype=torch.float32 - ) - self.register_buffer( - "per_var_std", - diff_std / torch.sqrt(feature_weights_t), - persistent=False, - ) - else: - self.per_var_std = None - # Standardization statistics used to normalize each batch on-device in # `on_after_batch_transfer`. WeatherDataset returns unstandardized # data, so state and forcing are normalized here rather than on CPU. @@ -207,9 +182,6 @@ def __init__( self.forcing_mean = None self.forcing_std = None - # Instantiate loss function - self.loss = metrics.get_metric(loss) - self.val_metrics: dict[str, list] = { "mse": [], } @@ -362,9 +334,10 @@ def training_step(self, batch): """ Perform a single training step. - The training objective is fully assembled by the wrapped forecaster; - this method injects the configured scoring rule and interior mask, - then logs the loss and any loss components the forecaster returns. + The training objective is fully assembled by the wrapped forecaster, + which owns its own scoring rule; this method injects the interior + mask, then logs the loss and any loss components the forecaster + returns. Parameters ---------- @@ -381,9 +354,7 @@ def training_step(self, batch): init_states, forcing_features, target_states, - score_metric=self.loss, interior_mask_bool=self.interior_mask_bool, - per_var_std=self.per_var_std, ) log_dict = { @@ -452,10 +423,10 @@ def validation_step(self, batch, batch_idx): """ prediction, target_states, pred_std, _ = self.common_step(batch) if pred_std is None: - pred_std = self.per_var_std + pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.loss( + self.forecaster.loss( prediction, target_states, pred_std, @@ -532,10 +503,10 @@ def test_step(self, batch, batch_idx): self.test_metrics["output_std"].append(mean_pred_std) if pred_std is None: - pred_std = self.per_var_std + pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.loss( + self.forecaster.loss( prediction, target_states, pred_std, @@ -572,7 +543,7 @@ def test_step(self, batch, batch_idx): ) self.test_metrics[metric_name].append(batch_metric_vals) - spatial_loss = self.loss( + spatial_loss = self.forecaster.loss( prediction, target_states, pred_std, average_grid=False ) log_spatial_losses = spatial_loss[ @@ -980,15 +951,20 @@ def on_load_checkpoint(self, checkpoint): # 1. Broad namespace remap: for pre-refactor checkpoints # The old ``ARModel`` was a flat LightningModule. Everything that # belonged to the predictor needs to be moved to - # 'forecaster.predictor.' + # 'forecaster.predictor.', while 'per_var_std' (now owned by the + # forecaster itself) moves to 'forecaster.per_var_std' and + # 'interior_mask_bool' (still owned by the module) stays as-is. old_keys = list(loaded_state_dict.keys()) for key in old_keys: - if not key.startswith("forecaster.") and key not in ( - "interior_mask_bool", - "per_var_std", - ): - new_key = f"forecaster.predictor.{key}" - loaded_state_dict[new_key] = loaded_state_dict.pop(key) + if key.startswith("forecaster.") or key == "interior_mask_bool": + continue + if key == "per_var_std": + loaded_state_dict["forecaster.per_var_std"] = ( + loaded_state_dict.pop(key) + ) + continue + new_key = f"forecaster.predictor.{key}" + loaded_state_dict[new_key] = loaded_state_dict.pop(key) # 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp # Will be under forecaster.predictor due to the remap above, or diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 24f83049..37a756d7 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -40,7 +40,7 @@ def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): uses the forecaster's configured ensemble size. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` - (``loss``, ``lr``, ...). + (``lr``, ...). """ super().__init__(*args, **kwargs) if eval_ensemble_size is not None and eval_ensemble_size < 1: diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index f98065c4..d5e86536 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -63,7 +63,9 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) return ForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, @@ -457,13 +459,14 @@ def main(input_args=None): mesh_up_gnn_type=args.mesh_up_gnn_type, mesh_down_gnn_type=args.mesh_down_gnn_type, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, lr=args.lr, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148..6f114043 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -50,12 +50,11 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, n_example_pred=1, val_steps_to_log=[1], diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 4b35840e..1941206b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -235,13 +235,14 @@ def _create_graph(): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 04c99003..789297d0 100644 --- a/tests/test_gnn_layers.py +++ b/tests/test_gnn_layers.py @@ -73,7 +73,7 @@ def _build_model_and_data( output_clamping_upper=config.training.output_clamping.upper, **gnn_kwargs, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config) B = 2 num_grid_nodes = predictor.num_grid_nodes diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index 8d516bfb..b063d626 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -26,7 +26,7 @@ def _build_module(datastore): ) ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config) return ForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 616d563d..970590be 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -467,13 +467,14 @@ class ModelArgs: output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, @@ -679,12 +680,11 @@ def _build_metrics_watch_module(datastore, config): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") return ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 73e2f905..6abc786c 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -97,13 +97,12 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1e-3, restore_opt=False, n_example_pred=1, @@ -193,8 +192,6 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) - # Use distinctive non-default values so we can detect silent fallback # to ForecasterModule's defaults during load. saved_loss = "mse" @@ -203,11 +200,14 @@ def test_forecaster_module_old_checkpoint(tmp_path): saved_val_steps = [2] saved_n_example_pred = 7 + forecaster = ARForecaster( + predictor, datastore, config=config, loss=saved_loss + ) + model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=saved_loss, lr=saved_lr, restore_opt=False, n_example_pred=saved_n_example_pred, @@ -269,7 +269,9 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss=saved_loss + ) # Load from hacked old checkpoint loaded_model = ForecasterModule.load_from_checkpoint( @@ -284,7 +286,6 @@ def test_forecaster_module_old_checkpoint(tmp_path): # Hyperparameters nested in the legacy 'args' namespace must round-trip # rather than silently falling back to ForecasterModule defaults. - assert loaded_model.hparams.loss == saved_loss assert loaded_model.hparams.lr == saved_lr assert loaded_model.hparams.val_steps_to_log == saved_val_steps assert loaded_model.create_gif is saved_create_gif diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index d22c68dc..1a8f9971 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -55,21 +55,21 @@ def _example_batch(datastore, B=2, pred_steps=3): def test_ar_forecaster_training_loss_matches_direct_score(): datastore = init_datastore_example("mdp") predictor = ZeroStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, loss="mse") init_states, forcing_features, target_states = _example_batch(datastore) score_metric = metrics.get_metric("mse") interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] - per_var_std = torch.ones(d_state) + # per_var_std is normally computed from config; override directly since + # this test only cares about the loss computation, not standardization. + forecaster.per_var_std = torch.ones(d_state) batch_loss, loss_components = forecaster.compute_training_loss( init_states, forcing_features, target_states, - score_metric=score_metric, interior_mask_bool=interior_mask_bool, - per_var_std=per_var_std, ) prediction, _ = forecaster(init_states, forcing_features, target_states) @@ -77,7 +77,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): score_metric( prediction, target_states, - per_var_std, + forecaster.per_var_std, mask=interior_mask_bool, ) ) @@ -139,21 +139,22 @@ def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 + predictor, datastore, ensemble_size=2, loss="mse" ) init_states, forcing_features, target_states = _example_batch(datastore) interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] + # per_var_std is normally computed from config; override directly since + # this test only cares about the loss computation, not standardization. + forecaster.per_var_std = torch.ones(d_state) torch.manual_seed(42) batch_loss, loss_components = forecaster.compute_training_loss( init_states, forcing_features, target_states, - score_metric=metrics.get_metric("mse"), interior_mask_bool=interior_mask_bool, - per_var_std=torch.ones(d_state), ) assert batch_loss.shape == () @@ -176,18 +177,17 @@ def test_probabilistic_forecaster_rejects_empty_ensemble(): def test_module_training_step_delegates_to_forecaster(): datastore = init_datastore_example("mdp") predictor = ZeroStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", ) init_states, forcing_features, target_states = _example_batch(datastore) @@ -200,9 +200,7 @@ def test_module_training_step_delegates_to_forecaster(): init_states, forcing_features, target_states, - score_metric=model.loss, interior_mask_bool=model.interior_mask_bool, - per_var_std=model.per_var_std, ) torch.testing.assert_close(batch_loss, expected_loss) @@ -232,7 +230,6 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", eval_ensemble_size=3, ) @@ -273,7 +270,6 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", eval_ensemble_size=0, ) @@ -293,7 +289,6 @@ def test_probabilistic_module_test_step_not_implemented(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", ) init_states, forcing_features, target_states = _example_batch(datastore) diff --git a/tests/test_training.py b/tests/test_training.py index bf1a5884..589e9d89 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -123,13 +123,12 @@ def run_simple_training( output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, From 511a6d5b493f941230bc1303d859ba38c97ff71f Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 09:47:21 +0530 Subject: [PATCH 4/9] Fail fast when a forecaster is missing per_var_std it needs 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. --- neural_lam/models/module.py | 8 ++++++++ tests/test_prediction_model_classes.py | 4 +++- tests/test_probabilistic_forecaster.py | 18 +++++++++--------- 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 4cebe6ca..e7420b52 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -122,6 +122,14 @@ def __init__( self.save_hyperparameters(ignore=["datastore", "forecaster"]) self.datastore = datastore self.forecaster = forecaster + if forecaster.per_var_std is None and not forecaster.predicts_std: + raise ValueError( + "forecaster.per_var_std is None but the forecaster does " + "not predict its own std (forecaster.predicts_std is " + "False), so training/validation/test scoring has no std " + "to use. Pass config to the forecaster's constructor so " + "it can compute the constant per-variable std." + ) self.matched_metrics: set = set() # Compute interior_mask_bool directly from datastore diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 6abc786c..9bc9d9c0 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -132,7 +132,9 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss="mse" + ) # Load from checkpoint loaded_model = ForecasterModule.load_from_checkpoint( diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 1a8f9971..515b9232 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -217,15 +217,15 @@ def sample_ensemble(self, *args, **kwargs): def test_probabilistic_module_validation_scores_ensemble_mean(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = MemberCountRecordingForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, ensemble_size=2, config=config + ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, @@ -256,14 +256,14 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): def test_probabilistic_module_rejects_empty_eval_ensemble(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2, config=config + ) with pytest.raises(ValueError, match="eval_ensemble_size"): ProbabilisticForecasterModule( @@ -277,14 +277,14 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): def test_probabilistic_module_test_step_not_implemented(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2, config=config + ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, From 56b3d6bd97f5a93b3d2627eefefada529581e644 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 10:47:31 +0530 Subject: [PATCH 5/9] Rename ensemble_std to per_member_std, document mixture semantics 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). --- .../models/forecasters/probabilistic.py | 43 +++++++++++++------ tests/test_probabilistic_forecaster.py | 4 +- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 2bfd0458..cfa6169b 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -23,6 +23,14 @@ class ProbabilisticForecaster(Forecaster): members are produced (auto-regressive sampling, diffusion, ...) is left to subclasses; consumers only rely on the shape of the returned ensemble. + + When ``sample_ensemble`` returns a ``per_member_std``, it is each + member's own predicted std, not a std describing the spread across + members. The predictive distribution is then a mixture of ``S`` + Gaussians, one per member: ``p(x) = mean_s N(x; ensemble[:, s], + per_member_std[:, s]**2)``, not a single Gaussian. In particular, the + variance of that mixture is not the average of the per-member + variances: it also includes the spread between the member means. """ @abstractmethod @@ -66,10 +74,12 @@ def sample_ensemble( Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. The sampled forecasts, stacked along the ensemble dimension ``S``. - ensemble_std : torch.Tensor or None - Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` - when the forecaster predicts an std, otherwise ``None``. Dims: - same as ``ensemble``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the forecaster predicts an std, otherwise + ``None``. Dims: same as ``ensemble``. """ @@ -168,10 +178,12 @@ def sample_ensemble( Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. The sampled forecasts, stacked along the ensemble dimension ``S``. - ensemble_std : torch.Tensor or None - Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` - when the wrapped predictor outputs an std, otherwise ``None``. - Dims: same as ``ensemble``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the wrapped predictor outputs an std, otherwise + ``None``. Dims: same as ``ensemble``. """ if num_members is None: num_members = self.ensemble_size @@ -187,10 +199,10 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) - ensemble_std = ( + per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) - return ensemble, ensemble_std + return ensemble, per_member_std def compute_training_loss( self, @@ -205,6 +217,11 @@ def compute_training_loss( Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the target states on interior nodes and averages over batch and time. + When members predict their own std, the std passed to ``self.loss`` + is the plain average of the per-member stds; this is a + simplification of the true mixture predictive variance, which + would also include the spread between the member means (see the + ``ProbabilisticForecaster`` class docstring). Parameters ---------- @@ -239,12 +256,12 @@ def compute_training_loss( loss_components : dict of {str: torch.Tensor} Empty; this objective has no separate components. """ - ensemble, ensemble_std = self.sample_ensemble( + ensemble, per_member_std = self.sample_ensemble( init_states, forcing_features, target_states ) ensemble_mean = ensemble.mean(dim=1) - if ensemble_std is not None: - pred_std = ensemble_std.mean(dim=1) + if per_member_std is not None: + pred_std = per_member_std.mean(dim=1) else: pred_std = self.per_var_std diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 515b9232..854f8bff 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -107,7 +107,7 @@ def test_sample_ensemble_shapes_and_member_variability(): d_state = target_states.shape[-1] torch.manual_seed(42) - ensemble, ensemble_std = forecaster.sample_ensemble( + ensemble, per_member_std = forecaster.sample_ensemble( init_states, forcing_features, target_states, @@ -121,7 +121,7 @@ def test_sample_ensemble_shapes_and_member_variability(): num_grid_nodes, d_state, ) - assert ensemble_std is None + assert per_member_std is None # Members carry independent samples on the interior node assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0]) From d64878cfd3d9fbc5c37b3429d77f811724d2d771 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 10:53:28 +0530 Subject: [PATCH 6/9] Update neural_lam/models/forecasters/probabilistic.py Co-authored-by: Joel Oskarsson --- neural_lam/models/forecasters/probabilistic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index cfa6169b..560ae599 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -199,6 +199,7 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) + # After stacking shape of ensemble is (B, S, pred_steps, num_grid_nodes, num_state_vars) per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) From 6fd050a18adda3d9e2cfa7fc3276229a59d7d89b Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:10:06 +0530 Subject: [PATCH 7/9] Leave ProbabilisticARForecaster.compute_training_loss abstract 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. --- .../models/forecasters/probabilistic.py | 86 +++++-------------- tests/test_probabilistic_forecaster.py | 60 +++++++++++-- 2 files changed, 74 insertions(+), 72 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 560ae599..ae5a15cb 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -91,11 +91,17 @@ class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): state, so the inherited ``ARForecaster.forward`` unrolls one sampled 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 using the - scoring rule passed to ``compute_training_loss`` (from - ``neural_lam.metrics``); forecasters with model-specific objectives - (ensemble scoring rules, variational objectives) override - ``compute_training_loss``. + + ``compute_training_loss`` is intentionally left abstract here (it does + not fall back to ``ARForecaster``'s single-rollout objective, which + would silently train on one stochastic sample). There is no default + objective that fits every stochastic model: 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 the ensemble to a point estimate. Concrete subclasses must + define an objective appropriate to how they are meant to be trained + (e.g. an ensemble scoring rule such as CRPS, or a variational + objective). """ def __init__( @@ -199,12 +205,14 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) - # After stacking shape of ensemble is (B, S, pred_steps, num_grid_nodes, num_state_vars) + # After stacking, ensemble has shape + # (B, S, pred_steps, num_grid_nodes, num_state_vars) per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) return ensemble, per_member_std + @abstractmethod def compute_training_loss( self, init_states: torch.Tensor, @@ -213,65 +221,11 @@ def compute_training_loss( interior_mask_bool: torch.Tensor, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with ``self.loss``. - - Samples an ensemble of ``self.ensemble_size`` forecasts, averages - the members into an ensemble mean forecast, scores it against the - target states on interior nodes and averages over batch and time. - When members predict their own std, the std passed to ``self.loss`` - is the plain average of the per-member stds; this is a - simplification of the true mixture predictive variance, which - would also include the spread between the member means (see the - ``ProbabilisticForecaster`` class docstring). - - Parameters - ---------- - init_states : torch.Tensor - Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial - states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: - ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), - ``num_grid_nodes`` is the number of spatial nodes, and - ``num_state_vars`` is the state feature dimension. - forcing_features : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. - External forcings provided at each predicted step. Dims: ``B`` - is batch size, ``pred_steps`` is the rollout length, - ``num_grid_nodes`` is the number of spatial nodes, and - ``num_forcing_vars`` is the forcing feature dimension (already - concatenated past/current/future windows). - target_states : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True - 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. - interior_mask_bool : torch.Tensor - Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``self.loss`` so that only interior - nodes are scored. + Compute the training objective for one batch. - Returns - ------- - batch_loss : torch.Tensor - Scalar. The scoring rule applied to the ensemble mean, averaged - over batch and time. - loss_components : dict of {str: torch.Tensor} - Empty; this objective has no separate components. + Left abstract; see the class docstring for why there is no default + objective. Concrete subclasses typically call ``sample_ensemble`` + and score the resulting members with an objective appropriate to + the model (see ``Forecaster.compute_training_loss`` for the + signature and general contract). """ - ensemble, per_member_std = self.sample_ensemble( - init_states, forcing_features, target_states - ) - ensemble_mean = ensemble.mean(dim=1) - if per_member_std is not None: - pred_std = per_member_std.mean(dim=1) - else: - pred_std = self.per_var_std - - batch_loss = torch.mean( - self.loss( - ensemble_mean, - target_states, - pred_std, - mask=interior_mask_bool, - ) - ) - return batch_loss, {} diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 854f8bff..f7cfe7de 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -37,6 +37,43 @@ def forward(self, prev_state, prev_prev_state, forcing): return pred_state, None +class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster): + """ + Test-only concrete ``ProbabilisticARForecaster``. + + ``ProbabilisticARForecaster`` leaves ``compute_training_loss`` abstract + (no single default objective fits every stochastic model), so tests + that only need a working forecaster to instantiate use this example + ensemble-mean objective rather than the base class directly. + """ + + def compute_training_loss( + self, + init_states, + forcing_features, + target_states, + interior_mask_bool, + ): + ensemble, per_member_std = self.sample_ensemble( + init_states, forcing_features, target_states + ) + ensemble_mean = ensemble.mean(dim=1) + pred_std = ( + per_member_std.mean(dim=1) + if per_member_std is not None + else self.per_var_std + ) + batch_loss = torch.mean( + self.loss( + ensemble_mean, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} + + def _example_batch(datastore, B=2, pred_steps=3): """Create constant example input tensors matching the datastore dims.""" num_grid_nodes = datastore.num_grid_points @@ -90,7 +127,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): def test_sample_ensemble_shapes_and_member_variability(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2 ) @@ -138,7 +175,7 @@ def test_sample_ensemble_shapes_and_member_variability(): def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, loss="mse" ) @@ -171,7 +208,18 @@ def test_probabilistic_forecaster_rejects_empty_ensemble(): predictor = NoisyStepPredictor(datastore=datastore, output_std=False) with pytest.raises(ValueError, match="ensemble_size"): - ProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + ConcreteProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + + +def test_probabilistic_ar_forecaster_is_abstract(): + """ProbabilisticARForecaster leaves compute_training_loss abstract, so + it cannot be instantiated directly; only a subclass that defines an + objective can.""" + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(TypeError, match="abstract"): + ProbabilisticARForecaster(predictor, datastore, ensemble_size=2) def test_module_training_step_delegates_to_forecaster(): @@ -206,7 +254,7 @@ def test_module_training_step_delegates_to_forecaster(): torch.testing.assert_close(batch_loss, expected_loss) -class MemberCountRecordingForecaster(ProbabilisticARForecaster): +class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster): """ProbabilisticARForecaster recording the requested member count.""" def sample_ensemble(self, *args, **kwargs): @@ -261,7 +309,7 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, config=config ) @@ -282,7 +330,7 @@ def test_probabilistic_module_test_step_not_implemented(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, config=config ) model = ProbabilisticForecasterModule( From a1350ff21848b0cf6afbb35f94019d091c8a5899 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:18:19 +0530 Subject: [PATCH 8/9] Require an explicit member count instead of a stored ensemble_size 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. --- .../models/forecasters/probabilistic.py | 34 +++++++-------- neural_lam/models/probabilistic_module.py | 9 ++-- tests/test_probabilistic_forecaster.py | 42 +++++++++++-------- 3 files changed, 43 insertions(+), 42 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index ae5a15cb..1c5aeeb4 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -39,7 +39,7 @@ def sample_ensemble( init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - num_members: int | None = None, + num_members: int, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Sample an ensemble of forecasts. @@ -64,9 +64,8 @@ def sample_ensemble( state values used only to overwrite boundary nodes at each predicted step, identically in every member. Dims: same as one member. - num_members : int or None - Number of ensemble members ``S`` to sample. When ``None``, the - forecaster's configured ensemble size is used. + num_members : int + Number of ensemble members ``S`` to sample. Returns ------- @@ -108,7 +107,6 @@ def __init__( self, predictor: StepPredictor, datastore: BaseDatastore, - ensemble_size: int, config: NeuralLAMConfig | None = None, loss: str = "wmse", ) -> None: @@ -122,9 +120,6 @@ def __init__( fresh sample of the next state. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. - ensemble_size : int - Number of ensemble members to sample when no explicit member - count is given, in particular for the training objective. config : NeuralLAMConfig or None Configuration used to compute the constant per-variable std substituted for ``pred_std`` when ``predictor`` does not output @@ -135,18 +130,13 @@ def __init__( ``compute_training_loss`` and stored as ``self.loss``. """ super().__init__(predictor, datastore, config=config, loss=loss) - if ensemble_size < 1: - raise ValueError( - f"ensemble_size must be at least 1, got {ensemble_size}" - ) - self.ensemble_size = ensemble_size def sample_ensemble( self, init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - num_members: int | None = None, + num_members: int, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Sample an ensemble of forecasts. @@ -174,9 +164,8 @@ def sample_ensemble( Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state values used only to overwrite boundary nodes at each AR step, identically in every member. Dims: same as one member. - num_members : int or None - Number of ensemble members ``S`` to sample. When ``None``, - ``self.ensemble_size`` is used. + num_members : int + Number of ensemble members ``S`` to sample. Returns ------- @@ -190,9 +179,16 @@ def sample_ensemble( why the ensemble is then a mixture, not this averaged with the others), when the wrapped predictor outputs an std, otherwise ``None``. Dims: same as ``ensemble``. + + Raises + ------ + ValueError + If ``num_members`` is less than 1. """ - if num_members is None: - num_members = self.ensemble_size + if num_members < 1: + raise ValueError( + f"num_members must be at least 1, got {num_members}" + ) member_list = [] member_std_list = [] diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 37a756d7..6bda75cd 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -25,7 +25,7 @@ class ProbabilisticForecasterModule(ForecasterModule): # The wrapped forecaster must be able to sample ensemble forecasts forecaster: ProbabilisticForecaster - def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): + def __init__(self, *args, eval_ensemble_size: int, **kwargs): """ Initialize the module and store the evaluation ensemble size. @@ -35,15 +35,14 @@ def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): Positional arguments forwarded to ``ForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). - eval_ensemble_size : int or None - Number of ensemble members sampled during validation. ``None`` - uses the forecaster's configured ensemble size. + eval_ensemble_size : int + Number of ensemble members sampled during validation. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` (``lr``, ...). """ super().__init__(*args, **kwargs) - if eval_ensemble_size is not None and eval_ensemble_size < 1: + if eval_ensemble_size < 1: raise ValueError( "eval_ensemble_size must be at least 1, " f"got {eval_ensemble_size}" diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index f7cfe7de..28006671 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -45,8 +45,14 @@ class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster): (no single default objective fits every stochastic model), so tests that only need a working forecaster to instantiate use this example ensemble-mean objective rather than the base class directly. + ``sample_ensemble`` always requires an explicit member count, so this + class takes its own ``train_num_members`` for the training objective. """ + def __init__(self, *args, train_num_members: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.train_num_members = train_num_members + def compute_training_loss( self, init_states, @@ -55,7 +61,10 @@ def compute_training_loss( interior_mask_bool, ): ensemble, per_member_std = self.sample_ensemble( - init_states, forcing_features, target_states + init_states, + forcing_features, + target_states, + num_members=self.train_num_members, ) ensemble_mean = ensemble.mean(dim=1) pred_std = ( @@ -127,9 +136,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): def test_sample_ensemble_shapes_and_member_variability(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) # Override masks to test boundary masking behaviour forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) @@ -165,18 +172,12 @@ def test_sample_ensemble_shapes_and_member_variability(): # Boundary nodes are overwritten with the true state in every member assert torch.all(ensemble[:, :, :, 1:] == 5.0) - # Without an explicit member count the configured ensemble_size is used - default_ensemble, _ = forecaster.sample_ensemble( - init_states, forcing_features, target_states - ) - assert default_ensemble.shape[1] == forecaster.ensemble_size - def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, loss="mse" + predictor, datastore, loss="mse", train_num_members=2 ) init_states, forcing_features, target_states = _example_batch(datastore) @@ -203,12 +204,16 @@ def test_probabilistic_training_loss_gradient_flow(): assert predictor.noise_scale.grad != 0.0 -def test_probabilistic_forecaster_rejects_empty_ensemble(): +def test_sample_ensemble_rejects_empty_member_count(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) + init_states, forcing_features, target_states = _example_batch(datastore) - with pytest.raises(ValueError, match="ensemble_size"): - ConcreteProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + with pytest.raises(ValueError, match="num_members"): + forecaster.sample_ensemble( + init_states, forcing_features, target_states, num_members=0 + ) def test_probabilistic_ar_forecaster_is_abstract(): @@ -219,7 +224,7 @@ def test_probabilistic_ar_forecaster_is_abstract(): predictor = NoisyStepPredictor(datastore=datastore, output_std=False) with pytest.raises(TypeError, match="abstract"): - ProbabilisticARForecaster(predictor, datastore, ensemble_size=2) + ProbabilisticARForecaster(predictor, datastore) def test_module_training_step_delegates_to_forecaster(): @@ -272,7 +277,7 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): ) ) forecaster = MemberCountRecordingForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, @@ -310,7 +315,7 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): ) ) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) with pytest.raises(ValueError, match="eval_ensemble_size"): @@ -331,12 +336,13 @@ def test_probabilistic_module_test_step_not_implemented(): ) ) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, datastore=datastore, + eval_ensemble_size=2, ) init_states, forcing_features, target_states = _example_batch(datastore) From 1cbded635ac606a4a4cab8c7a23de739198e5840 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:26:10 +0530 Subject: [PATCH 9/9] Implement ProbabilisticForecasterModule.test_step 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. --- neural_lam/models/probabilistic_module.py | 108 ++++++++++++++++------ tests/test_probabilistic_forecaster.py | 27 ++++-- 2 files changed, 99 insertions(+), 36 deletions(-) diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 6bda75cd..453b00c6 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -1,5 +1,8 @@ """Lightning module evaluating probabilistic forecasters as ensembles.""" +# Standard library +import warnings + # Third-party import torch @@ -14,9 +17,9 @@ class ProbabilisticForecasterModule(ForecasterModule): Lightning module for forecasters that sample ensemble forecasts. Training is inherited unchanged from ``ForecasterModule``: the wrapped - forecaster assembles its own training loss. Validation is ensemble - based instead of deterministic: an ensemble is sampled from the - forecaster and scored through its ensemble mean (root-mean-squared + forecaster assembles its own training loss. Validation and testing are + ensemble based instead of deterministic: an ensemble is sampled from + the forecaster and scored through its ensemble mean (root-mean-squared error of the ensemble mean). The module only assumes that the forecaster can sample ensemble forecasts of the correct shape; it makes no assumption on how the members are produced. @@ -36,7 +39,8 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ``ForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). eval_ensemble_size : int - Number of ensemble members sampled during validation. + Number of ensemble members sampled during validation and + testing. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` (``lr``, ...). @@ -49,23 +53,30 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ) self.eval_ensemble_size = eval_ensemble_size self.val_metrics = {"ens_mse": []} + self.test_metrics = {"ens_mse": []} - def validation_step(self, batch, batch_idx): + def _ensemble_step(self, batch, phase: str): """ - Perform a single ensemble validation step. + Sample an ensemble and score its mean against the target states. - Samples an ensemble from the forecaster and scores its ensemble - mean against the target states on interior nodes. Logs the - root-mean-squared error of the ensemble mean per configured rollout - step and averaged over the rollout, and collects per-variable - ensemble-mean MSE for epoch-end aggregation. + Shared by ``validation_step`` and ``test_step``: samples + ``self.eval_ensemble_size`` members, scores the ensemble mean with + plain (unweighted) MSE on interior nodes, logs the root-mean-squared + error per configured rollout step and averaged over the rollout + under the given phase's prefix. Parameters ---------- batch : tuple The batch of data. - batch_idx : int - The index of the batch. + phase : str + Logging phase, either ``"val"`` or ``"test"``. + + Returns + ------- + torch.Tensor + Per-variable ensemble-mean MSE, shape + ``(B, pred_steps, num_state_vars)``, for epoch-end aggregation. """ init_states, target_states, forcing_features, _ = batch ensemble, _ = self.forecaster.sample_ensemble( @@ -91,34 +102,55 @@ def validation_step(self, batch, batch_idx): ) time_step_rmse = torch.sqrt(time_step_mse) mean_rmse = torch.mean(time_step_rmse) - self._warn_skipped_val_steps(len(time_step_rmse), "val") + self._warn_skipped_val_steps(len(time_step_rmse), phase) - val_log_dict = { - f"val_loss_unroll{step}": time_step_rmse[step - 1] + log_dict = { + f"{phase}_loss_unroll{step}": time_step_rmse[step - 1] for step in self.hparams.val_steps_to_log if step <= len(time_step_rmse) } - val_log_dict["val_mean_loss"] = mean_rmse + log_dict[f"{phase}_mean_loss"] = mean_rmse self.log_dict( - val_log_dict, + log_dict, on_step=False, on_epoch=True, sync_dist=True, batch_size=batch[0].shape[0], ) - entry_mses = metrics.mse( + return metrics.mse( ensemble_mean, target_states, std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) + + def validation_step(self, batch, batch_idx): + """ + Perform a single ensemble validation step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "val") self.val_metrics["ens_mse"].append(entry_mses) def test_step(self, batch, batch_idx): """ - Not supported: ensemble test evaluation is not implemented. + Perform a single ensemble test step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. Parameters ---------- @@ -126,14 +158,32 @@ def test_step(self, batch, batch_idx): The batch of data. batch_idx : int The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "test") + self.test_metrics["ens_mse"].append(entry_mses) - Raises - ------ - NotImplementedError - Always; only training and ensemble validation are implemented - for probabilistic forecasters. + def on_test_epoch_end(self): """ - raise NotImplementedError( - "Ensemble test evaluation is not implemented for " - "probabilistic forecasters." - ) + Perform actions at the end of the test epoch. + + Aggregates ensemble test metrics. Overrides + ``ForecasterModule.on_test_epoch_end``, which also handles spatial + loss maps and example plots that ``test_step`` here does not + populate. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + if self.trainer.is_global_zero and self.hparams.metrics_watch: + unmatched = set(self.hparams.metrics_watch) - self.matched_metrics + if unmatched: + warnings.warn( + "The following metrics in --metrics_watch " + "were not found during test phase: " + f"{sorted(unmatched)}. Ensure the metric prefix " + "matches the evaluation mode (expected 'test_')." + ) + + self.matched_metrics = set() + + for metric_list in self.test_metrics.values(): + metric_list.clear() diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 28006671..aecdb09f 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -327,27 +327,40 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): ) -def test_probabilistic_module_test_step_not_implemented(): +def test_probabilistic_module_test_step_scores_ensemble_mean(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ConcreteProbabilisticARForecaster( + forecaster = MemberCountRecordingForecaster( predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - eval_ensemble_size=2, + eval_ensemble_size=3, ) - init_states, forcing_features, target_states = _example_batch(datastore) - batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) batch = (init_states, target_states, forcing_features, batch_times) - with pytest.raises(NotImplementedError): - model.test_step(batch, 0) + torch.manual_seed(42) + model.test_step(batch, 0) + + # Test samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.test_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses))