From 1713911ea38c1cc121deed3948fa16798b1c8915 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sun, 5 Jul 2026 09:31:42 +0530 Subject: [PATCH 01/16] 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 02/16] 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 03/16] 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 04/16] 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 05/16] 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 06/16] 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 07/16] 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 08/16] 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 09/16] 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)) From 703e53de6ad6155d05aab77df3215dc952137bb0 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 12:27:26 +0530 Subject: [PATCH 10/16] Address PR review: separate ensemble RMSE from validation loss naming Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the training loss, per review feedback. --- neural_lam/models/probabilistic_module.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 453b00c6..4a2843bd 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -63,7 +63,10 @@ def _ensemble_step(self, batch, phase: str): ``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. + under the given phase's prefix. This RMSE is a diagnostic metric, + not the training loss: it always scores the ensemble mean with + plain MSE, regardless of what objective ``compute_training_loss`` + actually trains on, which is not recomputed here. Parameters ---------- @@ -105,11 +108,11 @@ def _ensemble_step(self, batch, phase: str): self._warn_skipped_val_steps(len(time_step_rmse), phase) log_dict = { - f"{phase}_loss_unroll{step}": time_step_rmse[step - 1] + f"{phase}_ens_rmse_unroll{step}": time_step_rmse[step - 1] for step in self.hparams.val_steps_to_log if step <= len(time_step_rmse) } - log_dict[f"{phase}_mean_loss"] = mean_rmse + log_dict[f"{phase}_mean_ens_rmse"] = mean_rmse self.log_dict( log_dict, on_step=False, From 45ffdeb707874b3d3f52819ddd6d56b1e98b8ba6 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 12:41:17 +0530 Subject: [PATCH 11/16] Address PR review: drop redundant inline comments in probabilistic tests Remove explanatory comments around the per_var_std overrides; the assignments are clear on their own. --- tests/test_probabilistic_forecaster.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index aecdb09f..f5b87d3d 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -107,8 +107,6 @@ def test_ar_forecaster_training_loss_matches_direct_score(): 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 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( @@ -183,8 +181,6 @@ def test_probabilistic_training_loss_gradient_flow(): 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) From d020691100e5ace37172f958081a8c66ca191db2 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 13:08:57 +0530 Subject: [PATCH 12/16] Address PR review: split ForecasterModule into an abstract base + concrete deterministic/probabilistic modules Introduce BaseForecasterModule (abstract) under models/forecasters/ holding shared plumbing (training_step, common_step, batch standardization, checkpoint compatibility, plotting/aggregation helpers), with validation_step, test_step and on_test_epoch_end left abstract since they differ meaningfully between evaluation modes. Rename ForecasterModule to DeterministicForecasterModule and move it, alongside ProbabilisticForecasterModule, into forecasters/ as siblings implementing the shared contract, rather than one subclassing the other. --- neural_lam/models/__init__.py | 5 +- .../models/forecasters/autoregressive.py | 2 +- neural_lam/models/forecasters/base.py | 6 +- .../{module.py => forecasters/base_module.py} | 249 +++------------- .../forecasters/deterministic_module.py | 272 ++++++++++++++++++ .../{ => forecasters}/probabilistic_module.py | 38 +-- neural_lam/train_model.py | 12 +- neural_lam/weather_dataset.py | 2 +- tests/test_checkpoint.py | 8 +- tests/test_datasets.py | 4 +- tests/test_gpu_normalization.py | 15 +- tests/test_plotting.py | 12 +- tests/test_prediction_model_classes.py | 21 +- tests/test_probabilistic_forecaster.py | 4 +- tests/test_train_model_warnings.py | 8 +- tests/test_training.py | 18 +- 16 files changed, 398 insertions(+), 278 deletions(-) rename neural_lam/models/{module.py => forecasters/base_module.py} (79%) create mode 100644 neural_lam/models/forecasters/deterministic_module.py rename neural_lam/models/{ => forecasters}/probabilistic_module.py (82%) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 1bfe9eb5..cbeb1b01 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,12 +3,13 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.base_module import BaseForecasterModule +from .forecasters.deterministic_module import DeterministicForecasterModule from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, ) -from .module import ForecasterModule -from .probabilistic_module import ProbabilisticForecasterModule +from .forecasters.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/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 74f2dc77..a121bfef 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -141,7 +141,7 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the wrapped predictor outputs an std, otherwise ``None`` (in which - case ``ForecasterModule`` substitutes the constant + case ``DeterministicForecasterModule`` substitutes the constant per-variable std). Dims: same as ``prediction``. """ diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 8dfb0002..1e5b7db1 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -77,7 +77,7 @@ def forward( ``predicts_std`` is True, otherwise ``None``. Per-feature predicted standard deviation; when ``None``, the constant per-variable std is substituted upstream by - ``ForecasterModule``. Dims: same as ``prediction``. + ``DeterministicForecasterModule``. Dims: same as ``prediction``. """ @abstractmethod @@ -95,8 +95,8 @@ def compute_training_loss( to produce from the batch, which loss terms to compute from them and 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. + wrapping ``BaseForecasterModule`` only injects the interior mask, + logs the returned components and optimizes the returned loss. Parameters ---------- diff --git a/neural_lam/models/module.py b/neural_lam/models/forecasters/base_module.py similarity index 79% rename from neural_lam/models/module.py rename to neural_lam/models/forecasters/base_module.py index e7420b52..e32a99b4 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -1,9 +1,10 @@ -"""Lightning module handling training, validation and testing loops.""" +"""Abstract Lightning module shared by deterministic and probabilistic +forecaster modules.""" # Standard library import os import warnings -from typing import Any +from abc import ABC, abstractmethod # Third-party import matplotlib.pyplot as plt @@ -17,17 +18,26 @@ from neural_lam.utils import get_integer_time # Local -from .. import metrics, vis -from ..config import NeuralLAMConfig -from ..datastore import BaseDatastore -from ..weather_dataset import WeatherDataset -from .forecasters.base import Forecaster +from ... import vis +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ...weather_dataset import WeatherDataset +from .base import Forecaster -class ForecasterModule(pl.LightningModule): +class BaseForecasterModule(pl.LightningModule, ABC): """ - Lightning module handling training, validation and testing loops. - Wraps a Forecaster instance which performs the actual prediction. + Abstract Lightning module wrapping a ``Forecaster``. + + Owns everything that does not depend on whether the wrapped forecaster + produces a single deterministic forecast or samples an ensemble: + batch standardization, the training loop, optimizer configuration, + checkpoint compatibility, and the plotting/aggregation helpers used by + validation and testing. ``validation_step``, ``test_step`` and + ``on_test_epoch_end`` differ enough between the two evaluation modes + that they are left abstract; concrete subclasses implement them + independently (see ``DeterministicForecasterModule`` and + ``ProbabilisticForecasterModule``) rather than overriding one another. """ # pylint: disable=arguments-differ @@ -47,7 +57,7 @@ def __init__( args=None, ): """ - Initialize the ForecasterModule. + Initialize the BaseForecasterModule. Parameters ---------- @@ -190,16 +200,6 @@ def __init__( self.forcing_mean = None self.forcing_std = None - self.val_metrics: dict[str, list] = { - "mse": [], - } - self.test_metrics: dict[str, list] = { - "mse": [], - "mae": [], - } - if self.forecaster.predicts_std: - self.test_metrics["output_std"] = [] # Treat as metric - # For making restoring of optimizer state optional self.restore_opt = restore_opt @@ -208,9 +208,6 @@ def __init__( self.create_gif = create_gif self.plotted_examples = 0 - # For storing spatial loss maps during evaluation - self.spatial_loss_maps: list[Any] = [] - # Warn once per phase if val_steps_to_log exceeds the actual rollout self._val_steps_warn_issued = False self._test_steps_warn_issued = False @@ -418,10 +415,15 @@ def _warn_skipped_val_steps(self, pred_steps: int, phase: str) -> None: ) setattr(self, flag, True) + @abstractmethod def validation_step(self, batch, batch_idx): """ Perform a single validation step. + Concrete subclasses must both score the batch and populate + ``self.val_metrics`` for epoch-end aggregation by + ``on_validation_epoch_end``. + Parameters ---------- batch : tuple @@ -429,44 +431,6 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.forecaster.per_var_std - - time_step_loss = torch.mean( - self.forecaster.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - mean_loss = torch.mean(time_step_loss) - self._warn_skipped_val_steps(len(time_step_loss), "val") - - val_log_dict = { - f"val_loss_unroll{step}": time_step_loss[step - 1] - for step in self.hparams.val_steps_to_log - if step <= len(time_step_loss) - } - val_log_dict["val_mean_loss"] = mean_loss - 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( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.val_metrics["mse"].append(entry_mses) def on_validation_epoch_end(self): """ @@ -490,11 +454,15 @@ def on_validation_epoch_end(self): for metric_list in self.val_metrics.values(): metric_list.clear() - # pylint: disable-next=unused-argument + @abstractmethod def test_step(self, batch, batch_idx): """ Perform a single test step. + Concrete subclasses must both score the batch and populate + ``self.test_metrics`` for epoch-end aggregation by + ``on_test_epoch_end``. + Parameters ---------- batch : tuple @@ -502,83 +470,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - - if pred_std is not None: - mean_pred_std = torch.mean( - pred_std[..., self.interior_mask_bool, :], dim=-2 - ) - self.test_metrics["output_std"].append(mean_pred_std) - - if pred_std is None: - pred_std = self.forecaster.per_var_std - - time_step_loss = torch.mean( - self.forecaster.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - mean_loss = torch.mean(time_step_loss) - self._warn_skipped_val_steps(len(time_step_loss), "test") - - test_log_dict = { - f"test_loss_unroll{step}": time_step_loss[step - 1] - for step in self.hparams.val_steps_to_log - if step <= len(time_step_loss) - } - test_log_dict["test_mean_loss"] = mean_loss - - self.log_dict( - test_log_dict, - on_step=False, - on_epoch=True, - sync_dist=True, - batch_size=batch[0].shape[0], - ) - - for metric_name in ("mse", "mae"): - metric_func = metrics.get_metric(metric_name) - batch_metric_vals = metric_func( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.test_metrics[metric_name].append(batch_metric_vals) - - spatial_loss = self.forecaster.loss( - prediction, target_states, pred_std, average_grid=False - ) - log_spatial_losses = spatial_loss[ - :, - [ - step - 1 - for step in self.hparams.val_steps_to_log - if step <= spatial_loss.shape[1] - ], - ] - self.spatial_loss_maps.append(log_spatial_losses) - - if ( - self.trainer.is_global_zero - and self.plotted_examples < self.n_example_pred - ): - n_additional_examples = min( - prediction.shape[0], - self.n_example_pred - self.plotted_examples, - ) - - self.plot_examples( - batch, - n_additional_examples, - prediction=prediction, - split="test", - ) def plot_examples(self, batch, n_examples, split, prediction): """ @@ -867,82 +758,16 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix): plt.close("all") + @abstractmethod def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - Aggregates and plots test metrics and spatial loss maps. - """ - self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") - - spatial_loss_tensor = self.all_gather_cat( - torch.cat(self.spatial_loss_maps, dim=0) - ) - if self.trainer.is_global_zero: - mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) - - loss_map_figs = [ - vis.plot_spatial_error( - error=loss_map, - datastore=self.datastore, - title=f"Test loss, t={t_i} " - f"({(self.time_step_int * t_i)} {self.time_step_unit})", - ) - for t_i, loss_map in zip( - self.hparams.val_steps_to_log, mean_spatial_loss - ) - ] - - for i, fig in enumerate(loss_map_figs): - key = "test_loss" - if not isinstance(self.logger, pl.loggers.WandbLogger): - key = f"{key}_{i}" - if hasattr(self.logger, "log_image"): - self.logger.log_image(key=key, images=[fig]) - pdf_loss_map_figs = [ - vis.plot_spatial_error(error=loss_map, datastore=self.datastore) - for loss_map in mean_spatial_loss - ] - pdf_loss_maps_dir = os.path.join( - self.logger.save_dir, "spatial_loss_maps" - ) - os.makedirs(pdf_loss_maps_dir, exist_ok=True) - for t_i, fig in zip( - self.hparams.val_steps_to_log, pdf_loss_map_figs - ): - fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) - - torch.save( - mean_spatial_loss.cpu(), - os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), - ) - - if 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() - self.spatial_loss_maps.clear() - - # Clear stored test metrics so repeated `trainer.test()` calls on - # the same model instance start from a clean slate (otherwise the - # tensors accumulate and skew the aggregated metrics). - for metric_list in self.test_metrics.values(): - metric_list.clear() - - # Reset the example-plot counter so example prediction plots are - # generated again on every `trainer.test()` call, not just the - # first one (the guard `plotted_examples < n_example_pred` would - # otherwise stay permanently False). - self.plotted_examples = 0 + Concrete subclasses must at least aggregate and plot + ``self.test_metrics`` (typically via ``aggregate_and_plot_metrics``) + and reset any epoch-scoped state they accumulate during + ``test_step``. + """ def on_load_checkpoint(self, checkpoint): """ diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/forecasters/deterministic_module.py new file mode 100644 index 00000000..0cd1cbe8 --- /dev/null +++ b/neural_lam/models/forecasters/deterministic_module.py @@ -0,0 +1,272 @@ +"""Lightning module evaluating forecasters through a single deterministic +rollout per batch.""" + +# Standard library +import os +import warnings +from typing import Any + +# Third-party +import pytorch_lightning as pl +import torch + +# Local +from ... import metrics, vis +from .base_module import BaseForecasterModule + + +class DeterministicForecasterModule(BaseForecasterModule): + """ + Lightning module for a single deterministic forecast per batch. + + Validation and testing score the forecaster's own single-rollout + prediction directly with ``forecaster.loss``, as opposed to + ``ProbabilisticForecasterModule``, which samples and scores an + ensemble. Training is shared with that module unchanged (see + ``BaseForecasterModule.training_step``). + """ + + def __init__(self, *args, **kwargs): + """ + Initialize the module and its deterministic evaluation metrics. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + **kwargs + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` + (``lr``, ...). + """ + super().__init__(*args, **kwargs) + self.val_metrics: dict[str, list] = { + "mse": [], + } + self.test_metrics: dict[str, list] = { + "mse": [], + "mae": [], + } + if self.forecaster.predicts_std: + self.test_metrics["output_std"] = [] # Treat as metric + + # For storing spatial loss maps during evaluation + self.spatial_loss_maps: list[Any] = [] + + def validation_step(self, batch, batch_idx): + """ + Perform a single validation step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + if pred_std is None: + pred_std = self.forecaster.per_var_std + + time_step_loss = torch.mean( + self.forecaster.loss( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + self._warn_skipped_val_steps(len(time_step_loss), "val") + + val_log_dict = { + f"val_loss_unroll{step}": time_step_loss[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_loss) + } + val_log_dict["val_mean_loss"] = mean_loss + 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( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.val_metrics["mse"].append(entry_mses) + + # pylint: disable-next=unused-argument + def test_step(self, batch, batch_idx): + """ + Perform a single test step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + + if pred_std is not None: + mean_pred_std = torch.mean( + pred_std[..., self.interior_mask_bool, :], dim=-2 + ) + self.test_metrics["output_std"].append(mean_pred_std) + + if pred_std is None: + pred_std = self.forecaster.per_var_std + + time_step_loss = torch.mean( + self.forecaster.loss( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + self._warn_skipped_val_steps(len(time_step_loss), "test") + + test_log_dict = { + f"test_loss_unroll{step}": time_step_loss[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_loss) + } + test_log_dict["test_mean_loss"] = mean_loss + + self.log_dict( + test_log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + for metric_name in ("mse", "mae"): + metric_func = metrics.get_metric(metric_name) + batch_metric_vals = metric_func( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.test_metrics[metric_name].append(batch_metric_vals) + + spatial_loss = self.forecaster.loss( + prediction, target_states, pred_std, average_grid=False + ) + log_spatial_losses = spatial_loss[ + :, + [ + step - 1 + for step in self.hparams.val_steps_to_log + if step <= spatial_loss.shape[1] + ], + ] + self.spatial_loss_maps.append(log_spatial_losses) + + if ( + self.trainer.is_global_zero + and self.plotted_examples < self.n_example_pred + ): + n_additional_examples = min( + prediction.shape[0], + self.n_example_pred - self.plotted_examples, + ) + + self.plot_examples( + batch, + n_additional_examples, + prediction=prediction, + split="test", + ) + + def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + Aggregates and plots test metrics and spatial loss maps. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + spatial_loss_tensor = self.all_gather_cat( + torch.cat(self.spatial_loss_maps, dim=0) + ) + if self.trainer.is_global_zero: + mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) + + loss_map_figs = [ + vis.plot_spatial_error( + error=loss_map, + datastore=self.datastore, + title=f"Test loss, t={t_i} " + f"({(self.time_step_int * t_i)} {self.time_step_unit})", + ) + for t_i, loss_map in zip( + self.hparams.val_steps_to_log, mean_spatial_loss + ) + ] + + for i, fig in enumerate(loss_map_figs): + key = "test_loss" + if not isinstance(self.logger, pl.loggers.WandbLogger): + key = f"{key}_{i}" + if hasattr(self.logger, "log_image"): + self.logger.log_image(key=key, images=[fig]) + + pdf_loss_map_figs = [ + vis.plot_spatial_error(error=loss_map, datastore=self.datastore) + for loss_map in mean_spatial_loss + ] + pdf_loss_maps_dir = os.path.join( + self.logger.save_dir, "spatial_loss_maps" + ) + os.makedirs(pdf_loss_maps_dir, exist_ok=True) + for t_i, fig in zip( + self.hparams.val_steps_to_log, pdf_loss_map_figs + ): + fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) + + torch.save( + mean_spatial_loss.cpu(), + os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), + ) + + if 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() + self.spatial_loss_maps.clear() + + # Clear stored test metrics so repeated `trainer.test()` calls on + # the same model instance start from a clean slate (otherwise the + # tensors accumulate and skew the aggregated metrics). + for metric_list in self.test_metrics.values(): + metric_list.clear() + + # Reset the example-plot counter so example prediction plots are + # generated again on every `trainer.test()` call, not just the + # first one (the guard `plotted_examples < n_example_pred` would + # otherwise stay permanently False). + self.plotted_examples = 0 diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/forecasters/probabilistic_module.py similarity index 82% rename from neural_lam/models/probabilistic_module.py rename to neural_lam/models/forecasters/probabilistic_module.py index 4a2843bd..cf8d38b0 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/forecasters/probabilistic_module.py @@ -7,22 +7,22 @@ import torch # Local -from .. import metrics -from .forecasters.probabilistic import ProbabilisticForecaster -from .module import ForecasterModule +from ... import metrics +from .base_module import BaseForecasterModule +from .probabilistic import ProbabilisticForecaster -class ProbabilisticForecasterModule(ForecasterModule): +class ProbabilisticForecasterModule(BaseForecasterModule): """ Lightning module for forecasters that sample ensemble forecasts. - Training is inherited unchanged from ``ForecasterModule``: the wrapped - 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. + Training is inherited unchanged from ``BaseForecasterModule``: the + wrapped 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. """ # The wrapped forecaster must be able to sample ensemble forecasts @@ -36,13 +36,13 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ---------- *args Positional arguments forwarded to - ``ForecasterModule.__init__`` (``forecaster``, ``config``, + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). eval_ensemble_size : int Number of ensemble members sampled during validation and testing. **kwargs - Keyword arguments forwarded to ``ForecasterModule.__init__`` + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` (``lr``, ...). """ super().__init__(*args, **kwargs) @@ -52,8 +52,8 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): f"got {eval_ensemble_size}" ) self.eval_ensemble_size = eval_ensemble_size - self.val_metrics = {"ens_mse": []} - self.test_metrics = {"ens_mse": []} + self.val_metrics: dict[str, list] = {"ens_mse": []} + self.test_metrics: dict[str, list] = {"ens_mse": []} def _ensemble_step(self, batch, phase: str): """ @@ -169,10 +169,10 @@ def on_test_epoch_end(self): """ 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. + Aggregates ensemble test metrics. Implements + ``BaseForecasterModule.on_test_epoch_end`` without the spatial loss + maps and example plots that ``DeterministicForecasterModule`` adds, + since ``test_step`` here does not populate them. """ self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index d5e86536..b0709a89 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -19,7 +19,7 @@ from . import utils from .config import load_config_and_datastore from .gnn_layers import GNN_TYPES -from .models import MODELS, ARForecaster, ForecasterModule +from .models import MODELS, ARForecaster, DeterministicForecasterModule from .weather_dataset import WeatherDataModule @@ -40,8 +40,8 @@ def __init__(self, prog): def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): """ - Reconstruct a ForecasterModule from a checkpoint without requiring the - caller to know the original architecture kwargs. + Reconstruct a DeterministicForecasterModule from a checkpoint without + requiring the caller to know the original architecture kwargs. The checkpoint must have been saved with args in hyper_parameters (i.e. created via train_model.main), so that model class and architecture kwargs @@ -66,7 +66,7 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): forecaster = ARForecaster( predictor, datastore, config=config, loss=args.loss ) - return ForecasterModule.load_from_checkpoint( + return DeterministicForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, @@ -440,7 +440,7 @@ def main(input_args=None): raise ValueError("devices should be 'auto' or a list of integers") # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule predictor_class = MODELS[args.model] predictor = predictor_class( datastore=datastore, @@ -463,7 +463,7 @@ def main(input_args=None): predictor, datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 4168396a..c62e1e5a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -471,7 +471,7 @@ def __getitem__( target states, forcing and batch times. The returned data is unstandardized; normalization is applied on-device - in `ForecasterModule.on_after_batch_transfer`. + in `BaseForecasterModule.on_after_batch_transfer`. Parameters ---------- diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 6f114043..6644aae6 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -8,7 +8,11 @@ # First-party from neural_lam import config as nlconfig from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from tests.dummy_datastore import DummyDatastore @@ -51,7 +55,7 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_upper=config.training.output_clamping.upper, ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 1941206b..319e3372 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -12,7 +12,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore @@ -239,7 +239,7 @@ def _create_graph(): predictor, datastore=datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index b063d626..dcc8ba63 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -4,7 +4,11 @@ # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -13,7 +17,8 @@ class _MockStepPredictor(StepPredictor): - """Minimal predictor so a ForecasterModule can be built without a graph.""" + """Minimal predictor so a DeterministicForecasterModule can be built + without a graph.""" def forward(self, prev_state, prev_prev_state, forcing): return torch.zeros_like(prev_state), None @@ -27,7 +32,7 @@ def _build_module(datastore): ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) forecaster = ARForecaster(predictor, datastore, config=config) - return ForecasterModule( + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) @@ -111,7 +116,9 @@ def test_safe_std_clamps_near_zero(): eps = torch.finfo(torch.float32).eps with pytest.warns(UserWarning, match="near-zero std"): - std = ForecasterModule._safe_std([0.0, 1.0, 2.0], eps, "state") + std = DeterministicForecasterModule._safe_std( + [0.0, 1.0, 2.0], eps, "state" + ) assert std[0] == eps assert std[1] == 1.0 diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 970590be..f55755e5 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -18,7 +18,11 @@ from neural_lam import config as nlconfig from neural_lam import vis from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -471,7 +475,7 @@ class ModelArgs: predictor, datastore=datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -666,7 +670,7 @@ class _SimpleLogger: # Shared ModelArgs for metrics_watch regression tests (issue #302). # Kept at module level to avoid copy-paste duplication across tests. def _build_metrics_watch_module(datastore, config): - """Build a ForecasterModule wired for metrics_watch tests.""" + """Build a DeterministicForecasterModule wired for metrics_watch tests.""" predictor = GraphLAM( datastore=datastore, graph_name="1level", @@ -681,7 +685,7 @@ def _build_metrics_watch_module(datastore, config): output_clamping_upper=config.training.output_clamping.upper, ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - return ForecasterModule( + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 9bc9d9c0..b2d901e7 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -7,7 +7,11 @@ # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -81,7 +85,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS @@ -99,7 +103,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -137,7 +141,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) # Load from checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -195,7 +199,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): output_std=False, ) # Use distinctive non-default values so we can detect silent fallback - # to ForecasterModule's defaults during load. + # to DeterministicForecasterModule's defaults during load. saved_loss = "mse" saved_lr = 0.123 saved_create_gif = True @@ -206,7 +210,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): predictor, datastore, config=config, loss=saved_loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -276,7 +280,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): ) # Load from hacked old checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -287,7 +291,8 @@ def test_forecaster_module_old_checkpoint(tmp_path): assert loaded_model.forecaster.predictor.__class__.__name__ == "GraphLAM" # Hyperparameters nested in the legacy 'args' namespace must round-trip - # rather than silently falling back to ForecasterModule defaults. + # rather than silently falling back to DeterministicForecasterModule + # defaults. 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 f5b87d3d..d98cc345 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -8,7 +8,7 @@ from neural_lam import metrics from neural_lam.models import ( ARForecaster, - ForecasterModule, + DeterministicForecasterModule, ProbabilisticARForecaster, ProbabilisticForecasterModule, StepPredictor, @@ -233,7 +233,7 @@ def test_module_training_step_delegates_to_forecaster(): ) ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a..6bb0654f 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -44,7 +44,8 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): def test_create_gif_forwarded_to_forecaster_module(): - """--create_gif must be forwarded to ForecasterModule.__init__.""" + """--create_gif must be forwarded to + DeterministicForecasterModule.__init__.""" mock_args = MagicMock() mock_args.eval = None mock_args.load = None @@ -76,7 +77,8 @@ def capture_init(_self, **kwargs): patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), patch("neural_lam.train_model.ARForecaster"), patch( - "neural_lam.models.module.ForecasterModule.__init__", + "neural_lam.models.forecasters.deterministic_module." + "DeterministicForecasterModule.__init__", capture_init, ), pytest.raises(SystemExit), @@ -85,5 +87,5 @@ def capture_init(_self, **kwargs): assert ( "create_gif" in captured_kwargs - ), "create_gif was not forwarded to ForecasterModule" + ), "create_gif was not forwarded to DeterministicForecasterModule" assert captured_kwargs["create_gif"] is True diff --git a/tests/test_training.py b/tests/test_training.py index 589e9d89..aca432c5 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -13,7 +13,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -105,7 +105,7 @@ def run_simple_training( ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS, ARForecaster @@ -125,7 +125,7 @@ def run_simple_training( ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -175,9 +175,9 @@ def all_gather(self, tensor_to_gather, sync_grads=False): return tensor_to_gather module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) # Simulate a 3D metric tensor: (N_eval, pred_steps, d_f) @@ -206,9 +206,9 @@ def all_gather(self, tensor, sync_grads=False): return torch.stack([tensor, tensor], dim=0) module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) tensor = torch.randn(4, 3, 5) # (N_eval, pred_steps, d_f) From 98ab692ad6da66fd08c2df76467c2314a4357373 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 21:04:59 +0530 Subject: [PATCH 13/16] Ignore .idea directory in .gitignore Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries. --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index a2c5bb26..37242198 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,7 @@ tags # Coc configuration directory .vim .vscode +.idea # macos .DS_Store From f0e01b1efccea28e1d257911afba2429bfd5178d Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:15:14 +0530 Subject: [PATCH 14/16] Address PR review: delegate validation/test loss computation to Forecaster DeterministicForecasterModule.validation_step/test_step read self.forecaster.loss and self.forecaster.per_var_std directly, so the Module still knew how to compute a loss from a prediction. Add an abstract Forecaster.score (implemented on ARForecaster) that resolves the pred_std fallback and applies a scoring rule internally; the Module now only calls forecaster.score(...) and never touches loss/per_var_std itself. --- .../models/forecasters/autoregressive.py | 63 +++++++++++++++++++ neural_lam/models/forecasters/base.py | 61 +++++++++++++++++- neural_lam/models/forecasters/base_module.py | 9 +-- .../forecasters/deterministic_module.py | 20 +++--- tests/test_prediction_model_classes.py | 46 ++++++++++++++ 5 files changed, 180 insertions(+), 19 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index a121bfef..ca8166cd 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,5 +1,8 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" +# Standard library +from typing import Callable, Optional + # Third-party import torch @@ -245,3 +248,63 @@ def compute_training_loss( ) ) return batch_loss, {} + + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Substitutes ``self.per_var_std`` for ``pred_std`` when the latter is + ``None`` (predictor does not output its own std), then applies + ``metric`` (defaulting to ``self.loss``, the configured scoring + rule). + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``; ``None`` when + the wrapped predictor does not output one, in which case + ``self.per_var_std`` is substituted. + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to ``self.loss``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + """ + if pred_std is None: + pred_std = self.per_var_std + metric_fn = self.loss if metric is None else metric + return metric_fn( + prediction, + target_states, + pred_std, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 1e5b7db1..da63869b 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, Optional # Third-party import torch @@ -75,9 +76,10 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when ``predicts_std`` is True, otherwise ``None``. Per-feature - predicted standard deviation; when ``None``, the constant - per-variable std is substituted upstream by - ``DeterministicForecasterModule``. Dims: same as ``prediction``. + predicted standard deviation; when ``None``, the forecaster's + own constant per-variable std fallback is substituted by + ``compute_training_loss``/``score``, not by the caller. Dims: + same as ``prediction``. """ @abstractmethod @@ -134,3 +136,56 @@ def compute_training_loss( the training phase. Empty when the objective has no separate components worth logging. """ + + @abstractmethod + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Wrapping ``BaseForecasterModule`` subclasses use this for + validation/test logging and diagnostics instead of computing a loss + themselves: the forecaster owns both its scoring rule and its + ``pred_std`` fallback, so it is the only place that knows how to + turn a raw ``pred_std`` (possibly ``None``) into a valid one and + apply a metric to it. + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``, as returned + alongside it by ``forward``. When ``None``, implementations + substitute their own constant per-variable std fallback. + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to the forecaster's + own configured scoring rule when ``None``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + """ diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/forecasters/base_module.py index e32a99b4..85faa451 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -62,10 +62,11 @@ def __init__( Parameters ---------- forecaster : Forecaster - 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. + The forecaster model to use for predictions. Owns the training + objective (``compute_training_loss``) and validation/test + scoring (``score``); this module and its subclasses never + compute a loss themselves, only inject shared inputs (e.g. the + interior mask) and log what the forecaster returns. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/forecasters/deterministic_module.py index 0cd1cbe8..217d4937 100644 --- a/neural_lam/models/forecasters/deterministic_module.py +++ b/neural_lam/models/forecasters/deterministic_module.py @@ -20,7 +20,7 @@ class DeterministicForecasterModule(BaseForecasterModule): Lightning module for a single deterministic forecast per batch. Validation and testing score the forecaster's own single-rollout - prediction directly with ``forecaster.loss``, as opposed to + prediction via ``forecaster.score``, as opposed to ``ProbabilisticForecasterModule``, which samples and scores an ensemble. Training is shared with that module unchanged (see ``BaseForecasterModule.training_step``). @@ -66,11 +66,9 @@ def validation_step(self, batch, batch_idx): The index of the batch. """ prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.forecaster.loss( + self.forecaster.score( prediction, target_states, pred_std, @@ -95,10 +93,11 @@ def validation_step(self, batch, batch_idx): batch_size=batch[0].shape[0], ) - entry_mses = metrics.mse( + entry_mses = self.forecaster.score( prediction, target_states, pred_std, + metric=metrics.mse, mask=self.interior_mask_bool, sum_vars=False, ) @@ -124,11 +123,8 @@ def test_step(self, batch, batch_idx): ) self.test_metrics["output_std"].append(mean_pred_std) - if pred_std is None: - pred_std = self.forecaster.per_var_std - time_step_loss = torch.mean( - self.forecaster.loss( + self.forecaster.score( prediction, target_states, pred_std, @@ -155,17 +151,17 @@ def test_step(self, batch, batch_idx): ) for metric_name in ("mse", "mae"): - metric_func = metrics.get_metric(metric_name) - batch_metric_vals = metric_func( + batch_metric_vals = self.forecaster.score( prediction, target_states, pred_std, + metric=metrics.get_metric(metric_name), mask=self.interior_mask_bool, sum_vars=False, ) self.test_metrics[metric_name].append(batch_metric_vals) - spatial_loss = self.forecaster.loss( + spatial_loss = self.forecaster.score( prediction, target_states, pred_std, average_grid=False ) log_spatial_losses = spatial_loss[ diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index b2d901e7..51df3261 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -7,6 +7,7 @@ # First-party from neural_lam import config as nlconfig +from neural_lam import metrics from neural_lam.models import ( ARForecaster, DeterministicForecasterModule, @@ -75,6 +76,51 @@ def test_ar_forecaster_unroll(): assert torch.all(prediction[:, :, 1:, :] == 5.0) +def test_ar_forecaster_score(): + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = MockStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + mask = torch.ones(num_grid_nodes, dtype=torch.bool) + + # pred_std=None falls back to forecaster.per_var_std and applies the + # forecaster's own configured scoring rule (self.loss) + scored = forecaster.score(prediction, target, None, mask=mask) + expected = forecaster.loss( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored, expected) + + # An explicit metric overrides self.loss, still substituting the + # per_var_std fallback + scored_mse = forecaster.score( + prediction, target, None, metric=metrics.mse, mask=mask + ) + expected_mse = metrics.mse( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored_mse, expected_mse) + + # An explicit pred_std is used as-is, not overridden by per_var_std + explicit_std = torch.full((d_state,), 2.0) + scored_explicit = forecaster.score( + prediction, target, explicit_std, mask=mask + ) + expected_explicit = forecaster.loss( + prediction, target, explicit_std, mask=mask + ) + assert torch.equal(scored_explicit, expected_explicit) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") From 8b3ab331bc883c87f9aa2d6f0021524190cc7f35 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:29:42 +0530 Subject: [PATCH 15/16] Address PR review: move per_var_std/config validation into the Forecaster BaseForecasterModule.__init__ raised if forecaster.per_var_std was None and the forecaster didn't predict its own std -- a check on the Forecaster's configuration living in the wrong class. Move it into ARForecaster via a new _resolve_pred_std helper, shared by score() and compute_training_loss(): construction with config=None now always succeeds (a valid state for forecasters only ever used for inference), and the ValueError instead fires from the Forecaster itself, only once scoring is attempted without any std to use. --- .../models/forecasters/autoregressive.py | 71 ++++++++++++++++--- neural_lam/models/forecasters/base_module.py | 8 --- .../models/forecasters/probabilistic.py | 7 +- tests/test_prediction_model_classes.py | 42 +++++++++++ 4 files changed, 107 insertions(+), 21 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index ca8166cd..e7add53c 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -40,8 +40,11 @@ def __init__( 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. + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. loss : str, default "wmse" The scoring rule (from ``neural_lam.metrics``) used by ``compute_training_loss`` and stored as ``self.loss``. @@ -232,12 +235,18 @@ def compute_training_loss( batch and time. loss_components : dict of {str: torch.Tensor} Empty; the deterministic objective has no separate components. + + Raises + ------ + ValueError + If the predictor does not output its own std and no + ``per_var_std`` fallback is available; see + ``_resolve_pred_std``. """ prediction, pred_std = self( init_states, forcing_features, target_states ) - if pred_std is None: - pred_std = self.per_var_std + pred_std = self._resolve_pred_std(pred_std) batch_loss = torch.mean( self.loss( @@ -249,6 +258,41 @@ def compute_training_loss( ) return batch_loss, {} + def _resolve_pred_std( + self, pred_std: Optional[torch.Tensor] + ) -> torch.Tensor: + """ + Return ``pred_std``, or the constant ``per_var_std`` fallback. + + Parameters + ---------- + pred_std : torch.Tensor or None + Predicted standard deviation as returned by ``forward``, + possibly ``None``. + + Returns + ------- + torch.Tensor + ``pred_std`` unchanged when given; otherwise ``self.per_var_std``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available (``predictor.predicts_std`` is False and this + forecaster was constructed without ``config``). + """ + if pred_std is not None: + return pred_std + if self.per_var_std is None: + raise ValueError( + "No pred_std available for scoring: predictor.predicts_std " + "is False and this forecaster has no per_var_std fallback " + "(it was constructed without config). Pass config to the " + "constructor, or use a predictor that outputs its own std." + ) + return self.per_var_std + def score( self, prediction: torch.Tensor, @@ -262,10 +306,9 @@ def score( """ Score an already-produced prediction for reporting (not training). - Substitutes ``self.per_var_std`` for ``pred_std`` when the latter is - ``None`` (predictor does not output its own std), then applies - ``metric`` (defaulting to ``self.loss``, the configured scoring - rule). + Resolves ``pred_std`` via ``_resolve_pred_std`` (substituting + ``self.per_var_std`` when ``None``), then applies ``metric`` + (defaulting to ``self.loss``, the configured scoring rule). Parameters ---------- @@ -279,7 +322,8 @@ def score( Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. Predicted standard deviation for ``prediction``; ``None`` when the wrapped predictor does not output one, in which case - ``self.per_var_std`` is substituted. + ``self.per_var_std`` is substituted (see ``_resolve_pred_std`` + for when this raises instead). metric : callable or None, optional Scoring function with the ``neural_lam.metrics`` signature ``(pred, target, pred_std, mask=None, average_grid=True, @@ -296,9 +340,14 @@ def score( torch.Tensor The metric's output; shape depends on ``average_grid`` and ``sum_vars`` (see ``neural_lam.metrics``). + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available; see ``_resolve_pred_std``. """ - if pred_std is None: - pred_std = self.per_var_std + pred_std = self._resolve_pred_std(pred_std) metric_fn = self.loss if metric is None else metric return metric_fn( prediction, diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/forecasters/base_module.py index 85faa451..c77de238 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -133,14 +133,6 @@ 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/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 1c5aeeb4..57526d1c 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -123,8 +123,11 @@ def __init__( 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. + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. loss : str, default "wmse" The scoring rule (from ``neural_lam.metrics``) used by ``compute_training_loss`` and stored as ``self.loss``. diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 51df3261..85bb06d1 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -2,6 +2,7 @@ from argparse import Namespace # Third-party +import pytest import pytorch_lightning as pl import torch @@ -121,6 +122,47 @@ def test_ar_forecaster_score(): assert torch.equal(scored_explicit, expected_explicit) +def test_ar_forecaster_without_config_raises_on_use_not_construction(): + """A predictor that doesn't output std plus no config is a valid, + unambiguous state at construction time (the forecaster may only ever + be used for inference), so ARForecaster must not raise there. It + should only raise once scoring is actually attempted and has no + std to use, and the error should come from the forecaster itself, not + a wrapping module.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + # Construction succeeds even though predicts_std=False and config=None + forecaster = ARForecaster(predictor, datastore) + assert forecaster.per_var_std is None + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.score(prediction, target, None) + + 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 + ) + pred_steps = 3 + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + true_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.compute_training_loss( + init_states, + forcing_features, + true_states, + interior_mask_bool=torch.ones(num_grid_nodes, dtype=torch.bool), + ) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") From 7dc0b906d2615e3cfdcf56e3e3c0a386066b7fdf Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:43:39 +0530 Subject: [PATCH 16/16] Address PR review: move ForecasterModules out of forecasters/ into modules/ BaseForecasterModule, DeterministicForecasterModule and ProbabilisticForecasterModule (Lightning wrappers around a Forecaster) were mixed in with the Forecaster classes themselves under neural_lam/models/forecasters/. Move them into their own neural_lam/models/modules/ package (base.py, deterministic.py, probabilistic.py) so the two concerns: what a forecaster is vs how it's trained/evaluated by Lightning live in separate directories. Pure move: internal imports and neural_lam/models/__init__.py updated accordingly, no behavioural change. --- neural_lam/models/__init__.py | 6 +++--- neural_lam/models/modules/__init__.py | 8 ++++++++ .../{forecasters/base_module.py => modules/base.py} | 2 +- .../deterministic_module.py => modules/deterministic.py} | 2 +- .../probabilistic_module.py => modules/probabilistic.py} | 4 ++-- tests/test_train_model_warnings.py | 2 +- 6 files changed, 16 insertions(+), 8 deletions(-) create mode 100644 neural_lam/models/modules/__init__.py rename neural_lam/models/{forecasters/base_module.py => modules/base.py} (99%) rename neural_lam/models/{forecasters/deterministic_module.py => modules/deterministic.py} (99%) rename neural_lam/models/{forecasters/probabilistic_module.py => modules/probabilistic.py} (98%) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cbeb1b01..be27cfec 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,13 +3,13 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster -from .forecasters.base_module import BaseForecasterModule -from .forecasters.deterministic_module import DeterministicForecasterModule from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, ) -from .forecasters.probabilistic_module import ProbabilisticForecasterModule +from .modules.base import BaseForecasterModule +from .modules.deterministic import DeterministicForecasterModule +from .modules.probabilistic 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/modules/__init__.py b/neural_lam/models/modules/__init__.py new file mode 100644 index 00000000..b2e70700 --- /dev/null +++ b/neural_lam/models/modules/__init__.py @@ -0,0 +1,8 @@ +""" +Lightning modules wrapping forecasters for training and evaluation. +""" + +# Local +from .base import BaseForecasterModule +from .deterministic import DeterministicForecasterModule +from .probabilistic import ProbabilisticForecasterModule diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/modules/base.py similarity index 99% rename from neural_lam/models/forecasters/base_module.py rename to neural_lam/models/modules/base.py index c77de238..0f239ce3 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/modules/base.py @@ -22,7 +22,7 @@ from ...config import NeuralLAMConfig from ...datastore import BaseDatastore from ...weather_dataset import WeatherDataset -from .base import Forecaster +from ..forecasters.base import Forecaster class BaseForecasterModule(pl.LightningModule, ABC): diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/modules/deterministic.py similarity index 99% rename from neural_lam/models/forecasters/deterministic_module.py rename to neural_lam/models/modules/deterministic.py index 217d4937..2842c8f2 100644 --- a/neural_lam/models/forecasters/deterministic_module.py +++ b/neural_lam/models/modules/deterministic.py @@ -12,7 +12,7 @@ # Local from ... import metrics, vis -from .base_module import BaseForecasterModule +from .base import BaseForecasterModule class DeterministicForecasterModule(BaseForecasterModule): diff --git a/neural_lam/models/forecasters/probabilistic_module.py b/neural_lam/models/modules/probabilistic.py similarity index 98% rename from neural_lam/models/forecasters/probabilistic_module.py rename to neural_lam/models/modules/probabilistic.py index cf8d38b0..bcb61395 100644 --- a/neural_lam/models/forecasters/probabilistic_module.py +++ b/neural_lam/models/modules/probabilistic.py @@ -8,8 +8,8 @@ # Local from ... import metrics -from .base_module import BaseForecasterModule -from .probabilistic import ProbabilisticForecaster +from ..forecasters.probabilistic import ProbabilisticForecaster +from .base import BaseForecasterModule class ProbabilisticForecasterModule(BaseForecasterModule): diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index 6bb0654f..44a9f9cb 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -77,7 +77,7 @@ def capture_init(_self, **kwargs): patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), patch("neural_lam.train_model.ARForecaster"), patch( - "neural_lam.models.forecasters.deterministic_module." + "neural_lam.models.modules.deterministic." "DeterministicForecasterModule.__init__", capture_init, ),