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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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, ), From 3e46b5ccac0f6410400a21cb4e09e4ef15b3025d Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 10:22:39 +0530 Subject: [PATCH 17/41] Address PR review: root ensemble RMSE after averaging all samples _ensemble_step took the square root of each batch's MSE and let Lightning average those roots over the epoch, which does not give the RMSE: the root does not commute with the averaging, so the reported number was neither the RMSE nor a mean of RMSEs over a meaningful population. Accumulate only the per-variable squared errors per batch and reduce them once per epoch in a new _log_ensemble_rmse: gather across devices, average over every sample of the epoch and sum over variables, then take the root. Validation logs it from a new on_validation_epoch_end override (before the inherited implementation clears the metric lists) and testing from the existing on_test_epoch_end. This also drops the second metrics.mse call that recomputed, with sum_vars=False, the errors already computed at the top of _ensemble_step; sum_vars=True is exactly that result summed over the variable dimension, so the per-step values are now derived from the single computation. --- neural_lam/models/modules/probabilistic.py | 112 +++++++++++++-------- tests/test_probabilistic_forecaster.py | 52 ++++++++++ 2 files changed, 124 insertions(+), 40 deletions(-) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 83c7d50f..bccab449 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -57,16 +57,18 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): def _ensemble_step(self, batch, phase: str): """ - Sample an ensemble and score its mean against the target states. + Sample an ensemble and compute the per-variable MSE of its mean. 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. 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. + ``self.eval_ensemble_size`` members and scores the ensemble mean + with plain (unweighted) MSE on interior nodes. Only the squared + errors are computed here; they are reduced to an RMSE once per + epoch by ``_log_ensemble_rmse``, since the square root has to be + taken after averaging over every sample rather than per batch. + + This 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. Parameters ---------- @@ -94,40 +96,57 @@ def _ensemble_step(self, batch, phase: str): 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_steps(len(time_step_rmse), phase) + entry_mses = metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + sum_vars=False, + ) # (B, pred_steps, num_state_vars) + self._warn_skipped_steps(entry_mses.shape[1], phase) + + return entry_mses + + def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: + """ + Log the ensemble-mean RMSE accumulated over a full epoch. + + Averages the per-variable MSEs collected by ``_ensemble_step`` over + every sample of the epoch (across both batches and devices) and sums + them over variables, and only then takes the square root. Rooting + each batch's MSE and averaging those roots instead would report a + different quantity, since the square root does not commute with the + averaging. + + Parameters + ---------- + entry_mse_list : list of torch.Tensor + Per-batch per-variable ensemble-mean MSEs, each of shape + ``(B, pred_steps, num_state_vars)``, as returned by + ``_ensemble_step``. + phase : str + Logging phase, either ``"val"`` or ``"test"``. + """ + # Collective: must be reached by every rank, so it precedes the + # rank-zero check below. + entry_mses = self.all_gather_cat(torch.cat(entry_mse_list, dim=0)) + # (total_samples, pred_steps, num_state_vars) + + if not self.trainer.is_global_zero: + return + + time_step_rmse = torch.sqrt(entry_mses.sum(dim=-1).mean(dim=0)) + # (pred_steps,) log_dict = { 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_ens_rmse"] = mean_rmse - self.log_dict( - log_dict, - on_step=False, - on_epoch=True, - sync_dist=True, - batch_size=batch[0].shape[0], - ) - - return metrics.mse( - ensemble_mean, - target_states, - std_placeholder, - mask=self.interior_mask_bool, - sum_vars=False, - ) + log_dict[f"{phase}_mean_ens_rmse"] = torch.mean(time_step_rmse) + # No sync_dist: the values are already gathered above and only rank + # zero reaches this point, so a collective here would hang DDP. + self.log_dict(log_dict, rank_zero_only=True) def validation_step(self, batch, batch_idx): """ @@ -165,15 +184,28 @@ def test_step(self, batch, batch_idx): entry_mses = self._ensemble_step(batch, "test") self.test_metrics["ens_mse"].append(entry_mses) + def on_validation_epoch_end(self): + """ + Perform actions at the end of the validation epoch. + + Logs the epoch's ensemble-mean RMSE, then defers to + ``BaseForecasterModule.on_validation_epoch_end``, which aggregates + the same per-variable MSEs into heatmaps and clears them. + """ + self._log_ensemble_rmse(self.val_metrics["ens_mse"], "val") + super().on_validation_epoch_end() + def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - 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. + Logs the epoch's ensemble-mean RMSE and 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._log_ensemble_rmse(self.test_metrics["ens_mse"], "test") self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") if self.trainer.is_global_zero and self.hparams.metrics_watch: diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index d98cc345..d2b2cbd1 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -1,3 +1,7 @@ +# Standard library +import math +from types import SimpleNamespace + # Third-party import pytest import torch @@ -360,3 +364,51 @@ def test_probabilistic_module_test_step_scores_ensemble_mean(): (entry_mses,) = model.test_metrics["ens_mse"] assert entry_mses.shape == (B, pred_steps, d_state) assert torch.all(torch.isfinite(entry_mses)) + + +def test_ensemble_rmse_takes_root_after_averaging_all_samples(): + """The epoch RMSE roots the mean MSE over every sample, rather than + averaging per-batch roots (the two differ whenever batches disagree).""" + 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( + predictor, datastore, config=config + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=2, + ) + + # Two batches of one sample and one rollout step, with deliberately + # different squared errors so the two aggregation orders disagree. + d_state = datastore.get_num_data_vars(category="state") + model.val_metrics["ens_mse"] = [ + torch.full((1, 1, d_state), 1.0), + torch.full((1, 1, d_state), 9.0), + ] + + captured = {} + model.all_gather_cat = lambda tensor: tensor + model.log_dict = lambda log_dict, **kwargs: captured.update(log_dict) + model._trainer = SimpleNamespace(is_global_zero=True) + + model._log_ensemble_rmse(model.val_metrics["ens_mse"], "val") + + # Summed over variables the batches give MSEs of d_state and 9 * d_state, + # so the correct RMSE roots their mean, 5 * d_state. + assert captured["val_mean_ens_rmse"] == pytest.approx( + math.sqrt(5.0 * d_state) + ) + root_of_each_batch_averaged = ( + math.sqrt(d_state) + math.sqrt(9.0 * d_state) + ) / 2 + assert captured["val_mean_ens_rmse"] != pytest.approx( + root_of_each_batch_averaged + ) From 8f5006ed658e5e58d2185e8b182d9f33a9a97b80 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 10:56:29 +0530 Subject: [PATCH 18/41] Address PR review: separate the training objective from AR unrolling How a forecast is produced and how a training objective is computed from it are orthogonal, but ARForecaster bundled both: auto-regressive unrolling and the deterministic single-forecast loss. That left no way to express an auto-regressive forecaster trained a different way other than inheriting the deterministic objective and overriding it, which is why ProbabilisticARForecaster had to re-declare compute_training_loss abstract after inheriting a concrete one, and why it took a loss argument naming a pointwise scoring rule that will not apply to it. Split the two axes: - ARForecaster now covers only auto-regressive unrolling (predictor, boundary masks, forward) and leaves compute_training_loss abstract. - DeterministicForecaster (new) supplies the objective half: score a single forecast with a configured scoring rule, plus the reporting score() and the pred_std fallback it needs. It makes no assumption about how the forecast is produced. - DeterministicARForecaster combines the two and is the deterministic model the CLI builds; this is the rename of the old ARForecaster. - ProbabilisticARForecaster is now ARForecaster + ProbabilisticForecaster, so it never inherits a concrete compute_training_loss to re-abstract, and takes neither config nor loss: a concrete probabilistic forecaster brings whatever configuration its own objective needs. The per-variable std computation moves to loss_weighting.get_per_var_std so both objective families can reuse it, and the abstract score() drops off Forecaster since it is specific to the deterministic objective (only DeterministicForecasterModule calls it; the probabilistic module scores ensembles with metrics.mse directly). --- neural_lam/loss_weighting.py | 38 +++ neural_lam/models/__init__.py | 4 + neural_lam/models/forecasters/__init__.py | 1 + .../models/forecasters/autoregressive.py | 225 +------------- neural_lam/models/forecasters/base.py | 73 +---- .../models/forecasters/deterministic.py | 283 ++++++++++++++++++ .../models/forecasters/probabilistic.py | 71 +---- neural_lam/models/modules/deterministic.py | 13 +- neural_lam/train_model.py | 10 +- tests/test_checkpoint.py | 6 +- tests/test_datasets.py | 4 +- tests/test_gnn_layers.py | 4 +- tests/test_gpu_normalization.py | 4 +- tests/test_plotting.py | 12 +- tests/test_prediction_model_classes.py | 34 ++- tests/test_probabilistic_forecaster.py | 39 ++- tests/test_train_model_warnings.py | 2 +- tests/test_training.py | 6 +- 18 files changed, 440 insertions(+), 389 deletions(-) create mode 100644 neural_lam/models/forecasters/deterministic.py diff --git a/neural_lam/loss_weighting.py b/neural_lam/loss_weighting.py index 8a37d07f..5a9910e3 100644 --- a/neural_lam/loss_weighting.py +++ b/neural_lam/loss_weighting.py @@ -1,5 +1,8 @@ """Utility functions for configuring state-feature loss weighting.""" +# Third-party +import torch + # Local from .config import ( ManualStateFeatureWeighting, @@ -118,3 +121,38 @@ def get_state_feature_weighting( ) return weights + + +def get_per_var_std( + config: NeuralLAMConfig, datastore: BaseDatastore +) -> torch.Tensor: + """ + Return the constant per-variable standard deviation of the one-step + difference, weighted by the configured state feature weighting. + + Forecasters whose predictor does not output its own standard deviation + substitute this for ``pred_std`` when applying a scoring rule. + + Parameters + ---------- + config : NeuralLAMConfig + Configuration object for neural-lam, supplying the state feature + weighting. + datastore : BaseDatastore + Datastore object containing the state standardization statistics. + + Returns + ------- + torch.Tensor + Shape ``(num_state_vars,)``. Per-variable standard deviation. + """ + da_state_stats = datastore.get_standardization_dataarray(category="state") + diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + feature_weights = torch.tensor( + get_state_feature_weighting(config=config, datastore=datastore), + dtype=torch.float32, + ) + return diff_std / torch.sqrt(feature_weights) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index be27cfec..f1b7a2de 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,6 +3,10 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.deterministic import ( + DeterministicARForecaster, + DeterministicForecaster, +) from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 254c4ba0..9590b33f 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -5,4 +5,5 @@ # Local from .autoregressive import ARForecaster from .base import Forecaster +from .deterministic import DeterministicARForecaster, DeterministicForecaster from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index e7add53c..671c1400 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,32 +1,32 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" -# Standard library -from typing import Callable, Optional - # 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 class ARForecaster(Forecaster): """ - Subclass of Forecaster that uses an auto-regressive strategy to - unroll a forecast. Makes use of a StepPredictor at each AR step. + Forecaster that produces a forecast by auto-regressive unrolling, using + a StepPredictor at each AR step. + + This class fixes only *how forecasts are produced*, and deliberately says + nothing about how a training objective is computed from them: it leaves + ``compute_training_loss`` abstract. The two are orthogonal, so the + objective is mixed in separately (see ``DeterministicForecaster`` and + ``ProbabilisticForecaster``), which lets an auto-regressive forecaster be + trained deterministically or probabilistically, and equally lets a + non-auto-regressive forecaster reuse either objective. """ def __init__( self, predictor: StepPredictor, datastore: BaseDatastore, - config: NeuralLAMConfig | None = None, - loss: str = "wmse", ) -> None: """ Initialize the ARForecaster. @@ -37,17 +37,6 @@ 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``). 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``. """ super().__init__() self.predictor = predictor @@ -63,31 +52,6 @@ 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: """ @@ -188,172 +152,3 @@ 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, - interior_mask_bool: torch.Tensor, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - """ - 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 - 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. - 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. - - 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. - - 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 - ) - pred_std = self._resolve_pred_std(pred_std) - - batch_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=interior_mask_bool, - ) - ) - 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, - 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). - - 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 - ---------- - 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 (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, - 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``). - - Raises - ------ - ValueError - If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is - available; see ``_resolve_pred_std``. - """ - pred_std = self._resolve_pred_std(pred_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 da63869b..ccd8ff6e 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, Optional # Third-party import torch @@ -76,10 +75,9 @@ 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 forecaster's - own constant per-variable std fallback is substituted by - ``compute_training_loss``/``score``, not by the caller. Dims: - same as ``prediction``. + predicted standard deviation; when ``None``, substituting a + fallback std is left to whatever consumes the forecast, not to + the caller of ``forward``. Dims: same as ``prediction``. """ @abstractmethod @@ -95,10 +93,14 @@ def compute_training_loss( 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, using its own - ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The - wrapping ``BaseForecasterModule`` only injects the interior mask, - logs the returned components and optimizes the returned loss. + how to combine those terms into a single scalar. The wrapping + ``BaseForecasterModule`` only injects the interior mask, logs the + returned components and optimizes the returned loss. + + How the objective is computed is orthogonal to how forecasts are + produced, so implementations are mixed in separately from the + ``forward`` implementation; see ``DeterministicForecaster`` for the + single-forecast scoring-rule objective. Parameters ---------- @@ -136,56 +138,3 @@ 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/deterministic.py b/neural_lam/models/forecasters/deterministic.py new file mode 100644 index 00000000..e17fe82f --- /dev/null +++ b/neural_lam/models/forecasters/deterministic.py @@ -0,0 +1,283 @@ +"""Forecasters trained by scoring a single deterministic forecast.""" + +# Standard library +from typing import Callable, Optional + +# Third-party +import torch + +# Local +from ... import metrics +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ...loss_weighting import get_per_var_std +from ..step_predictors.base import StepPredictor +from .autoregressive import ARForecaster +from .base import Forecaster + + +class DeterministicForecaster(Forecaster): + """ + Forecaster whose training objective is a scoring rule applied to a + single forecast. + + Supplies the objective half of a forecaster: ``compute_training_loss`` + produces one forecast and scores it, and ``score`` applies a metric to + an already-produced forecast for reporting. Neither makes any assumption + about *how* that forecast is produced, so this composes with any way of + implementing ``forward`` (see ``DeterministicARForecaster`` for the + auto-regressive combination). + + Concrete subclasses must call ``_configure_scoring`` from their + ``__init__`` to set up the scoring rule and the ``pred_std`` fallback. + """ + + def _configure_scoring( + self, + datastore: BaseDatastore, + config: NeuralLAMConfig | None, + loss: str, + ) -> None: + """ + Set up the scoring rule and the constant ``pred_std`` fallback. + + Called by concrete subclasses from ``__init__``, after + ``torch.nn.Module`` initialization (buffers are registered here). + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing the state standardization statistics + used to compute ``per_var_std``. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when the forecast carries no std of + its own. Required in that case for ``score`` and + ``compute_training_loss`` to work (they raise ``ValueError`` via + ``_resolve_pred_std`` otherwise); forecasters used purely for + inference can omit it. + loss : str + The scoring rule (from ``neural_lam.metrics``) applied by + ``compute_training_loss``, stored as ``self.loss``. + """ + self.loss = metrics.get_metric(loss) + + # Store per_var_std only if the forecast carries no std of its own + if not self.predicts_std and config is not None: + self.register_buffer( + "per_var_std", + get_per_var_std(config=config, datastore=datastore), + persistent=False, + ) + else: + self.per_var_std = None + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score a single forecast with ``self.loss``. + + Produces one 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 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 while forecasting. + Dims: same as the prediction. + 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. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the forecast, averaged over + batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; the deterministic objective has no separate components. + + Raises + ------ + ValueError + If the forecast carries no std of its own and no + ``per_var_std`` fallback is available; see + ``_resolve_pred_std``. + """ + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + pred_std = self._resolve_pred_std(pred_std) + + batch_loss = torch.mean( + self.loss( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + 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 (``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, + 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). + + 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 + ---------- + 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 forecast carries no std, in which case ``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, + 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``). + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available; see ``_resolve_pred_std``. + """ + pred_std = self._resolve_pred_std(pred_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, + ) + + +class DeterministicARForecaster(DeterministicForecaster, ARForecaster): + """ + Auto-regressive forecaster trained by scoring its single rollout. + + Combines the two orthogonal halves: ``ARForecaster`` supplies the + auto-regressive ``forward``, ``DeterministicForecaster`` supplies the + single-forecast training objective and the reporting ``score``. + """ + + def __init__( + self, + predictor: StepPredictor, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + ) -> None: + """ + Initialize the DeterministicARForecaster. + + Parameters + ---------- + predictor : StepPredictor + The predictor to use for each AR 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. Required in that case for ``score`` and + ``compute_training_loss`` to work (they raise ``ValueError`` + otherwise); forecasters used purely for inference (``forward``) + can omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) applied by + ``compute_training_loss``. + """ + # DeterministicForecaster defines no __init__, so this initializes + # the AR half (and torch.nn.Module) before any buffer is registered + super().__init__(predictor, datastore) + self._configure_scoring(datastore, config=config, loss=loss) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 57526d1c..33f5eaff 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -7,9 +7,6 @@ import torch # Local -from ...config import NeuralLAMConfig -from ...datastore import BaseDatastore -from ..step_predictors.base import StepPredictor from .autoregressive import ARForecaster from .base import Forecaster @@ -91,49 +88,17 @@ class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): trajectory. This class adds ensemble forecasting on top: unrolling several trajectories and stacking them along an ensemble dimension. - ``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). + It supplies no training objective, and so remains abstract in + ``compute_training_loss``. There is no default 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 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), along with + whatever configuration that objective needs. """ - def __init__( - self, - predictor: StepPredictor, - datastore: BaseDatastore, - config: NeuralLAMConfig | None = None, - loss: str = "wmse", - ) -> 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. - 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``). 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``. - """ - super().__init__(predictor, datastore, config=config, loss=loss) - def sample_ensemble( self, init_states: torch.Tensor, @@ -210,21 +175,3 @@ def sample_ensemble( 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, - forcing_features: torch.Tensor, - target_states: torch.Tensor, - interior_mask_bool: torch.Tensor, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - """ - Compute the training objective for one batch. - - 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). - """ diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 817abdc9..8833d237 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -12,6 +12,7 @@ # Local from ... import metrics, vis +from ..forecasters.deterministic import DeterministicForecaster from .base import BaseForecasterModule @@ -19,13 +20,15 @@ class DeterministicForecasterModule(BaseForecasterModule): """ Lightning module for a single deterministic forecast per batch. - Validation and testing score the forecaster's own single-rollout - 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``). + Validation and testing score the forecaster's own single 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``). """ + # score() is supplied by the deterministic objective mixin + forecaster: DeterministicForecaster + def __init__(self, *args, **kwargs): """ Initialize the module and its deterministic evaluation metrics. diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 26e70d94..7514c8bf 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -19,7 +19,11 @@ from . import utils from .config import load_config_and_datastore from .gnn_layers import GNN_TYPES -from .models import MODELS, ARForecaster, DeterministicForecasterModule +from .models import ( + MODELS, + DeterministicARForecaster, + DeterministicForecasterModule, +) from .weather_dataset import WeatherDataModule @@ -63,7 +67,7 @@ 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( + forecaster = DeterministicARForecaster( predictor, datastore, config=config, loss=args.loss ) return DeterministicForecasterModule.load_from_checkpoint( @@ -476,7 +480,7 @@ 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( + forecaster = DeterministicARForecaster( predictor, datastore, config=config, loss=args.loss ) diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 6644aae6..40ee58f5 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -9,7 +9,7 @@ from neural_lam import config as nlconfig from neural_lam.create_graph import create_graph_from_datastore from neural_lam.models import ( - ARForecaster, + DeterministicARForecaster, DeterministicForecasterModule, GraphLAM, ) @@ -54,7 +54,9 @@ 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, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) model = DeterministicForecasterModule( forecaster=forecaster, config=config, diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 319e3372..91788d2b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -219,7 +219,7 @@ def _create_graph(): dataset = WeatherDataset(datastore=datastore, split=split, ar_steps=2) # First-party - from neural_lam.models import MODELS, ARForecaster + from neural_lam.models import MODELS, DeterministicARForecaster predictor_class = MODELS["graph_lam"] predictor = predictor_class( @@ -235,7 +235,7 @@ def _create_graph(): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster( + forecaster = DeterministicARForecaster( predictor, datastore=datastore, config=config, loss=args.loss ) diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 789297d0..dbd0e86b 100644 --- a/tests/test_gnn_layers.py +++ b/tests/test_gnn_layers.py @@ -8,7 +8,7 @@ from neural_lam import config as nlconfig from neural_lam.create_graph import create_graph_from_datastore from neural_lam.gnn_layers import InteractionNet, PropagationNet -from neural_lam.models import MODELS, ARForecaster +from neural_lam.models import MODELS, DeterministicARForecaster from tests.conftest import init_datastore_example @@ -73,7 +73,7 @@ def _build_model_and_data( output_clamping_upper=config.training.output_clamping.upper, **gnn_kwargs, ) - forecaster = ARForecaster(predictor, datastore, config=config) + forecaster = DeterministicARForecaster(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 dcc8ba63..4e0b8c31 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -5,7 +5,7 @@ # First-party from neural_lam import config as nlconfig from neural_lam.models import ( - ARForecaster, + DeterministicARForecaster, DeterministicForecasterModule, StepPredictor, ) @@ -31,7 +31,7 @@ def _build_module(datastore): ) ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore, config=config) + forecaster = DeterministicARForecaster(predictor, datastore, config=config) return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index f55755e5..f6e3d495 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -19,7 +19,7 @@ from neural_lam import vis from neural_lam.create_graph import create_graph_from_datastore from neural_lam.models import ( - ARForecaster, + DeterministicARForecaster, DeterministicForecasterModule, GraphLAM, ) @@ -454,7 +454,7 @@ class ModelArgs: # Create model # First-party - from neural_lam.models import MODELS, ARForecaster + from neural_lam.models import MODELS, DeterministicARForecaster args = ModelArgs() predictor_class = MODELS["graph_lam"] @@ -471,7 +471,7 @@ class ModelArgs: output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster( + forecaster = DeterministicARForecaster( predictor, datastore=datastore, config=config, loss=args.loss ) @@ -531,7 +531,7 @@ def test_plot_examples_integration_saves_figure( ), f"Expected time_step_unit={time_unit}, got {model.time_step_unit}" # Generate prediction - (init_states, target, forcing_features, _batch_times) = batch + init_states, target, forcing_features, _batch_times = batch prediction, _ = model.forecaster(init_states, forcing_features, target) # Rescale to original data scale @@ -684,7 +684,9 @@ 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, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) return DeterministicForecasterModule( forecaster=forecaster, config=config, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 85bb06d1..0cfbdcca 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -10,7 +10,7 @@ from neural_lam import config as nlconfig from neural_lam import metrics from neural_lam.models import ( - ARForecaster, + DeterministicARForecaster, DeterministicForecasterModule, StepPredictor, ) @@ -46,7 +46,7 @@ def test_ar_forecaster_unroll(): output_std=False, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = DeterministicARForecaster(predictor, datastore) # Override masks to test boundary masking behaviour forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) @@ -85,7 +85,9 @@ def test_ar_forecaster_score(): ) ) predictor = MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) B, num_grid_nodes = 2, predictor.num_grid_nodes d_state = datastore.get_num_data_vars(category="state") @@ -125,15 +127,15 @@ def test_ar_forecaster_score(): 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.""" + be used for inference), so DeterministicARForecaster 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) + forecaster = DeterministicARForecaster(predictor, datastore) assert forecaster.per_var_std is None B, num_grid_nodes = 2, predictor.num_grid_nodes @@ -189,7 +191,9 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) model = DeterministicForecasterModule( forecaster=forecaster, @@ -224,7 +228,7 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster( + load_forecaster = DeterministicARForecaster( load_predictor, datastore, config=config, loss="mse" ) @@ -294,7 +298,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): saved_val_steps = [2] saved_n_example_pred = 7 - forecaster = ARForecaster( + forecaster = DeterministicARForecaster( predictor, datastore, config=config, loss=saved_loss ) @@ -363,7 +367,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster( + load_forecaster = DeterministicARForecaster( load_predictor, datastore, config=config, loss=saved_loss ) @@ -451,7 +455,7 @@ def get_dataarray(self, category, split=None, standardize=False): assert predictor.grid_static_features.shape[1] == 0 - forecaster = ARForecaster(predictor, datastore) + forecaster = DeterministicARForecaster(predictor, datastore) B = 2 num_grid_nodes = predictor.num_grid_nodes d_state = base_datastore.get_num_data_vars(category="state") @@ -485,8 +489,8 @@ def test_step_predictor_no_static_features(): 0, ) - # Verify a forward pass works end-to-end via ARForecaster - forecaster = ARForecaster(predictor, datastore) + # Verify a forward pass works end-to-end via DeterministicARForecaster + forecaster = DeterministicARForecaster(predictor, datastore) B, num_grid_nodes = 2, predictor.num_grid_nodes d_state = datastore.get_num_data_vars(category="state") d_forcing = datastore.get_num_data_vars(category="forcing") diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index d2b2cbd1..9cc09674 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -10,8 +10,9 @@ # First-party from neural_lam import config as nlconfig from neural_lam import metrics +from neural_lam.loss_weighting import get_per_var_std from neural_lam.models import ( - ARForecaster, + DeterministicARForecaster, DeterministicForecasterModule, ProbabilisticARForecaster, ProbabilisticForecasterModule, @@ -45,16 +46,30 @@ 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. - ``sample_ensemble`` always requires an explicit member count, so this - class takes its own ``train_num_members`` for the training objective. + ``ProbabilisticARForecaster`` supplies no training objective (no single + default 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. Being the concrete class, it also + owns whatever that objective needs: the scoring rule, the constant + per-variable std, and the member count to train on (``sample_ensemble`` + always requires an explicit one). """ - def __init__(self, *args, train_num_members: int = 2, **kwargs): - super().__init__(*args, **kwargs) + def __init__( + self, + predictor, + datastore, + config=None, + loss: str = "wmse", + train_num_members: int = 2, + ): + super().__init__(predictor, datastore) + self.loss = metrics.get_metric(loss) + self.per_var_std = ( + get_per_var_std(config=config, datastore=datastore) + if config is not None + else None + ) self.train_num_members = train_num_members def compute_training_loss( @@ -105,7 +120,7 @@ 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, loss="mse") + forecaster = DeterministicARForecaster(predictor, datastore, loss="mse") init_states, forcing_features, target_states = _example_batch(datastore) score_metric = metrics.get_metric("mse") @@ -236,7 +251,9 @@ def test_module_training_step_delegates_to_forecaster(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) model = DeterministicForecasterModule( forecaster=forecaster, config=config, diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index 270c04cf..9642acfe 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -83,7 +83,7 @@ def capture_init(_self, **kwargs): ), patch("neural_lam.train_model.WeatherDataModule"), patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), - patch("neural_lam.train_model.ARForecaster"), + patch("neural_lam.train_model.DeterministicARForecaster"), patch( "neural_lam.models.modules.deterministic." "DeterministicForecasterModule.__init__", diff --git a/tests/test_training.py b/tests/test_training.py index aca432c5..e660eef4 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -107,7 +107,7 @@ def run_simple_training( # Build predictor and forecaster externally, then inject into # DeterministicForecasterModule # First-party - from neural_lam.models import MODELS, ARForecaster + from neural_lam.models import MODELS, DeterministicARForecaster predictor_class = MODELS["graph_lam"] predictor = predictor_class( @@ -123,7 +123,9 @@ def run_simple_training( output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) model = DeterministicForecasterModule( forecaster=forecaster, From 5a6bd6d8deb01619867245a38a575e34b4e483e6 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 11:33:52 +0530 Subject: [PATCH 19/41] Address PR review: restore per-step training loss logging Moving the objective onto the Forecaster left --train_steps_to_log accepted but inert, since compute_training_loss returns only a scalar. That is the right general contract (an ELBO need not decompose over rollout steps), but it dropped a breakdown every model in the repo can actually produce. Give the decomposition a home on the class that knows the objective has one: DeterministicForecaster.compute_step_losses returns the scoring rule per predicted step, and compute_training_loss is now its mean, so the two cannot drift. DeterministicForecasterModule overrides training_step to log the breakdown through the existing _log_step_loss helper, producing the same train_loss_unroll{i} keys as before. The general BaseForecasterModule.training_step still logs only the scalar, so a forecaster whose objective does not decompose is unaffected. This deliberately sits on the deterministic objective rather than on the auto-regressive rollout: a direct (non-AR) forecaster scored per lead time decomposes just as well, so AR-ness is not what makes the breakdown available. --- .../models/forecasters/deterministic.py | 81 ++++++++++++++++--- neural_lam/models/modules/base.py | 14 ++-- neural_lam/models/modules/deterministic.py | 36 +++++++++ tests/test_probabilistic_forecaster.py | 53 ++++++++++++ 4 files changed, 166 insertions(+), 18 deletions(-) diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index e17fe82f..9aec50bc 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -72,6 +72,66 @@ def _configure_scoring( else: self.per_var_std = None + def compute_step_losses( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> torch.Tensor: + """ + Score a single forecast with ``self.loss``, per predicted step. + + This objective is a per-step scoring rule averaged over the rollout, + so it decomposes into the contribution of each predicted step, which + callers can report individually. ``compute_training_loss`` is the + mean of what this returns. + + 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. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. + 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 while forecasting. + 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. + + Returns + ------- + torch.Tensor + Shape ``(pred_steps,)``. The scoring rule at each predicted + step, averaged over the batch. + + Raises + ------ + ValueError + If the forecast carries no std of its own and no + ``per_var_std`` fallback is available; see + ``_resolve_pred_std``. + """ + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + pred_std = self._resolve_pred_std(pred_std) + + return torch.mean( + self.loss( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ), + dim=0, + ) + def compute_training_loss( self, init_states: torch.Tensor, @@ -84,6 +144,8 @@ def compute_training_loss( Produces one forecast over the full rollout, scores it against the target states on interior nodes and averages over batch and time. + Callers wanting the per-step breakdown of this same objective + should use ``compute_step_losses`` instead, which this averages. Parameters ---------- @@ -125,20 +187,13 @@ def compute_training_loss( ``per_var_std`` fallback is available; see ``_resolve_pred_std``. """ - prediction, pred_std = self( - init_states, forcing_features, target_states - ) - pred_std = self._resolve_pred_std(pred_std) - - batch_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=interior_mask_bool, - ) + step_losses = self.compute_step_losses( + init_states, + forcing_features, + target_states, + interior_mask_bool, ) - return batch_loss, {} + return torch.mean(step_losses), {} def _resolve_pred_std( self, pred_std: Optional[torch.Tensor] diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index 94823c6f..fb7d28fc 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -348,11 +348,15 @@ def training_step(self, batch): 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. Unlike validation/test, no per-step ``train_loss_unroll{i}`` - breakdown is logged here even when ``train_steps_to_log`` is set: - ``compute_training_loss`` intentionally returns only a scalar (a - forecaster's objective, e.g. an ELBO, may not decompose per rollout - step at all), so there is no per-step tensor to select from. + returns. + + This general implementation logs only the scalar objective, since + ``compute_training_loss`` returns only a scalar: a forecaster's + objective (e.g. an ELBO) need not decompose over rollout steps, so + there is no per-step tensor to select from with + ``train_steps_to_log``. Subclasses whose forecaster does expose such + a decomposition override this to log the breakdown as well; see + ``DeterministicForecasterModule.training_step``. Parameters ---------- diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 8833d237..d257ba3b 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -57,6 +57,42 @@ def __init__(self, *args, **kwargs): # For storing spatial loss maps during evaluation self.spatial_loss_maps: list[Any] = [] + def training_step(self, batch): + """ + Perform a single training step, logging the per-step breakdown. + + Overrides ``BaseForecasterModule.training_step``, which logs only the + scalar objective because a forecaster's training loss need not + decompose over rollout steps. The deterministic objective is a + per-step scoring rule averaged over the rollout, so it does, and + ``--train_steps_to_log`` selects which of those steps to report as + ``train_loss_unroll{i}``. The logged ``train_loss`` is the mean of + the per-step losses, i.e. exactly + ``forecaster.compute_training_loss``. + + Parameters + ---------- + batch : tuple + The batch of data. + + Returns + ------- + torch.Tensor + The computed loss for the training step. + """ + init_states, target_states, forcing_features, _ = batch + time_step_loss = self.forecaster.compute_step_losses( + init_states, + forcing_features, + target_states, + interior_mask_bool=self.interior_mask_bool, + ) + batch_loss = torch.mean(time_step_loss) + self._log_step_loss( + time_step_loss, batch_loss, "train", batch[0].shape[0] + ) + return batch_loss + def validation_step(self, batch, batch_idx): """ Perform a single validation step. diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 9cc09674..b512d976 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -276,6 +276,59 @@ def test_module_training_step_delegates_to_forecaster(): torch.testing.assert_close(batch_loss, expected_loss) +def test_deterministic_training_step_logs_per_step_losses(): + """--train_steps_to_log selects per-step training losses to report, + which the deterministic objective can supply because it decomposes + over rollout steps.""" + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = DeterministicARForecaster( + predictor, datastore, config=config, loss="mse" + ) + model = DeterministicForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + train_steps_to_log=[1, 3], + ) + + pred_steps = 3 + init_states, forcing_features, target_states = _example_batch( + datastore, pred_steps=pred_steps + ) + batch_times = torch.zeros(init_states.shape[0], pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + captured = {} + model.log_dict = lambda log_dict, **kwargs: captured.update(log_dict) + + batch_loss = model.training_step(batch) + + assert set(captured) == { + "train_loss", + "train_loss_unroll1", + "train_loss_unroll3", + } + torch.testing.assert_close(captured["train_loss"], batch_loss) + + # The reported steps are the corresponding entries of the same + # decomposition the logged train_loss averages + step_losses = forecaster.compute_step_losses( + init_states, + forcing_features, + target_states, + interior_mask_bool=model.interior_mask_bool, + ) + torch.testing.assert_close(captured["train_loss_unroll1"], step_losses[0]) + torch.testing.assert_close(captured["train_loss_unroll3"], step_losses[2]) + torch.testing.assert_close(batch_loss, torch.mean(step_losses)) + + class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster): """ProbabilisticARForecaster recording the requested member count.""" From e6961eda6c23ee9fcbcec66605cfde265bc0fed7 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 11:49:17 +0530 Subject: [PATCH 20/41] Address PR review: restore _compute_prediction_and_loss for val/test Deleting this helper during the module split reintroduced the duplication #675 had removed: validation_step and test_step each repeated the same common_step, score, mean-over-batch sequence. Restore it on DeterministicForecasterModule, which is the only module that scores a single prediction this way (the probabilistic one samples and scores ensembles instead, and nothing else calls common_step). training_step is deliberately not folded back in, so the helper now serves two callers rather than the original three. It goes through the forecaster's training objective rather than the reporting scoring rule used here; those two happen to coincide for the deterministic objective, but routing training through score() would put the choice of objective back in the module, which is what this PR moves onto the Forecaster. --- neural_lam/models/modules/deterministic.py | 63 ++++++++++++++++------ 1 file changed, 48 insertions(+), 15 deletions(-) diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index d257ba3b..1142cb8a 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -93,16 +93,39 @@ def training_step(self, batch): ) return batch_loss - def validation_step(self, batch, batch_idx): + def _compute_prediction_and_loss( + self, + batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ - Perform a single validation step. + Compute predicted mean, standard deviation, and step-wise loss. + Also extract and return corresponding target from batch. + + Shared by ``validation_step`` and ``test_step``. ``training_step`` + deliberately does not use this: it goes through the forecaster's + training objective rather than the reporting scoring rule applied + here, and those are only interchangeable for this particular + objective. Parameters ---------- - batch : tuple + batch : tuple of torch.Tensor The batch of data. - batch_idx : int - The index of the batch. + + Returns + ------- + prediction : torch.Tensor + Model predictions, shape + ``(B, pred_steps, num_grid_nodes, num_state_vars)``. + target_states : torch.Tensor + Target states, shape + ``(B, pred_steps, num_grid_nodes, num_state_vars)``. + pred_std : torch.Tensor + Predicted or pre-defined standard deviation, shape + ``(B, pred_steps, num_grid_nodes, num_state_vars)`` or + ``(num_state_vars,)``. + time_step_loss : torch.Tensor + Loss for each unroll step, shape ``(pred_steps,)``. """ prediction, target_states, pred_std, _ = self.common_step(batch) @@ -115,6 +138,23 @@ def validation_step(self, batch, batch_idx): ), dim=0, ) + return prediction, target_states, pred_std, time_step_loss + + 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, time_step_loss = ( + self._compute_prediction_and_loss(batch) + ) + mean_loss = torch.mean(time_step_loss) batch_size = batch[0].shape[0] self._log_step_loss(time_step_loss, mean_loss, "val", batch_size) @@ -141,7 +181,9 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, _ = self.common_step(batch) + prediction, target_states, pred_std, time_step_loss = ( + self._compute_prediction_and_loss(batch) + ) if pred_std is not None: mean_pred_std = torch.mean( @@ -149,15 +191,6 @@ def test_step(self, batch, batch_idx): ) self.test_metrics["output_std"].append(mean_pred_std) - time_step_loss = torch.mean( - self.forecaster.score( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) mean_loss = torch.mean(time_step_loss) batch_size = batch[0].shape[0] self._log_step_loss(time_step_loss, mean_loss, "test", batch_size) From 56b72770adb9eebb0a69161d1d49e36d68bd0386 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 12:01:37 +0530 Subject: [PATCH 21/41] Address PR review: fix stale val_metrics docs, justify abstract eval steps The abstract validation_step/test_step docstrings pointed at self.val_metrics/self.test_metrics as if this class owned them. It does consume both in its epoch-end hooks but never creates them, since which metrics are collected depends on the evaluation mode. Declare them as annotations here to state that contract explicitly, and reword the docstrings to refer to the ones the subclass creates. Keep the abstract declarations, and record why inline. LightningModule does define validation_step and test_step, but as no-op stubs rather than abstract methods (only training_step even warns), so dropping these would let a module that omits them instantiate happily and silently skip evaluation instead of failing at construction. Added a test covering that. --- neural_lam/models/modules/base.py | 29 ++++++++++++++++++++------ tests/test_probabilistic_forecaster.py | 21 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index fb7d28fc..9eaff8dc 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -40,6 +40,12 @@ class BaseForecasterModule(pl.LightningModule, ABC): ``ProbabilisticForecasterModule``) rather than overriding one another. """ + # Which metrics are collected differs per evaluation mode, so concrete + # subclasses create these in __init__; the epoch-end hooks here consume + # whatever they contain. Declared for the contract only, not assigned. + val_metrics: dict[str, list] + test_metrics: dict[str, list] + # pylint: disable=arguments-differ def __init__( @@ -472,9 +478,15 @@ 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``. + Concrete subclasses must both score the batch and populate the + ``val_metrics`` they created, which ``on_validation_epoch_end`` + below aggregates. + + Kept abstract even though ``LightningModule`` defines this method: + that definition is a no-op stub, not an abstract method, so without + this declaration a subclass that omitted it would instantiate + happily and silently skip validation. Declaring it abstract turns + that into a ``TypeError`` at construction. Parameters ---------- @@ -511,9 +523,14 @@ 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``. + Concrete subclasses must both score the batch and populate the + ``test_metrics`` they created, which their ``on_test_epoch_end`` + aggregates. + + Kept abstract for the same reason as ``validation_step``: the + ``LightningModule`` definition is a no-op stub rather than an + abstract method, so omitting it would silently skip testing instead + of failing at construction. Parameters ---------- diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index b512d976..0ca1a2c1 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -12,6 +12,7 @@ from neural_lam import metrics from neural_lam.loss_weighting import get_per_var_std from neural_lam.models import ( + BaseForecasterModule, DeterministicARForecaster, DeterministicForecasterModule, ProbabilisticARForecaster, @@ -242,6 +243,26 @@ def test_probabilistic_ar_forecaster_is_abstract(): ProbabilisticARForecaster(predictor, datastore) +def test_forecaster_module_evaluation_steps_are_abstract(): + """A module omitting validation_step/test_step fails at construction. + + LightningModule defines both as no-op stubs rather than abstract + methods, so without BaseForecasterModule declaring them abstract such a + module would instantiate happily and silently skip evaluation. + """ + + class MissingEvaluationSteps(BaseForecasterModule): + pass + + assert { + "validation_step", + "test_step", + } <= MissingEvaluationSteps.__abstractmethods__ + + with pytest.raises(TypeError, match="validation_step|test_step"): + MissingEvaluationSteps(forecaster=None, config=None, datastore=None) + + def test_module_training_step_delegates_to_forecaster(): datastore = init_datastore_example("mdp") predictor = ZeroStepPredictor(datastore=datastore, output_std=False) From a5dc2927996dfabdbfcc7c9c473f9ef4f611920d Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 12:23:10 +0530 Subject: [PATCH 22/41] Address PR review: say forecast rather than rollout where AR is not implied Apply the suggested module docstring wording, and carry the same correction through the surrounding docs: DeterministicForecasterModule and DeterministicForecaster describe a single forecast scored per predicted step, which need not have been produced by unrolling. "Rollout" is kept where it is accurate, i.e. on DeterministicARForecaster and the AR machinery itself. --- neural_lam/models/forecasters/deterministic.py | 16 ++++++++-------- neural_lam/models/modules/base.py | 2 +- neural_lam/models/modules/deterministic.py | 8 ++++---- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 9aec50bc..5de1b385 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -82,10 +82,10 @@ def compute_step_losses( """ Score a single forecast with ``self.loss``, per predicted step. - This objective is a per-step scoring rule averaged over the rollout, - so it decomposes into the contribution of each predicted step, which - callers can report individually. ``compute_training_loss`` is the - mean of what this returns. + This objective is a per-step scoring rule averaged over the + predicted steps, so it decomposes into the contribution of each of + them, which callers can report individually. + ``compute_training_loss`` is the mean of what this returns. Parameters ---------- @@ -142,9 +142,9 @@ def compute_training_loss( """ Score a single forecast with ``self.loss``. - Produces one forecast over the full rollout, scores it against the - target states on interior nodes and averages over batch and time. - Callers wanting the per-step breakdown of this same objective + Produces one forecast over every predicted step, scores it against + the target states on interior nodes and averages over batch and + time. Callers wanting the per-step breakdown of this same objective should use ``compute_step_losses`` instead, which this averages. Parameters @@ -158,7 +158,7 @@ def compute_training_loss( 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, + is batch size, ``pred_steps`` is the number of predicted steps, ``num_grid_nodes`` is the number of spatial nodes, and ``num_forcing_vars`` is the forcing feature dimension (already concatenated past/current/future windows). diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index 9eaff8dc..c269662a 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -358,7 +358,7 @@ def training_step(self, batch): This general implementation logs only the scalar objective, since ``compute_training_loss`` returns only a scalar: a forecaster's - objective (e.g. an ELBO) need not decompose over rollout steps, so + objective (e.g. an ELBO) need not decompose over predicted steps, so there is no per-step tensor to select from with ``train_steps_to_log``. Subclasses whose forecaster does expose such a decomposition override this to log the breakdown as well; see diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 1142cb8a..f112557b 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -1,5 +1,5 @@ """Lightning module evaluating forecasters through a single deterministic -rollout per batch.""" +forecast per batch.""" # Standard library import os @@ -63,9 +63,9 @@ def training_step(self, batch): Overrides ``BaseForecasterModule.training_step``, which logs only the scalar objective because a forecaster's training loss need not - decompose over rollout steps. The deterministic objective is a - per-step scoring rule averaged over the rollout, so it does, and - ``--train_steps_to_log`` selects which of those steps to report as + decompose over predicted steps. The deterministic objective is a + per-step scoring rule averaged over those steps, so it does, and + ``--train_steps_to_log`` selects which of them to report as ``train_loss_unroll{i}``. The logged ``train_loss`` is the mean of the per-step losses, i.e. exactly ``forecaster.compute_training_loss``. From ade1929743463583f4c749db4d4ad80b1fa7cd17 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 31 Jul 2026 13:47:17 +0530 Subject: [PATCH 23/41] Address PR review: spell out module constructor args, note sequential sampling Replace the *args/**kwargs passthrough on DeterministicForecasterModule and ProbabilisticForecasterModule with the full explicit signature and parameter docs, so each module's constructor is self-describing rather than pointing at the base class. eval_ensemble_size becomes keyword-only, matching that it is required and has no sensible default. Writing out the signatures exposed a latent fragility in hparam handling: save_hyperparameters() collects the arguments of whichever __init__ frame calls it, which was the base class's only because the subclasses forwarded opaquely. With explicit subclass signatures it instead captured their arguments as passed, i.e. before the base resolves mutable defaults and unpacks a legacy args namespace, leaving hparams.val_steps_to_log as None. Write the resolved values back after saving so hparams no longer depends on which frame happens to make the call. Also document that sample_ensemble draws members sequentially, and that this is an implementation choice rather than a constraint, since members are independent given the inputs. --- .../models/forecasters/probabilistic.py | 6 ++ neural_lam/models/modules/base.py | 25 +++++- neural_lam/models/modules/deterministic.py | 69 ++++++++++++++-- neural_lam/models/modules/probabilistic.py | 82 ++++++++++++++++--- 4 files changed, 160 insertions(+), 22 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 33f5eaff..04ab75d9 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -113,6 +113,12 @@ def sample_ensemble( randomness at every step, and stacks them along a new ensemble dimension after the batch dimension. + This implementation draws the members sequentially, one full rollout + at a time, so cost grows linearly with ``num_members``. Members are + independent given the inputs, so this is only an implementation + choice: it could be batched by folding the member dimension into the + batch dimension, at proportionally higher peak memory. + Parameters ---------- init_states : torch.Tensor diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index c269662a..abae68e8 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -89,9 +89,10 @@ def __init__( val_steps_to_log : list of int, optional Specific rollout steps to log during validation/testing. train_steps_to_log : list of int, optional - Specific rollout steps to log during training. Accepted for CLI - and checkpoint compatibility; ``training_step`` only ever logs - the aggregate ``train_loss`` (see its docstring for why). + Specific predicted steps to log during training. Only has an + effect for a forecaster whose objective decomposes per step; + this class's ``training_step`` logs the aggregate ``train_loss`` + alone (see its docstring for why). metrics_watch : list of str, optional List of metrics to watch and log specifically. var_leads_metrics_watch : dict of {int: list of int}, optional @@ -147,6 +148,24 @@ def __init__( # graph_name, hidden_dim, etc. so the caller can reconstruct the # exact forecaster architecture from the checkpoint alone. self.save_hyperparameters(ignore=["datastore", "forecaster"]) + # save_hyperparameters collects the arguments of the __init__ frame + # it is called from, which for a subclass with its own explicit + # signature is that subclass's frame, holding the values as passed. + # Write back the ones resolved above (mutable defaults, and the args + # namespace unpacking) so hparams never disagrees with what this + # module actually uses. + for name, resolved in ( + ("lr", lr), + ("restore_opt", restore_opt), + ("n_example_pred", n_example_pred), + ("create_gif", create_gif), + ("val_steps_to_log", val_steps_to_log), + ("train_steps_to_log", train_steps_to_log), + ("metrics_watch", metrics_watch), + ("var_leads_metrics_watch", var_leads_metrics_watch), + ): + self.hparams[name] = resolved + self.datastore = datastore self.forecaster = forecaster self.matched_metrics: set = set() diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index f112557b..29942acc 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -12,6 +12,8 @@ # Local from ... import metrics, vis +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore from ..forecasters.deterministic import DeterministicForecaster from .base import BaseForecasterModule @@ -29,21 +31,70 @@ class DeterministicForecasterModule(BaseForecasterModule): # score() is supplied by the deterministic objective mixin forecaster: DeterministicForecaster - def __init__(self, *args, **kwargs): + def __init__( + self, + forecaster: DeterministicForecaster, + config: NeuralLAMConfig, + datastore: BaseDatastore, + lr: float = 1e-3, + restore_opt: bool = False, + n_example_pred: int = 1, + create_gif: bool = False, + val_steps_to_log: list[int] | None = None, + train_steps_to_log: list[int] | None = None, + metrics_watch: list[str] | None = None, + var_leads_metrics_watch: dict[int, list[int]] | None = None, + args=None, + ): """ 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``, ...). + forecaster : DeterministicForecaster + The forecaster to evaluate. Must supply ``score``, i.e. carry + the deterministic objective, since validation and testing score + a single prediction through it. + config : NeuralLAMConfig + Configuration object for the neural LAM model. + datastore : BaseDatastore + Datastore providing grid metadata and data access. + lr : float, default 1e-3 + Learning rate for the optimizer. + restore_opt : bool, default False + Whether to restore optimizer state from checkpoint. + n_example_pred : int, default 1 + Number of example predictions to plot during testing. + create_gif : bool, default False + Whether to create GIFs of example predictions. + val_steps_to_log : list of int, optional + Specific predicted steps to log during validation/testing. + train_steps_to_log : list of int, optional + Specific predicted steps to log during training, reported as + ``train_loss_unroll{i}`` (see ``training_step``). + metrics_watch : list of str, optional + List of metrics to watch and log specifically. + var_leads_metrics_watch : dict of {int: list of int}, optional + Mapping from variable index to a list of predicted steps to log + individually for the configured metrics. + args : argparse.Namespace, optional + Pre-refactor ``ARModel`` checkpoint hyperparameters; see + ``BaseForecasterModule.__init__``. """ - super().__init__(*args, **kwargs) + super().__init__( + forecaster=forecaster, + config=config, + datastore=datastore, + lr=lr, + restore_opt=restore_opt, + n_example_pred=n_example_pred, + create_gif=create_gif, + val_steps_to_log=val_steps_to_log, + train_steps_to_log=train_steps_to_log, + metrics_watch=metrics_watch, + var_leads_metrics_watch=var_leads_metrics_watch, + args=args, + ) self.val_metrics: dict[str, list] = { "mse": [], } diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index bccab449..cd539d75 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -8,6 +8,8 @@ # Local from ... import metrics +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore from ..forecasters.probabilistic import ProbabilisticForecaster from .base import BaseForecasterModule @@ -28,24 +30,84 @@ class ProbabilisticForecasterModule(BaseForecasterModule): # The wrapped forecaster must be able to sample ensemble forecasts forecaster: ProbabilisticForecaster - def __init__(self, *args, eval_ensemble_size: int, **kwargs): + def __init__( + self, + forecaster: ProbabilisticForecaster, + config: NeuralLAMConfig, + datastore: BaseDatastore, + *, + eval_ensemble_size: int, + lr: float = 1e-3, + restore_opt: bool = False, + n_example_pred: int = 1, + create_gif: bool = False, + val_steps_to_log: list[int] | None = None, + train_steps_to_log: list[int] | None = None, + metrics_watch: list[str] | None = None, + var_leads_metrics_watch: dict[int, list[int]] | None = None, + args=None, + ): """ Initialize the module and store the evaluation ensemble size. Parameters ---------- - *args - Positional arguments forwarded to - ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, - ``datastore``, ...). + forecaster : ProbabilisticForecaster + The forecaster to evaluate. Must supply ``sample_ensemble``, + since validation and testing score a sampled ensemble. + config : NeuralLAMConfig + Configuration object for the neural LAM model. + datastore : BaseDatastore + Datastore providing grid metadata and data access. eval_ensemble_size : int Number of ensemble members sampled during validation and - testing. - **kwargs - Keyword arguments forwarded to ``BaseForecasterModule.__init__`` - (``lr``, ...). + testing. Keyword-only and required: how many members to draw is + a choice this module cannot sensibly default. + lr : float, default 1e-3 + Learning rate for the optimizer. + restore_opt : bool, default False + Whether to restore optimizer state from checkpoint. + n_example_pred : int, default 1 + Number of example predictions to plot during testing. Unused + here, since ``test_step`` plots no examples. + create_gif : bool, default False + Whether to create GIFs of example predictions. Unused here, for + the same reason. + val_steps_to_log : list of int, optional + Specific predicted steps to log the ensemble-mean RMSE for, + during both validation and testing. + train_steps_to_log : list of int, optional + Specific predicted steps to log during training. Has no effect + unless the forecaster's objective decomposes per step; see + ``BaseForecasterModule.training_step``. + metrics_watch : list of str, optional + List of metrics to watch and log specifically. + var_leads_metrics_watch : dict of {int: list of int}, optional + Mapping from variable index to a list of predicted steps to log + individually for the configured metrics. + args : argparse.Namespace, optional + Pre-refactor ``ARModel`` checkpoint hyperparameters; see + ``BaseForecasterModule.__init__``. + + Raises + ------ + ValueError + If ``eval_ensemble_size`` is less than 1. """ - super().__init__(*args, **kwargs) + super().__init__( + forecaster=forecaster, + config=config, + datastore=datastore, + lr=lr, + restore_opt=restore_opt, + n_example_pred=n_example_pred, + create_gif=create_gif, + val_steps_to_log=val_steps_to_log, + train_steps_to_log=train_steps_to_log, + metrics_watch=metrics_watch, + var_leads_metrics_watch=var_leads_metrics_watch, + args=args, + ) if eval_ensemble_size < 1: raise ValueError( "eval_ensemble_size must be at least 1, " From 68c55ce93a0a6755c5bbd3a1df9bd9b8f8e40e32 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Mon, 3 Aug 2026 14:16:29 +0530 Subject: [PATCH 24/41] Name saved hyperparameters explicitly instead of inspecting the stack save_hyperparameters() decides what to record by walking the constructor chain and letting the most derived __init__ win, on the assumption that its arguments are the authoritative record of how the object was built. BaseForecasterModule breaks that assumption: it unpacks a legacy args namespace and resolves mutable defaults after receiving them, so the values it runs on are not the values a subclass was called with. While the modules forwarded through *args/**kwargs they had no such variables to find and the base's frame was read instead, which is why this only surfaced once the constructors were written out and hparams.val_steps_to_log arrived as None. Passing a mapping makes save_hyperparameters use it verbatim and skip the inspection, so what is recorded equals what is used regardless of how a subclass writes its signature -- removing the undocumented requirement that subclasses not declare these parameters. This also replaces the write-back added in ade1929, which could not fully work: the snapshot into _hparams_initial is the final statement of save_hyperparameters, so correcting values after it returns reaches only the live hparams. Verified equal before and after on key sets, values and checkpoint contents, with _hparams_initial now consistent where it previously diverged on four parameters (six with a legacy args namespace). Skipping the inspection means a subclass's own hyperparameters are no longer collected, so ProbabilisticForecasterModule records eval_ensemble_size itself; a second call merges rather than replaces. Note this was only ever saved as a side effect of writing out its signature in ade1929, so keeping it is a deliberate choice -- it changes evaluation results and should round-trip through a checkpoint. --- neural_lam/models/modules/base.py | 48 +++++++++++--------- neural_lam/models/modules/probabilistic.py | 4 ++ tests/test_probabilistic_forecaster.py | 53 ++++++++++++++++++++++ 3 files changed, 84 insertions(+), 21 deletions(-) diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index abae68e8..4a3f9c2f 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -142,29 +142,35 @@ def __init__( if var_leads_metrics_watch is None: var_leads_metrics_watch = {} - # datastore and forecaster are excluded from saved hparams and must - # be provided explicitly when calling load_from_checkpoint. Saving - # args makes the checkpoint self-describing: it carries model, + # Hyperparameters are named explicitly rather than left to + # save_hyperparameters' default of inspecting the constructor chain. + # That inspection records the arguments of the most derived + # __init__, i.e. as a subclass received them, whereas the values + # this module runs on are the ones resolved just above (the args + # namespace, then the mutable defaults). Passing them keeps what is + # recorded equal to what is used, whatever a subclass's signature + # looks like. Subclasses add their own with a further + # save_hyperparameters call, which merges into these. + # + # datastore and forecaster are deliberately absent and must be + # provided explicitly when calling load_from_checkpoint. Saving args + # makes the checkpoint self-describing: it carries model, # graph_name, hidden_dim, etc. so the caller can reconstruct the # exact forecaster architecture from the checkpoint alone. - self.save_hyperparameters(ignore=["datastore", "forecaster"]) - # save_hyperparameters collects the arguments of the __init__ frame - # it is called from, which for a subclass with its own explicit - # signature is that subclass's frame, holding the values as passed. - # Write back the ones resolved above (mutable defaults, and the args - # namespace unpacking) so hparams never disagrees with what this - # module actually uses. - for name, resolved in ( - ("lr", lr), - ("restore_opt", restore_opt), - ("n_example_pred", n_example_pred), - ("create_gif", create_gif), - ("val_steps_to_log", val_steps_to_log), - ("train_steps_to_log", train_steps_to_log), - ("metrics_watch", metrics_watch), - ("var_leads_metrics_watch", var_leads_metrics_watch), - ): - self.hparams[name] = resolved + self.save_hyperparameters( + { + "config": config, + "lr": lr, + "restore_opt": restore_opt, + "n_example_pred": n_example_pred, + "create_gif": create_gif, + "val_steps_to_log": val_steps_to_log, + "train_steps_to_log": train_steps_to_log, + "metrics_watch": metrics_watch, + "var_leads_metrics_watch": var_leads_metrics_watch, + "args": args, + } + ) self.datastore = datastore self.forecaster = forecaster diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index cd539d75..345282fb 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -113,6 +113,10 @@ def __init__( "eval_ensemble_size must be at least 1, " f"got {eval_ensemble_size}" ) + # The base class names the hyperparameters it knows about; this one + # is ours, so record it here. A second call merges rather than + # replaces, so both end up saved. + self.save_hyperparameters({"eval_ensemble_size": eval_ensemble_size}) self.eval_ensemble_size = eval_ensemble_size self.val_metrics: dict[str, list] = {"ens_mse": []} self.test_metrics: dict[str, list] = {"ens_mse": []} diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 0ca1a2c1..7142121c 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -243,6 +243,59 @@ def test_probabilistic_ar_forecaster_is_abstract(): ProbabilisticARForecaster(predictor, datastore) +def test_saved_hparams_hold_resolved_values(): + """Saved hyperparameters are the values the module runs on. + + ``save_hyperparameters`` defaults to inspecting the constructor chain + and recording the arguments of the most derived ``__init__``, which for + a subclass with its own signature are the values as passed, before + ``BaseForecasterModule`` resolves them. ``hparams_initial`` is asserted + alongside ``hparams`` because it is snapshotted inside + ``save_hyperparameters``: any attempt to correct the values after that + call returns reaches only ``hparams`` and would fail here. + """ + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + module = DeterministicForecasterModule( + forecaster=DeterministicARForecaster( + predictor, datastore, config=config + ), + config=config, + datastore=datastore, + ) + + # Defaults resolved by the base class, not the None it was called with + assert module.hparams.val_steps_to_log == [1] + assert module.hparams.train_steps_to_log == [] + assert module.hparams.metrics_watch == [] + assert module.hparams.var_leads_metrics_watch == {} + assert dict(module.hparams_initial) == dict(module.hparams) + + # The forecaster and datastore are supplied on load, never saved + assert "forecaster" not in module.hparams + assert "datastore" not in module.hparams + + # A subclass's own hyperparameters survive alongside the base's + prob_module = ProbabilisticForecasterModule( + forecaster=ConcreteProbabilisticARForecaster( + NoisyStepPredictor(datastore=datastore, output_std=False), + datastore, + config=config, + ), + config=config, + datastore=datastore, + eval_ensemble_size=3, + ) + assert prob_module.hparams.eval_ensemble_size == 3 + assert prob_module.hparams.val_steps_to_log == [1] + assert dict(prob_module.hparams_initial) == dict(prob_module.hparams) + + def test_forecaster_module_evaluation_steps_are_abstract(): """A module omitting validation_step/test_step fails at construction. From 195e8f24ff5670e24a98c8f22027882a9f017521 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Tue, 4 Aug 2026 17:51:58 +0530 Subject: [PATCH 25/41] Address PR review: make forecaster constructors cooperate via kwargs DeterministicForecaster required subclasses to call _configure_scoring to set up the scoring rule and buffers, which are its own concern. Give it a real __init__ instead and have every forecaster mix-in consume the keyword arguments it owns and forward the rest along the MRO. datastore is taken by Forecaster, since both mix-ins need it. Arguments no mix-in claims reach nn.Module and raise rather than being silently dropped. Also drop the one cross-mix-in read at construction time: the per_var_std fallback was allocated based on predicts_std, which the mix-in supplying forward answers, so the two only composed with the objective listed first. Register it whenever a config is given and let _resolve_pred_std settle per call whether it is used, so the mix-ins compose in either order. --- .../models/forecasters/autoregressive.py | 12 ++- neural_lam/models/forecasters/base.py | 41 +++++++++ .../models/forecasters/deterministic.py | 71 ++++++++------- tests/test_prediction_model_classes.py | 89 +++++++++++++++++++ 4 files changed, 178 insertions(+), 35 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 671c1400..b075f5db 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 Any + # Third-party import torch @@ -27,6 +30,7 @@ def __init__( self, predictor: StepPredictor, datastore: BaseDatastore, + **kwargs: Any, ) -> None: """ Initialize the ARForecaster. @@ -36,9 +40,13 @@ def __init__( predictor : StepPredictor The predictor to use for each step. datastore : BaseDatastore - The datastore providing grid metadata and boundary masks. + The datastore providing grid metadata and boundary masks. Also + forwarded on, since mix-ins later in the MRO need it too. + **kwargs : Any + Arguments belonging to the mix-ins this is combined with, + forwarded unchanged along the MRO. See ``Forecaster``. """ - super().__init__() + super().__init__(datastore=datastore, **kwargs) self.predictor = predictor # Register boundary/interior masks on the forecaster, not the predictor diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index ccd8ff6e..b48dfa6c 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,19 +2,60 @@ # Standard library from abc import ABC, abstractmethod +from typing import Any # Third-party import torch from torch import nn +# Local +from ...datastore import BaseDatastore + class Forecaster(nn.Module, ABC): """ Generic forecaster capable of mapping from a set of initial states, forcing and forces and previous states into a full forecast of the requested length. + + Concrete forecasters are assembled from mix-ins that each configure one + aspect: how forecasts are produced (``ARForecaster``) and which training + objective is used (``DeterministicForecaster``). Their constructors + cooperate along the MRO: each consumes the keyword arguments it needs + and forwards the rest with ``super().__init__(**kwargs)``, so a mix-in + never has to know which others it is combined with. + + Because each mix-in forwards *before* running its own setup, the bodies + run in reverse MRO order, so a mix-in reading another's attributes would + only work for one ordering of the bases. No mix-in does, and none + should: keeping their constructors independent is what lets them be + combined in any order. """ + def __init__(self, datastore: BaseDatastore, **kwargs: Any) -> None: + """ + Initialize the forecaster and end the cooperative ``__init__`` chain. + + Runs before every mix-in body, since each forwards to + ``super().__init__`` ahead of its own setup. + + ``datastore`` is taken here rather than by the mix-in that happens to + need it, because several of them do; each forwards it on so that the + ones after it in the MRO still receive it. + + Parameters + ---------- + datastore : BaseDatastore + The datastore this forecaster is built for, providing grid + metadata, boundary masks and standardization statistics. + **kwargs : Any + Forwarded to ``nn.Module``, which raises ``TypeError`` on + anything left, so an argument no mix-in claimed is not silently + dropped. + """ + super().__init__(**kwargs) + self.datastore = datastore + @property @abstractmethod def predicts_std(self) -> bool: diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 5de1b385..13bca699 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -1,7 +1,7 @@ """Forecasters trained by scoring a single deterministic forecast.""" # Standard library -from typing import Callable, Optional +from typing import Any, Callable, Optional # Third-party import torch @@ -27,50 +27,52 @@ class DeterministicForecaster(Forecaster): about *how* that forecast is produced, so this composes with any way of implementing ``forward`` (see ``DeterministicARForecaster`` for the auto-regressive combination). - - Concrete subclasses must call ``_configure_scoring`` from their - ``__init__`` to set up the scoring rule and the ``pred_std`` fallback. """ - def _configure_scoring( + def __init__( self, datastore: BaseDatastore, - config: NeuralLAMConfig | None, - loss: str, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + **kwargs: Any, ) -> None: """ Set up the scoring rule and the constant ``pred_std`` fallback. - Called by concrete subclasses from ``__init__``, after - ``torch.nn.Module`` initialization (buffers are registered here). - Parameters ---------- datastore : BaseDatastore The datastore providing the state standardization statistics - used to compute ``per_var_std``. - config : NeuralLAMConfig or None + used to compute ``per_var_std``. Also forwarded on, since + mix-ins later in the MRO need it too. + config : NeuralLAMConfig or None, optional Configuration used to compute the constant per-variable std substituted for ``pred_std`` when the forecast carries no std of its own. Required in that case for ``score`` and ``compute_training_loss`` to work (they raise ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters used purely for - inference can omit it. - loss : str + inference can omit it. Default ``None``. + loss : str, optional The scoring rule (from ``neural_lam.metrics``) applied by - ``compute_training_loss``, stored as ``self.loss``. + ``compute_training_loss``, stored as ``self.loss``. Default + ``"wmse"``. + **kwargs : Any + Arguments belonging to the mix-ins this is combined with, + forwarded unchanged along the MRO. See ``Forecaster``. """ + super().__init__(datastore=datastore, **kwargs) self.loss = metrics.get_metric(loss) - # Store per_var_std only if the forecast carries no std of its own - if not self.predicts_std and config is not None: - self.register_buffer( - "per_var_std", - get_per_var_std(config=config, datastore=datastore), - persistent=False, - ) - else: - self.per_var_std = None + # Registered whenever a config is given, rather than only when the + # forecast lacks its own std: that is settled per call in + # _resolve_pred_std, and checking it here would mean reading state a + # mix-in later in the MRO has yet to set up + per_var_std = ( + get_per_var_std(config=config, datastore=datastore) + if config is not None + else None + ) + self.register_buffer("per_var_std", per_var_std, persistent=False) def compute_step_losses( self, @@ -216,16 +218,15 @@ def _resolve_pred_std( ------ ValueError If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is - available (``predicts_std`` is False and this forecaster was - constructed without ``config``). + available (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 " + "No pred_std available for scoring: the forecast carries no " + "std 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 @@ -332,7 +333,11 @@ def __init__( The scoring rule (from ``neural_lam.metrics``) applied by ``compute_training_loss``. """ - # DeterministicForecaster defines no __init__, so this initializes - # the AR half (and torch.nn.Module) before any buffer is registered - super().__init__(predictor, datastore) - self._configure_scoring(datastore, config=config, loss=loss) + # Named rather than positional: each argument is consumed by + # whichever half of the MRO declares it + super().__init__( + predictor=predictor, + datastore=datastore, + config=config, + loss=loss, + ) diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 0cfbdcca..f6a4607b 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -10,8 +10,11 @@ from neural_lam import config as nlconfig from neural_lam import metrics from neural_lam.models import ( + ARForecaster, DeterministicARForecaster, + DeterministicForecaster, DeterministicForecasterModule, + Forecaster, StepPredictor, ) from tests.conftest import init_datastore_example @@ -165,6 +168,92 @@ def test_ar_forecaster_without_config_raises_on_use_not_construction(): ) +class MeanAbsObjective(Forecaster): + """Test-only objective mix-in taking a constructor argument of its own. + + Stands in for a future objective (CRPS, a variational loss, ...) to + check that adding one requires nothing beyond declaring its arguments + and forwarding the rest. + """ + + def __init__(self, datastore, scale: float = 1.0, **kwargs): + super().__init__(datastore=datastore, **kwargs) + self.scale = scale + + def compute_training_loss( + self, init_states, forcing_features, target_states, interior_mask_bool + ): + prediction, _ = self(init_states, forcing_features, target_states) + loss = self.scale * torch.mean(torch.abs(prediction - target_states)) + return loss, {} + + +class MeanAbsARForecaster(MeanAbsObjective, ARForecaster): + """Auto-regressive forecast production plus the mean-abs objective.""" + + +def test_objective_mixin_composes_without_manual_wiring(): + """A new objective mix-in only declares its own constructor arguments + and forwards the rest; combining it with ARForecaster initializes both + halves, with no setup call for the concrete class to remember.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + forecaster = MeanAbsARForecaster( + predictor=predictor, datastore=datastore, scale=2.0 + ) + + # Objective half + assert forecaster.scale == 2.0 + # Forecast-production half, initialized before the objective's body ran + assert forecaster.predictor is predictor + assert forecaster.boundary_mask.shape[1] == predictor.num_grid_nodes + # Shared argument, consumed at the end of the chain + assert forecaster.datastore is datastore + + +def test_objective_mixin_composes_in_either_order(): + """Mix-in constructors must not read state another mix-in sets up. + DeterministicForecaster used to decide the per_var_std fallback from + predicts_std, which ARForecaster answers from the predictor it stores, + so it only worked when the objective was listed first.""" + 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) + + class ObjectiveFirst(DeterministicForecaster, ARForecaster): + pass + + class ProductionFirst(ARForecaster, DeterministicForecaster): + pass + + for cls in (ObjectiveFirst, ProductionFirst): + forecaster = cls( + predictor=predictor, + datastore=datastore, + config=config, + loss="mse", + ) + assert forecaster.predictor is predictor + assert forecaster.per_var_std is not None + + +def test_unclaimed_constructor_argument_raises(): + """An argument no mix-in in the MRO declares must not be swallowed by + the **kwargs forwarding.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(TypeError, match="not_a_real_arg"): + MeanAbsARForecaster( + predictor=predictor, datastore=datastore, not_a_real_arg=1 + ) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") From 1ad3881c7cca4b21c07b7b10a43307c2bf26931e Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Tue, 4 Aug 2026 18:17:05 +0530 Subject: [PATCH 26/41] Address PR review: compute reporting metrics in the module, not via score Drop the metric argument from Forecaster.score, leaving it to apply the forecaster's own scoring rule. Compute mse and mae in DeterministicForecasterModule directly from neural_lam.metrics, as ProbabilisticForecasterModule already does. Values are unchanged: both metrics replace the std argument with ones internally. --- .../models/forecasters/deterministic.py | 30 ++++++++-------- neural_lam/models/modules/deterministic.py | 36 +++++++++++++------ tests/test_prediction_model_classes.py | 11 ------ 3 files changed, 41 insertions(+), 36 deletions(-) diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 13bca699..d0e7273a 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -1,7 +1,7 @@ """Forecasters trained by scoring a single deterministic forecast.""" # Standard library -from typing import Any, Callable, Optional +from typing import Any, Optional # Third-party import torch @@ -236,17 +236,22 @@ def score( 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). + Apply this forecaster's scoring rule to an already-produced forecast. 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). + ``self.per_var_std`` when ``None``), then applies ``self.loss``. Used + to report the loss on a forecast the caller already has, at a + reduction of its choosing, without recomputing it. + + Only the loss goes through here. Metrics unrelated to the training + objective depend on nothing but the shapes of the three tensors + below, so callers compute those directly from + ``neural_lam.metrics``. Parameters ---------- @@ -262,21 +267,17 @@ def score( the forecast carries no std, in which case ``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, - sum_vars=True) -> torch.Tensor``. Defaults to ``self.loss``. mask : torch.Tensor or None, optional - Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``self.loss``. average_grid : bool, optional - Forwarded to ``metric``. + Forwarded to ``self.loss``. Default ``True``. sum_vars : bool, optional - Forwarded to ``metric``. + Forwarded to ``self.loss``. Default ``True``. Returns ------- torch.Tensor - The metric's output; shape depends on ``average_grid`` and + The scoring rule's output; shape depends on ``average_grid`` and ``sum_vars`` (see ``neural_lam.metrics``). Raises @@ -286,8 +287,7 @@ def score( available; see ``_resolve_pred_std``. """ pred_std = self._resolve_pred_std(pred_std) - metric_fn = self.loss if metric is None else metric - return metric_fn( + return self.loss( prediction, target_states, pred_std, diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 29942acc..7d4b978f 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -22,10 +22,16 @@ class DeterministicForecasterModule(BaseForecasterModule): """ Lightning module for a single deterministic forecast per batch. - Validation and testing score the forecaster's own single 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``). + Validation and testing evaluate the forecaster's own single prediction, + as opposed to ``ProbabilisticForecasterModule``, which samples and scores + an ensemble. Training is shared with that module unchanged (see + ``BaseForecasterModule.training_step``). + + The reported loss comes from ``forecaster.score``, since only the + forecaster knows its objective. The reported metrics (mse, mae) are + computed here from ``neural_lam.metrics``: they are fixed regardless of + what the forecaster trains on, so routing them through it would add a + layer without adding meaning. """ # score() is supplied by the deterministic objective mixin @@ -210,11 +216,16 @@ def validation_step(self, batch, batch_idx): batch_size = batch[0].shape[0] self._log_step_loss(time_step_loss, mean_loss, "val", batch_size) - entry_mses = self.forecaster.score( + # Reported independently of the training objective, so computed here + # rather than through the forecaster. metrics.mse ignores the std + # argument, but requires one + std_placeholder = torch.ones( + target_states.shape[-1], device=target_states.device + ) + entry_mses = metrics.mse( prediction, target_states, - pred_std, - metric=metrics.mse, + std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) @@ -246,12 +257,17 @@ def test_step(self, batch, batch_idx): batch_size = batch[0].shape[0] self._log_step_loss(time_step_loss, mean_loss, "test", batch_size) + # Reported independently of the training objective, so computed here + # rather than through the forecaster. Both ignore the std argument, + # but require one + std_placeholder = torch.ones( + target_states.shape[-1], device=target_states.device + ) for metric_name in ("mse", "mae"): - batch_metric_vals = self.forecaster.score( + batch_metric_vals = metrics.get_metric(metric_name)( prediction, target_states, - pred_std, - metric=metrics.get_metric(metric_name), + std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index f6a4607b..e9fe9334 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -8,7 +8,6 @@ # First-party from neural_lam import config as nlconfig -from neural_lam import metrics from neural_lam.models import ( ARForecaster, DeterministicARForecaster, @@ -106,16 +105,6 @@ def test_ar_forecaster_score(): ) 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( From 3ca990dffdc5bb95c4dd8427fb129f633470556d Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Tue, 4 Aug 2026 19:20:55 +0530 Subject: [PATCH 27/41] Address PR review: share _compute_prediction_and_loss with training_step The per-step losses it returns are the ones compute_step_losses produces for this objective, gradients included, so all three steps can use it. --- neural_lam/models/modules/deterministic.py | 24 +++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 7d4b978f..0b8db197 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -127,6 +127,11 @@ def training_step(self, batch): the per-step losses, i.e. exactly ``forecaster.compute_training_loss``. + Shares ``_compute_prediction_and_loss`` with the validation and test + steps: for this objective the per-step losses it returns are the same + ones ``forecaster.compute_step_losses`` produces, down to the + gradients. + Parameters ---------- batch : tuple @@ -137,13 +142,7 @@ def training_step(self, batch): torch.Tensor The computed loss for the training step. """ - init_states, target_states, forcing_features, _ = batch - time_step_loss = self.forecaster.compute_step_losses( - init_states, - forcing_features, - target_states, - interior_mask_bool=self.interior_mask_bool, - ) + _, _, _, time_step_loss = self._compute_prediction_and_loss(batch) batch_loss = torch.mean(time_step_loss) self._log_step_loss( time_step_loss, batch_loss, "train", batch[0].shape[0] @@ -158,11 +157,12 @@ def _compute_prediction_and_loss( Compute predicted mean, standard deviation, and step-wise loss. Also extract and return corresponding target from batch. - Shared by ``validation_step`` and ``test_step``. ``training_step`` - deliberately does not use this: it goes through the forecaster's - training objective rather than the reporting scoring rule applied - here, and those are only interchangeable for this particular - objective. + Shared by ``training_step``, ``validation_step`` and ``test_step``. + The scoring rule applied here is the forecaster's own, so for this + objective the per-step losses match + ``forecaster.compute_step_losses`` exactly; a forecaster whose + objective is not a per-step scoring rule needs the plain + ``BaseForecasterModule.training_step`` instead. Parameters ---------- From ff5c511ff88ffc8b8df070c3af07d82fd4cd3044 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Tue, 4 Aug 2026 19:28:34 +0530 Subject: [PATCH 28/41] Address PR review: drop the unenforced metrics dict declarations The class-level annotations declared a contract nothing checked. Subclasses already create both dicts, and the docstrings on validation_step and test_step say so. --- neural_lam/models/modules/base.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index 4a3f9c2f..c9043393 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -40,12 +40,6 @@ class BaseForecasterModule(pl.LightningModule, ABC): ``ProbabilisticForecasterModule``) rather than overriding one another. """ - # Which metrics are collected differs per evaluation mode, so concrete - # subclasses create these in __init__; the epoch-end hooks here consume - # whatever they contain. Declared for the contract only, not assigned. - val_metrics: dict[str, list] - test_metrics: dict[str, list] - # pylint: disable=arguments-differ def __init__( From a6586a2cd5a320688b50b2d8055713fffde321ab Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Tue, 4 Aug 2026 22:58:09 +0530 Subject: [PATCH 29/41] Address PR review: gate ensemble RMSE logging with rank_zero_only Split the reduction and logging into their own method so the decorator can gate it. The all_gather stays in the caller, which every rank must reach. --- neural_lam/models/modules/probabilistic.py | 35 +++++++++++++++------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 345282fb..cd079df1 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -5,6 +5,7 @@ # Third-party import torch +from pytorch_lightning.utilities import rank_zero_only # Local from ... import metrics @@ -177,12 +178,9 @@ def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: """ Log the ensemble-mean RMSE accumulated over a full epoch. - Averages the per-variable MSEs collected by ``_ensemble_step`` over - every sample of the epoch (across both batches and devices) and sums - them over variables, and only then takes the square root. Rooting - each batch's MSE and averaging those roots instead would report a - different quantity, since the square root does not commute with the - averaging. + Gathers across devices, which every rank must take part in, then + hands the result to ``_log_gathered_ensemble_rmse`` for the rank-zero + half. Parameters ---------- @@ -193,14 +191,31 @@ def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: phase : str Logging phase, either ``"val"`` or ``"test"``. """ - # Collective: must be reached by every rank, so it precedes the - # rank-zero check below. entry_mses = self.all_gather_cat(torch.cat(entry_mse_list, dim=0)) # (total_samples, pred_steps, num_state_vars) + self._log_gathered_ensemble_rmse(entry_mses, phase) - if not self.trainer.is_global_zero: - return + @rank_zero_only + def _log_gathered_ensemble_rmse( + self, entry_mses: torch.Tensor, phase: str + ) -> None: + """ + Reduce the gathered MSEs to an RMSE and log it, on rank zero only. + + Averages the per-variable MSEs over every sample of the epoch (across + both batches and devices) and sums them over variables, and only then + takes the square root. Rooting each batch's MSE and averaging those + roots instead would report a different quantity, since the square + root does not commute with the averaging. + Parameters + ---------- + entry_mses : torch.Tensor + Shape ``(total_samples, pred_steps, num_state_vars)``. + Per-variable ensemble-mean MSEs gathered over the whole epoch. + phase : str + Logging phase, either ``"val"`` or ``"test"``. + """ time_step_rmse = torch.sqrt(entry_mses.sum(dim=-1).mean(dim=0)) # (pred_steps,) From 21be5ed7a50cbabce8aae13f90b8abd48a8e428d Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 7 Aug 2026 12:01:15 +0530 Subject: [PATCH 30/41] Address PR review: keep ARForecaster as the only constructor mix-in DeterministicForecaster takes a fixed signature again; ARForecaster alone forwards **kwargs and is listed first in the bases. Trim the docstrings and comments that narrated the mix-in mechanics. --- .../models/forecasters/autoregressive.py | 19 +++---- neural_lam/models/forecasters/base.py | 33 ++--------- .../models/forecasters/deterministic.py | 31 ++++------ neural_lam/models/modules/deterministic.py | 2 +- tests/test_prediction_model_classes.py | 57 ++++--------------- 5 files changed, 36 insertions(+), 106 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index b075f5db..2aaad08a 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -17,13 +17,11 @@ class ARForecaster(Forecaster): Forecaster that produces a forecast by auto-regressive unrolling, using a StepPredictor at each AR step. - This class fixes only *how forecasts are produced*, and deliberately says - nothing about how a training objective is computed from them: it leaves - ``compute_training_loss`` abstract. The two are orthogonal, so the - objective is mixed in separately (see ``DeterministicForecaster`` and - ``ProbabilisticForecaster``), which lets an auto-regressive forecaster be - trained deterministically or probabilistically, and equally lets a - non-auto-regressive forecaster reuse either objective. + This class fixes only *how forecasts are produced* and leaves the + training objective (``compute_training_loss``) abstract. It is a + mix-in: combine it with an objective class, listing it first so its + ``**kwargs`` forwarding reaches the objective's constructor, as in + ``DeterministicARForecaster(ARForecaster, DeterministicForecaster)``. """ def __init__( @@ -40,11 +38,10 @@ def __init__( predictor : StepPredictor The predictor to use for each step. datastore : BaseDatastore - The datastore providing grid metadata and boundary masks. Also - forwarded on, since mix-ins later in the MRO need it too. + The datastore providing grid metadata and boundary masks. **kwargs : Any - Arguments belonging to the mix-ins this is combined with, - forwarded unchanged along the MRO. See ``Forecaster``. + Constructor arguments of the objective class this mix-in is + combined with, forwarded to it unchanged. """ super().__init__(datastore=datastore, **kwargs) self.predictor = predictor diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index b48dfa6c..30f3e46b 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 Any # Third-party import torch @@ -18,42 +17,22 @@ class Forecaster(nn.Module, ABC): forcing and forces and previous states into a full forecast of the requested length. - Concrete forecasters are assembled from mix-ins that each configure one - aspect: how forecasts are produced (``ARForecaster``) and which training - objective is used (``DeterministicForecaster``). Their constructors - cooperate along the MRO: each consumes the keyword arguments it needs - and forwards the rest with ``super().__init__(**kwargs)``, so a mix-in - never has to know which others it is combined with. - - Because each mix-in forwards *before* running its own setup, the bodies - run in reverse MRO order, so a mix-in reading another's attributes would - only work for one ordering of the bases. No mix-in does, and none - should: keeping their constructors independent is what lets them be - combined in any order. + Concrete forecasters combine a way of producing forecasts (the + ``ARForecaster`` mix-in) with a training objective + (``DeterministicForecaster``). """ - def __init__(self, datastore: BaseDatastore, **kwargs: Any) -> None: + def __init__(self, datastore: BaseDatastore) -> None: """ - Initialize the forecaster and end the cooperative ``__init__`` chain. - - Runs before every mix-in body, since each forwards to - ``super().__init__`` ahead of its own setup. - - ``datastore`` is taken here rather than by the mix-in that happens to - need it, because several of them do; each forwards it on so that the - ones after it in the MRO still receive it. + Initialize the forecaster. Parameters ---------- datastore : BaseDatastore The datastore this forecaster is built for, providing grid metadata, boundary masks and standardization statistics. - **kwargs : Any - Forwarded to ``nn.Module``, which raises ``TypeError`` on - anything left, so an argument no mix-in claimed is not silently - dropped. """ - super().__init__(**kwargs) + super().__init__() self.datastore = datastore @property diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index d0e7273a..8f4a289b 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -1,7 +1,7 @@ """Forecasters trained by scoring a single deterministic forecast.""" # Standard library -from typing import Any, Optional +from typing import Optional # Third-party import torch @@ -21,12 +21,10 @@ class DeterministicForecaster(Forecaster): Forecaster whose training objective is a scoring rule applied to a single forecast. - Supplies the objective half of a forecaster: ``compute_training_loss`` - produces one forecast and scores it, and ``score`` applies a metric to - an already-produced forecast for reporting. Neither makes any assumption - about *how* that forecast is produced, so this composes with any way of - implementing ``forward`` (see ``DeterministicARForecaster`` for the - auto-regressive combination). + ``compute_training_loss`` produces one forecast and scores it, and + ``score`` applies the same metric to an already-produced forecast for + reporting. ``forward`` is left abstract; see + ``DeterministicARForecaster`` for the auto-regressive combination. """ def __init__( @@ -34,7 +32,6 @@ def __init__( datastore: BaseDatastore, config: NeuralLAMConfig | None = None, loss: str = "wmse", - **kwargs: Any, ) -> None: """ Set up the scoring rule and the constant ``pred_std`` fallback. @@ -43,8 +40,7 @@ def __init__( ---------- datastore : BaseDatastore The datastore providing the state standardization statistics - used to compute ``per_var_std``. Also forwarded on, since - mix-ins later in the MRO need it too. + used to compute ``per_var_std``. config : NeuralLAMConfig or None, optional Configuration used to compute the constant per-variable std substituted for ``pred_std`` when the forecast carries no std of @@ -56,17 +52,12 @@ def __init__( The scoring rule (from ``neural_lam.metrics``) applied by ``compute_training_loss``, stored as ``self.loss``. Default ``"wmse"``. - **kwargs : Any - Arguments belonging to the mix-ins this is combined with, - forwarded unchanged along the MRO. See ``Forecaster``. """ - super().__init__(datastore=datastore, **kwargs) + super().__init__(datastore=datastore) self.loss = metrics.get_metric(loss) - # Registered whenever a config is given, rather than only when the - # forecast lacks its own std: that is settled per call in - # _resolve_pred_std, and checking it here would mean reading state a - # mix-in later in the MRO has yet to set up + # Whether the fallback is needed is settled per call in + # _resolve_pred_std per_var_std = ( get_per_var_std(config=config, datastore=datastore) if config is not None @@ -297,7 +288,7 @@ def score( ) -class DeterministicARForecaster(DeterministicForecaster, ARForecaster): +class DeterministicARForecaster(ARForecaster, DeterministicForecaster): """ Auto-regressive forecaster trained by scoring its single rollout. @@ -333,8 +324,6 @@ def __init__( The scoring rule (from ``neural_lam.metrics``) applied by ``compute_training_loss``. """ - # Named rather than positional: each argument is consumed by - # whichever half of the MRO declares it super().__init__( predictor=predictor, datastore=datastore, diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 0b8db197..60e6262c 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -34,7 +34,7 @@ class DeterministicForecasterModule(BaseForecasterModule): layer without adding meaning. """ - # score() is supplied by the deterministic objective mixin + # Narrowed from Forecaster: this module calls forecaster.score() forecaster: DeterministicForecaster def __init__( diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index e9fe9334..bcadb1e9 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -11,7 +11,6 @@ from neural_lam.models import ( ARForecaster, DeterministicARForecaster, - DeterministicForecaster, DeterministicForecasterModule, Forecaster, StepPredictor, @@ -158,15 +157,14 @@ def test_ar_forecaster_without_config_raises_on_use_not_construction(): class MeanAbsObjective(Forecaster): - """Test-only objective mix-in taking a constructor argument of its own. + """Test-only objective class taking a constructor argument of its own. Stands in for a future objective (CRPS, a variational loss, ...) to - check that adding one requires nothing beyond declaring its arguments - and forwarding the rest. + check that adding one requires nothing beyond declaring its arguments. """ - def __init__(self, datastore, scale: float = 1.0, **kwargs): - super().__init__(datastore=datastore, **kwargs) + def __init__(self, datastore, scale: float = 1.0): + super().__init__(datastore=datastore) self.scale = scale def compute_training_loss( @@ -177,14 +175,14 @@ def compute_training_loss( return loss, {} -class MeanAbsARForecaster(MeanAbsObjective, ARForecaster): +class MeanAbsARForecaster(ARForecaster, MeanAbsObjective): """Auto-regressive forecast production plus the mean-abs objective.""" -def test_objective_mixin_composes_without_manual_wiring(): - """A new objective mix-in only declares its own constructor arguments - and forwards the rest; combining it with ARForecaster initializes both - halves, with no setup call for the concrete class to remember.""" +def test_objective_class_composes_without_manual_wiring(): + """An objective class only declares its own constructor arguments; the + ARForecaster mix-in forwards the rest to it, so combining the two + initializes both halves with no setup call to remember.""" datastore = init_datastore_example("mdp") predictor = MockStepPredictor(datastore=datastore, output_std=False) @@ -192,48 +190,15 @@ def test_objective_mixin_composes_without_manual_wiring(): predictor=predictor, datastore=datastore, scale=2.0 ) - # Objective half assert forecaster.scale == 2.0 - # Forecast-production half, initialized before the objective's body ran assert forecaster.predictor is predictor assert forecaster.boundary_mask.shape[1] == predictor.num_grid_nodes - # Shared argument, consumed at the end of the chain assert forecaster.datastore is datastore -def test_objective_mixin_composes_in_either_order(): - """Mix-in constructors must not read state another mix-in sets up. - DeterministicForecaster used to decide the per_var_std fallback from - predicts_std, which ARForecaster answers from the predictor it stores, - so it only worked when the objective was listed first.""" - 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) - - class ObjectiveFirst(DeterministicForecaster, ARForecaster): - pass - - class ProductionFirst(ARForecaster, DeterministicForecaster): - pass - - for cls in (ObjectiveFirst, ProductionFirst): - forecaster = cls( - predictor=predictor, - datastore=datastore, - config=config, - loss="mse", - ) - assert forecaster.predictor is predictor - assert forecaster.per_var_std is not None - - def test_unclaimed_constructor_argument_raises(): - """An argument no mix-in in the MRO declares must not be swallowed by - the **kwargs forwarding.""" + """An argument neither ARForecaster nor the objective class declares + must raise, not be silently swallowed by the **kwargs forwarding.""" datastore = init_datastore_example("mdp") predictor = MockStepPredictor(datastore=datastore, output_std=False) From 80a4b0e95f97b5266bd0c13774b24b5b231e5aa1 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 7 Aug 2026 12:15:11 +0530 Subject: [PATCH 31/41] Address PR review: fold compute_step_losses into compute_training_loss The deterministic module gets its per-step breakdown from score() since 3ca990d, leaving compute_step_losses with no caller outside the class. Forecaster now exposes two loss entry points: compute_training_loss, which produces its own forecast and returns the finished objective, and score, which applies the same rule to a forecast the caller already has. --- .../models/forecasters/deterministic.py | 75 +++---------------- neural_lam/models/modules/deterministic.py | 20 ++--- tests/test_probabilistic_forecaster.py | 16 ++-- 3 files changed, 26 insertions(+), 85 deletions(-) diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 8f4a289b..711c5932 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -65,66 +65,6 @@ def __init__( ) self.register_buffer("per_var_std", per_var_std, persistent=False) - def compute_step_losses( - self, - init_states: torch.Tensor, - forcing_features: torch.Tensor, - target_states: torch.Tensor, - interior_mask_bool: torch.Tensor, - ) -> torch.Tensor: - """ - Score a single forecast with ``self.loss``, per predicted step. - - This objective is a per-step scoring rule averaged over the - predicted steps, so it decomposes into the contribution of each of - them, which callers can report individually. - ``compute_training_loss`` is the mean of what this returns. - - 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. - forcing_features : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. - External forcings provided at each predicted step. - 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 while forecasting. - 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. - - Returns - ------- - torch.Tensor - Shape ``(pred_steps,)``. The scoring rule at each predicted - step, averaged over the batch. - - Raises - ------ - ValueError - If the forecast carries no std of its own and no - ``per_var_std`` fallback is available; see - ``_resolve_pred_std``. - """ - prediction, pred_std = self( - init_states, forcing_features, target_states - ) - pred_std = self._resolve_pred_std(pred_std) - - return torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=interior_mask_bool, - ), - dim=0, - ) - def compute_training_loss( self, init_states: torch.Tensor, @@ -137,8 +77,8 @@ def compute_training_loss( Produces one forecast over every predicted step, scores it against the target states on interior nodes and averages over batch and - time. Callers wanting the per-step breakdown of this same objective - should use ``compute_step_losses`` instead, which this averages. + time. Callers that already hold a forecast should score that one + with ``score`` rather than producing another here. Parameters ---------- @@ -180,11 +120,14 @@ def compute_training_loss( ``per_var_std`` fallback is available; see ``_resolve_pred_std``. """ - step_losses = self.compute_step_losses( - init_states, - forcing_features, + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + step_losses = self.score( + prediction, target_states, - interior_mask_bool, + pred_std, + mask=interior_mask_bool, ) return torch.mean(step_losses), {} diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 60e6262c..5480c86b 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -120,17 +120,10 @@ def training_step(self, batch): Overrides ``BaseForecasterModule.training_step``, which logs only the scalar objective because a forecaster's training loss need not - decompose over predicted steps. The deterministic objective is a - per-step scoring rule averaged over those steps, so it does, and - ``--train_steps_to_log`` selects which of them to report as - ``train_loss_unroll{i}``. The logged ``train_loss`` is the mean of - the per-step losses, i.e. exactly - ``forecaster.compute_training_loss``. - - Shares ``_compute_prediction_and_loss`` with the validation and test - steps: for this objective the per-step losses it returns are the same - ones ``forecaster.compute_step_losses`` produces, down to the - gradients. + decompose over predicted steps. The deterministic objective does, so + ``--train_steps_to_log`` can report individual steps as + ``train_loss_unroll{i}`` and the logged ``train_loss`` is their mean, + equal to ``forecaster.compute_training_loss``. Parameters ---------- @@ -158,9 +151,8 @@ def _compute_prediction_and_loss( Also extract and return corresponding target from batch. Shared by ``training_step``, ``validation_step`` and ``test_step``. - The scoring rule applied here is the forecaster's own, so for this - objective the per-step losses match - ``forecaster.compute_step_losses`` exactly; a forecaster whose + The scoring rule applied here is the forecaster's own, so the mean + of the per-step losses is its training objective; a forecaster whose objective is not a per-step scoring rule needs the plain ``BaseForecasterModule.training_step`` instead. diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 7142121c..9c48fe00 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -392,11 +392,17 @@ def test_deterministic_training_step_logs_per_step_losses(): # The reported steps are the corresponding entries of the same # decomposition the logged train_loss averages - step_losses = forecaster.compute_step_losses( - init_states, - forcing_features, - target_states, - interior_mask_bool=model.interior_mask_bool, + prediction, pred_std = forecaster( + init_states, forcing_features, target_states + ) + step_losses = torch.mean( + forecaster.score( + prediction, + target_states, + pred_std, + mask=model.interior_mask_bool, + ), + dim=0, ) torch.testing.assert_close(captured["train_loss_unroll1"], step_losses[0]) torch.testing.assert_close(captured["train_loss_unroll3"], step_losses[2]) From 9156e021bc75f03511969276cbee9f7b97c68b32 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 7 Aug 2026 12:35:49 +0530 Subject: [PATCH 32/41] Revert ensemble RMSE logging to a single rank-zero-gated method The all_gather_cat has to run on every rank, so the decorator could not gate the whole method and splitting it in two just to apply it was worse than the guard it replaced. --- neural_lam/models/modules/probabilistic.py | 35 +++++++--------------- 1 file changed, 10 insertions(+), 25 deletions(-) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index cd079df1..345282fb 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -5,7 +5,6 @@ # Third-party import torch -from pytorch_lightning.utilities import rank_zero_only # Local from ... import metrics @@ -178,9 +177,12 @@ def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: """ Log the ensemble-mean RMSE accumulated over a full epoch. - Gathers across devices, which every rank must take part in, then - hands the result to ``_log_gathered_ensemble_rmse`` for the rank-zero - half. + Averages the per-variable MSEs collected by ``_ensemble_step`` over + every sample of the epoch (across both batches and devices) and sums + them over variables, and only then takes the square root. Rooting + each batch's MSE and averaging those roots instead would report a + different quantity, since the square root does not commute with the + averaging. Parameters ---------- @@ -191,31 +193,14 @@ def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: phase : str Logging phase, either ``"val"`` or ``"test"``. """ + # Collective: must be reached by every rank, so it precedes the + # rank-zero check below. entry_mses = self.all_gather_cat(torch.cat(entry_mse_list, dim=0)) # (total_samples, pred_steps, num_state_vars) - self._log_gathered_ensemble_rmse(entry_mses, phase) - @rank_zero_only - def _log_gathered_ensemble_rmse( - self, entry_mses: torch.Tensor, phase: str - ) -> None: - """ - Reduce the gathered MSEs to an RMSE and log it, on rank zero only. - - Averages the per-variable MSEs over every sample of the epoch (across - both batches and devices) and sums them over variables, and only then - takes the square root. Rooting each batch's MSE and averaging those - roots instead would report a different quantity, since the square - root does not commute with the averaging. + if not self.trainer.is_global_zero: + return - Parameters - ---------- - entry_mses : torch.Tensor - Shape ``(total_samples, pred_steps, num_state_vars)``. - Per-variable ensemble-mean MSEs gathered over the whole epoch. - phase : str - Logging phase, either ``"val"`` or ``"test"``. - """ time_step_rmse = torch.sqrt(entry_mses.sum(dim=-1).mean(dim=0)) # (pred_steps,) From b530106497c8fad014a1b17ac85311728bc609b2 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Fri, 7 Aug 2026 12:46:34 +0530 Subject: [PATCH 33/41] Log the forecaster's objective at evaluation instead of ensemble RMSE Summing the ensemble-mean MSE over variables reported a quantity with no clear meaning, and the per-lead per-variable path through aggregate_and_plot_metrics already covers the metric properly. Drop it and log compute_training_loss as {phase}_mean_loss instead, which restores the scalar ModelCheckpoint monitors and keeps it the model's own objective. --- neural_lam/models/modules/probabilistic.py | 122 +++++++++------------ tests/test_probabilistic_forecaster.py | 85 ++++++++------ 2 files changed, 108 insertions(+), 99 deletions(-) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 345282fb..877bd81f 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -21,10 +21,10 @@ class ProbabilisticForecasterModule(BaseForecasterModule): 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. + sampled from the forecaster and its mean scored per lead time and + variable, alongside the forecaster's own objective. 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 @@ -74,8 +74,10 @@ def __init__( Whether to create GIFs of example predictions. Unused here, for the same reason. val_steps_to_log : list of int, optional - Specific predicted steps to log the ensemble-mean RMSE for, - during both validation and testing. + Specific predicted steps to log during validation/testing. + Unused here: nothing this module reports decomposes per step, + since the ensemble metrics are logged for every step at once as + heatmaps and the objective is a single scalar. train_steps_to_log : list of int, optional Specific predicted steps to log during training. Has no effect unless the forecaster's objective decomposes per step; see @@ -121,27 +123,27 @@ def __init__( self.val_metrics: dict[str, list] = {"ens_mse": []} self.test_metrics: dict[str, list] = {"ens_mse": []} - def _ensemble_step(self, batch, phase: str): + def _ensemble_step(self, batch): """ Sample an ensemble and compute the per-variable MSE of its mean. Shared by ``validation_step`` and ``test_step``: samples ``self.eval_ensemble_size`` members and scores the ensemble mean with plain (unweighted) MSE on interior nodes. Only the squared - errors are computed here; they are reduced to an RMSE once per - epoch by ``_log_ensemble_rmse``, since the square root has to be - taken after averaging over every sample rather than per batch. + errors are computed here; ``aggregate_and_plot_metrics`` reduces + them to a per-lead per-variable RMSE once per epoch, since the + square root has to be taken after averaging over every sample + rather than per batch. This 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. + objective ``compute_training_loss`` actually trains on. The + objective itself is logged by ``_log_objective``. Parameters ---------- batch : tuple The batch of data. - phase : str - Logging phase, either ``"val"`` or ``"test"``. Returns ------- @@ -169,58 +171,53 @@ def _ensemble_step(self, batch, phase: str): mask=self.interior_mask_bool, sum_vars=False, ) # (B, pred_steps, num_state_vars) - self._warn_skipped_steps(entry_mses.shape[1], phase) return entry_mses - def _log_ensemble_rmse(self, entry_mse_list, phase: str) -> None: + def _log_objective(self, batch, phase: str) -> None: """ - Log the ensemble-mean RMSE accumulated over a full epoch. + Log the forecaster's own training objective for a batch. - Averages the per-variable MSEs collected by ``_ensemble_step`` over - every sample of the epoch (across both batches and devices) and sums - them over variables, and only then takes the square root. Rooting - each batch's MSE and averaging those roots instead would report a - different quantity, since the square root does not commute with the - averaging. + Reported as ``{phase}_mean_loss``, mirroring + ``DeterministicForecasterModule``, so that ``ModelCheckpoint`` has a + scalar to monitor. What that objective is stays entirely up to the + forecaster, and it is recomputed here rather than derived from the + sampled ensemble, since the two need not agree on either the member + count or the scoring rule. Parameters ---------- - entry_mse_list : list of torch.Tensor - Per-batch per-variable ensemble-mean MSEs, each of shape - ``(B, pred_steps, num_state_vars)``, as returned by - ``_ensemble_step``. + batch : tuple + The batch of data. phase : str Logging phase, either ``"val"`` or ``"test"``. """ - # Collective: must be reached by every rank, so it precedes the - # rank-zero check below. - entry_mses = self.all_gather_cat(torch.cat(entry_mse_list, dim=0)) - # (total_samples, pred_steps, num_state_vars) - - if not self.trainer.is_global_zero: - return - - time_step_rmse = torch.sqrt(entry_mses.sum(dim=-1).mean(dim=0)) - # (pred_steps,) + init_states, target_states, forcing_features, _ = batch + batch_loss, loss_components = self.forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=self.interior_mask_bool, + ) log_dict = { - 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) + f"{phase}_{name}": value for name, value in loss_components.items() } - log_dict[f"{phase}_mean_ens_rmse"] = torch.mean(time_step_rmse) - # No sync_dist: the values are already gathered above and only rank - # zero reaches this point, so a collective here would hang DDP. - self.log_dict(log_dict, rank_zero_only=True) + log_dict[f"{phase}_mean_loss"] = batch_loss + self.log_dict( + log_dict, + on_epoch=True, + sync_dist=True, + batch_size=init_states.shape[0], + ) 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. + Logs the forecaster's objective as ``val_mean_loss``, scores the + ensemble mean against the target states (see ``_ensemble_step``) and + collects per-variable ensemble-mean MSE for epoch-end aggregation. Parameters ---------- @@ -229,16 +226,17 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - entry_mses = self._ensemble_step(batch, "val") + self._log_objective(batch, "val") + entry_mses = self._ensemble_step(batch) self.val_metrics["ens_mse"].append(entry_mses) def test_step(self, batch, batch_idx): """ 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. + Logs the forecaster's objective as ``test_mean_loss``, scores the + ensemble mean against the target states (see ``_ensemble_step``) and + collects per-variable ensemble-mean MSE for epoch-end aggregation. Parameters ---------- @@ -247,31 +245,19 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - entry_mses = self._ensemble_step(batch, "test") + self._log_objective(batch, "test") + entry_mses = self._ensemble_step(batch) self.test_metrics["ens_mse"].append(entry_mses) - def on_validation_epoch_end(self): - """ - Perform actions at the end of the validation epoch. - - Logs the epoch's ensemble-mean RMSE, then defers to - ``BaseForecasterModule.on_validation_epoch_end``, which aggregates - the same per-variable MSEs into heatmaps and clears them. - """ - self._log_ensemble_rmse(self.val_metrics["ens_mse"], "val") - super().on_validation_epoch_end() - def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - Logs the epoch's ensemble-mean RMSE and 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. + Aggregates the 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._log_ensemble_rmse(self.test_metrics["ens_mse"], "test") self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") if self.trainer.is_global_zero and self.hparams.metrics_watch: diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 9c48fe00..b5ecae5f 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -1,7 +1,3 @@ -# Standard library -import math -from types import SimpleNamespace - # Third-party import pytest import torch @@ -417,6 +413,31 @@ def sample_ensemble(self, *args, **kwargs): return super().sample_ensemble(*args, **kwargs) +class ComponentReportingForecaster(MemberCountRecordingForecaster): + """Forecaster splitting its objective into named loss components. + + Records every requested member count in order, so that the objective's + own ensemble can be told apart from the evaluation one. + """ + + def __init__(self, *args, **kwargs): + """Initialize the forecaster and its recording state.""" + super().__init__(*args, **kwargs) + self.num_members_seen: list[int] = [] + self.last_batch_loss = None + self.last_components: dict[str, torch.Tensor] = {} + + def sample_ensemble(self, *args, **kwargs): + self.num_members_seen.append(kwargs.get("num_members")) + return super().sample_ensemble(*args, **kwargs) + + def compute_training_loss(self, *args, **kwargs): + batch_loss, _ = super().compute_training_loss(*args, **kwargs) + self.last_batch_loss = batch_loss + self.last_components = {"kl": torch.tensor(0.5)} + return batch_loss, self.last_components + + def test_probabilistic_module_validation_scores_ensemble_mean(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) @@ -516,9 +537,12 @@ def test_probabilistic_module_test_step_scores_ensemble_mean(): assert torch.all(torch.isfinite(entry_mses)) -def test_ensemble_rmse_takes_root_after_averaging_all_samples(): - """The epoch RMSE roots the mean MSE over every sample, rather than - averaging per-batch roots (the two differ whenever batches disagree).""" +@pytest.mark.parametrize( + "step_name, phase", [("validation_step", "val"), ("test_step", "test")] +) +def test_probabilistic_module_logs_forecaster_objective(step_name, phase): + """Evaluation reports the forecaster's own training objective, giving + ModelCheckpoint a val_mean_loss to monitor.""" datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) config = nlconfig.NeuralLAMConfig( @@ -526,39 +550,38 @@ def test_ensemble_rmse_takes_root_after_averaging_all_samples(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, config=config + forecaster = ComponentReportingForecaster( + predictor, datastore, config=config, train_num_members=2 ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - eval_ensemble_size=2, + eval_ensemble_size=3, ) - # Two batches of one sample and one rollout step, with deliberately - # different squared errors so the two aggregation orders disagree. - d_state = datastore.get_num_data_vars(category="state") - model.val_metrics["ens_mse"] = [ - torch.full((1, 1, d_state), 1.0), - torch.full((1, 1, d_state), 9.0), - ] + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch = ( + init_states, + target_states, + forcing_features, + torch.zeros(B, pred_steps), + ) captured = {} - model.all_gather_cat = lambda tensor: tensor model.log_dict = lambda log_dict, **kwargs: captured.update(log_dict) - model._trainer = SimpleNamespace(is_global_zero=True) - model._log_ensemble_rmse(model.val_metrics["ens_mse"], "val") + torch.manual_seed(42) + getattr(model, step_name)(batch, 0) + + # The objective is the forecaster's own, reported under the same name + # the deterministic module uses, alongside any components it splits out + assert captured[f"{phase}_mean_loss"] is forecaster.last_batch_loss + assert captured[f"{phase}_kl"] is forecaster.last_components["kl"] - # Summed over variables the batches give MSEs of d_state and 9 * d_state, - # so the correct RMSE roots their mean, 5 * d_state. - assert captured["val_mean_ens_rmse"] == pytest.approx( - math.sqrt(5.0 * d_state) - ) - root_of_each_batch_averaged = ( - math.sqrt(d_state) + math.sqrt(9.0 * d_state) - ) / 2 - assert captured["val_mean_ens_rmse"] != pytest.approx( - root_of_each_batch_averaged - ) + # The objective samples its own ensemble, with the member count it + # trains on, before the evaluation ensemble is drawn + assert forecaster.num_members_seen == [2, 3] From e64d2bd6ecdb5e829253da09375674d245cfd4ae Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Sat, 8 Aug 2026 14:45:07 +0530 Subject: [PATCH 34/41] Log the evaluation objective at validation only Nothing monitors a test-phase loss, so computing it there spent a forward pass per batch for a number no one reads. --- neural_lam/models/modules/probabilistic.py | 39 ++++++++++----------- test_ens_rmse.csv | 3 ++ test_ens_rmse.pdf | Bin 0 -> 16342 bytes tests/test_probabilistic_forecaster.py | 24 ++++++++----- 4 files changed, 37 insertions(+), 29 deletions(-) create mode 100644 test_ens_rmse.csv create mode 100644 test_ens_rmse.pdf diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 877bd81f..2149963e 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -22,9 +22,10 @@ class ProbabilisticForecasterModule(BaseForecasterModule): wrapped forecaster assembles its own training loss. Validation and testing are ensemble based instead of deterministic: an ensemble is sampled from the forecaster and its mean scored per lead time and - variable, alongside the forecaster's own objective. 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. + variable, with validation additionally reporting the forecaster's own + objective. 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 @@ -174,23 +175,22 @@ def _ensemble_step(self, batch): return entry_mses - def _log_objective(self, batch, phase: str) -> None: + def _log_objective(self, batch) -> None: """ - Log the forecaster's own training objective for a batch. + Log the forecaster's own training objective as ``val_mean_loss``. - Reported as ``{phase}_mean_loss``, mirroring - ``DeterministicForecasterModule``, so that ``ModelCheckpoint`` has a - scalar to monitor. What that objective is stays entirely up to the - forecaster, and it is recomputed here rather than derived from the - sampled ensemble, since the two need not agree on either the member - count or the scoring rule. + Named as ``DeterministicForecasterModule`` names it, so that + ``ModelCheckpoint`` has a scalar to monitor. What that objective is + stays entirely up to the forecaster, and it costs a forward pass of + its own rather than being derived from the sampled ensemble, since + the two need not agree on either the member count or the scoring + rule. Validation only: nothing monitors the test phase, so paying + that pass again there would buy nothing. Parameters ---------- batch : tuple The batch of data. - phase : str - Logging phase, either ``"val"`` or ``"test"``. """ init_states, target_states, forcing_features, _ = batch batch_loss, loss_components = self.forecaster.compute_training_loss( @@ -201,9 +201,9 @@ def _log_objective(self, batch, phase: str) -> None: ) log_dict = { - f"{phase}_{name}": value for name, value in loss_components.items() + f"val_{name}": value for name, value in loss_components.items() } - log_dict[f"{phase}_mean_loss"] = batch_loss + log_dict["val_mean_loss"] = batch_loss self.log_dict( log_dict, on_epoch=True, @@ -226,7 +226,7 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - self._log_objective(batch, "val") + self._log_objective(batch) entry_mses = self._ensemble_step(batch) self.val_metrics["ens_mse"].append(entry_mses) @@ -234,9 +234,9 @@ def test_step(self, batch, batch_idx): """ Perform a single ensemble test step. - Logs the forecaster's objective as ``test_mean_loss``, scores the - ensemble mean against the target states (see ``_ensemble_step``) and - collects per-variable ensemble-mean MSE for epoch-end aggregation. + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. Parameters ---------- @@ -245,7 +245,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - self._log_objective(batch, "test") entry_mses = self._ensemble_step(batch) self.test_metrics["ens_mse"].append(entry_mses) diff --git a/test_ens_rmse.csv b/test_ens_rmse.csv new file mode 100644 index 00000000..eba95fc5 --- /dev/null +++ b/test_ens_rmse.csv @@ -0,0 +1,3 @@ +1.254199600219726562e+01,6.226966381072998047e+00,1.941929012537002563e-01,3.444211006164550781e+00 +1.272980690002441406e+01,5.614918231964111328e+00,1.835764646530151367e-01,3.391820907592773438e+00 +1.302767562866210938e+01,4.653012752532958984e+00,1.753242611885070801e-01,3.230829238891601562e+00 diff --git a/test_ens_rmse.pdf b/test_ens_rmse.pdf new file mode 100644 index 0000000000000000000000000000000000000000..812d639d414ecbafcbb2638756fa7d4106bfc8e5 GIT binary patch literal 16342 zcmb_@30#d&^mi#sSBexZs;fw$?tbe^v@cTHC@HPCeebptg^)_Cgi5IhMWG!n+V>r8 zC@s>WqUCPyJU4znOaJ%h|9;+l#+_%LSJet2m?PiA{?P)L9%diwS^;eENqFMa5Si(2S0ECMs%@;D$4xSK-JZS1jjHc zAdHS!IuNZ$aOveo6zl5FU z{K>h9^;S%8h}wi6rY74q#8L@olcqeI^ITp8mw&wHvBjQ4DG~VfeQ%ikxZK}~idkNX zHvOji%Ej=TFge+@q%=5VY|Lu_?!DPi<3#lN?PKFwYR`Lpx}@Ge_?Rti^M(I&z3{g- zlefhK`+z?ao9y8_dAMx9ql`Z3)!yw6FNPF% zd@y2ainsZaPdoVfW4z7aHR)k9a%|F@tKN%MCB`>2+t2a_kBY3?+{JX1#C29a`91g7 z8>>rCpS@i#!bQB#e1CL@JltDE$TmB9{OcY*>(=0arXZ(|(?gNF4EnPU47{f1H`WFE zE`6hN~zDRwSaQ(3Wylye3bE=Ka2+TmfrBD$$eF-PB8 zGW(;$MT)EySw|zEwkmA&X00qzCKv@sYhAM$a4VVcxP96a*;^$T!R}Y1b}_5fWCQSgx);$n6ZCsnEV`n?LK@dmi* zp0b&8X2FKAJMqhJ#No}Vd@ZuJ?q9gB%myUCD@84$e^nLrbEiH{EMbxYmtzfm@>%X^~S;7dD}*sps+U$a^0u6@Ry za~fvHHQMgj4`S>(&>d%jAN$@&Epu_6Jlwi(&;6(;SKxdcm8v8{+nn;8Loqey+4|z1 zGyRVaOnj5}+!dX3>|$E0j#&5mW64%bXLi^uly6iW++uj66mcM9EgjN|-?#O5{m9S( zkX|52qkrg3hVFpMP5=e-quxsT186hEe<=^GW(AJ3$aSox0cK^J#mn{RR$*(Y^v+b( zf>?WekB-^(>kY%qRvyRg$0rKSP=`mIOx0FR>SLBRr9V+fubDo72X@Jz2$EJ-_oNcijHX z!Q3SsQJ$iGU(zd@L$s^!CP-EXyw1D0=wQQz^tL}L;TRlip!&3_(k?Ah>9eea7Y7GVW_|bd&F#BYh;aJmnCsgbSjM^T z_TX}=l*4UPOOhn@eV=7n<}29hNnhdmF!Ww|w~F_cr#x4(CHQYnJ$@OPV2p2n@*0&w zvCQ+%jlbDu?&@4IMl^h9e9Kv$I=#YgbqT2-5tv1Sk~dzn1_PQMjoB$Qu>9tM00Fo!;Wky@JC z>TO~1MoSKDArJ4fXVjQT+8ZYB@^MK{@aal2Sk!uRy!-i`ZBzZWr>aLdH?m78tUx&g zwtk}=jYt1CF6eUYyZN>{bN(ZLID6`_Wd?d1 z^|ENw=cdXJZXRY#QIFH3Z+{?lDJM6t7i>Gajro}PxB{t^mt>$SW!#Y17Eya?9p8Fs zW&N#ftiJS=vU}ExO3%X|1%8M!(!F!^QKE8i?{?0V*8*oJw%!U>W9r8*yflL2|IafB9%}ltij>gM1BK&=vGZgqk?x8LB*dpprdTclv8bnoVs64y7 z8`;zy@SxyIvm{?y;z`+81m-l~T-92GON#UBww2xb(4bLXp>iZVnC(@YY=PE~%F zyR*@9ik*&o^-f{dcHK`U)Zg>S+&iIDno%KkbqH5;@4(j97}LVK8C927L2bgBr=uUu zUmi?7v#X>EA>GKaDDputIxzIx(+WS?tizoTE8IvuivTjdfSxAcXIK;~+)`|KA+M=_y3)K&>n8-m_48C`ewtzx_e$ z$zgi;v&2tIY=*1yf*uw+y>=iEwmS%pvy;bd5b>Vt@2xU3Zd$`@oj;sXxwpL%d!tm+RN^N7 zc=j6kdFP>7N&q^ZXYveUU%~2%8x~_04~!1Z=RDhfw5O}kKg{V%q^?NtZ*j`jIJe zdpYBos7!_}w)Pi%nztPoObl;%#U{2!2?tk3)cW7F9Y5V7D(d557UuK5pP-PTl&s|7 zj=Hq%tb^>1$sN&x&9Mptk`+<;t0Ojyn!8`vPFUTd?ZKvPD;Uxb<25otsJg5>^8NljpHqm}^ptnls7{`8{M6?}yd3wXZkhS2cq8dJeezuK z;v(nt%~YEemJ;&pep^a3?k|>7OZSIiT9Rj)$Q{6_v??vp9VJ#{Mz_PQ^tmMN?nF)8 z)(jKjh!a|h=5bOwb1_}GD1K=ct#{HAD9xl`LKz#CFb1qq@`P>%@_1KOtRv|4{`BC z_vK`5U{1f^4sXB7clkX6FL?2FW&s&iB9>`&uX#U{h#xK|(Yd~|*r!;Zt1NpJ=E<7m z{S_?*JR!}iPwuvrZLtWoIPcgPR@F3zyr8MJN`*CEqmy4%jPzK$om}64PJKi1#jjWT zlk=ZVPX*u`RKV1zk}aed6B@`t5tfTrW0s?KYS8ELP zemptN?EIWsE5)=Zf4oG3P-$PUnKE1ux88Kq_7@zdugN@}<@ffEA?C%nUwwR>$D90K z*rhCIpL6p4hI#&pQ5Ir1`9<}>!fECPCQY3`nSogW0^|My9$@f!I!GNPN*9U7>mYNn zx=1`yUmZ1s4tnfSXl_X`-=cn$Bc6?mMI^o~ZVl6#S`OR_gRcNzQTV?YT}4NM3*hf! zGAv16M|Yq-F;VC(-Ar1?HlW2Cu55|piThOIw1LlBm>ZqE+b6nnwF8<2(;=GI>Y9gJ z@^sj)doGqcSgd%?VDD$X9;vF$*_RDeE;igeUoI=7`&scVs%}~q~SuL{P+u$y#vC}L=WCx zVXiA+VKm_{;EYOiy;miiU3uOhh3otaovWT6x#~*;HajSvS0^#^k#@e$yQIOAEE!`G z_hNXuum(vBc@X0l84xWg`lw8``_knB6^Xlr#^YSpH3jU1 z`K&sgAkUP6R&L>n7InX;EMpXY9-7Rn&3o|)xij0gG&CeyrbMW>Odgw_b?0!I_1mr7 zv>vy16vpE4X1AjuYqtJcKjJ8j3ZPV+E!#wrRlr5^i@g%{3nh|^Y z)_Gj|_5N6qw%B0((A+tlYvUcci#QeAS<&bFRZd#vgdRLO=-J)BZ^B*dc=ql~sS2Fk zA__;FbmQC-oligbx&}KJ$Wv9FSS7!wT=31;0}jWegR1%;P4OQbBnw{)3oKZH*(;!4 z4CXJu4h}dTEGUNKAgm`KmbeBS+(dx|*L(u=eb;~n5Xg*ck59-RQYwC1psSwvYu_(4O#f{XTR$`W>6 z=UPXv&6|%t#4XV2aGmcs-eWQR^8T9RtT_Z6-mHARd(*wklIOhdoAFAFpF160-X~zV z@Xq^o@Rp7P;^X`huY_$4q6@a)BG-wj_8|l_Vyh+5B-z_{UTTUQ+_dMy(D|iL^4X#{ zGjh-E_dnN~RkWd_bNATH_D248;=J8mJlj$n+eeE7YyBT36Uickc~|>@`g?nF?E;HG z;I)M&pLZmZl-CGP?fOuc{WzxGD#f!Zb?|A3K6&$~T(~vg7wb`TLzzmW{m1fphHBYd z$Yl@i$Zu*buJIn3ORMjVPKG1;^lwttjyRHOYsV~4d#q z0WisFrg6uBr2X8_t?7;$J6Er|6r~bspcNxQ_;^PuL2OLQOz^>%HF8HZHl`JP9tfFPEy3Q36|Rq$*I9p2wbUy`Z@Qz^%`@4e zKhnq9bYJZGeL_}}4V^BDZz?KlxNhw$t8VeUVkNPdXRA#*`DW2I;B&5C1#WLLedC(ZtgWo1!MZ)L!Hu77N6(hr|J2fdMsg@K(G5wnb_#-D(#EiJy$@d7q!fc{5$*jr%xzhGq z0kL8K=By5!$#}2=FJ^wvg8cmk%f|zlOal%K{iSMoAK`EcO7T*N+(q^T{TCg?2?-vt z0sn*xLsE+-p=u(M@g`FHG(9p(q&CDomHZ^-Iez(l1JmkHHS4C2Z$5B6=n(tnK8Y83 zk2wd&bL~gcQ~C0$O$S9!RN{=|zRcO5Z?aCy$|m&RQ_MF=C)*}e=a~Q z3%6+6)?#E%WW6QdsNaxvC)+$^-9fZGQA1)NWs}n?jkw9OksD9whNu%goB`pDpVB)m z1y3Yh+Mw~c*VMxK78mwSD3TiXIY1!p{rQz0t&V6|YqE@3nPMqbb-pn%^9ZI4$lqH~_ zJlcd&NE$uOH}{%Py*gZ|UjP^K;%>-V`JqrmmFW71PVUblRCZoO}ISxE3T0w$WLb&L^K;@F!Ct%>vhukbd3JRXs;dMGGJU6+I?h`G$PJEQ8_0s&_ zg^FDwTY8U_oQ`}G@o0K)0iTq8(4mr$db^=`uTmeJUhS#tq2A{wc5x-$Z(-X4dwt|a za&5iA;!wYX#nQ%@=bB#?Vf$aRTcO=_$c0zkBNBs_49`?iSqkmMS2(=Ft^iI^;Oo|Z znhKRj1un(y?BIMeRqeL@Z8sso>I?m69PbVKlr;+&w^7jXURN?XFNl^;_QP5yOu37D z9&DNw;XnN;4Hc3z?)E0DGQ_$(QZ~FaPLYvycD>ihTDA+9@i`e!lIo*hCpv=P0 z%kQnQb2_=>8Oxb>CU-q)77JTqdsLK#%4@f-*_mo1k{?cMyqTNEnoBHt;9vnShrAnsS2F`3-UuJe-lQrY4hdF$m zLk^3b9SKJ?I=X9mki0ubp5mIPrq;XgDc$MRd3+12P?>w&&NU|2a+Tnw&KJ%530$UY zcUB*X)t25{$5kaf`z?>Zgb#k-(J?etxne$U5A9BnFs}y&BW!%3XO6AEDMY3~K8iJA zgQ&_2Kh3SlfxZ>_mlsNcUCzj8Ul2Ep4tjvCYd^q07CCdJ(dMyh?wA@|c2!q)49|t$ zsN@a8+-jeMwbn1D%W2e8Z>{1HIcD6u%Kw!}(~?n{d8wyEnFQO;PqlWJ5!)_Sj_-~0 zN|!XaCfB|2;zf3|nHA@m&nv9Y3SeIPFW{wC1B#;*&aTudcXxUdd_zd@-TV8_lOJt} z3oUx>y0rgNumvGJBs=0sok+1&=ca&A>@9}j9-sZ1Kb}x&$*Do_j=ZteoXLG|>)1@RWamt})Eg;5 zv4Nsq)!Y~?6#&8%|)+sD8v%YeRfyy4ie*A;b7xXzqJ%`GO&~JyY-n-4S(~xHC||YC5w+NF z#{LA}WPh}0Tw91A;YmoVhxB(wN&QcZ&tK^^k7fE3W;Z ziFp^+2J2<-bQc@3v&g%*uFmV8j!<4dT7;N!Kla?q9%bLkw9gxcJzBLzLT>l3Pvv~G z=B`XR=X;r3fv5XQ9Z!BW$x-eKLdu%>oq7n#9No(oc|~|Wh9ZvVN6x?V zxxD%iw_dQ`xk=@OcZD4ZIV_UNhIv*WJoy&+3!{!~K6xn5dS^{6-@0RZUY0{+heW;3 zZ(ER9=$jr5_1tY}^QLb;D6VYMAZ*lR#>{BTjlPGw)^UpCo{DTp6yN`(O&uX?W{|(# z{(<|5bMBZ-%KXH-qaGo(qg8{75wkTV3kr&F6WVF586%e5(Jw}CP(+ghxdh@L`vyoO zPS2g68QhRAw8nUZ<)Pqh>sq1QyMyh4o2+M!Y*Q1k*EvrPNq+kdygSWd}CvAG}%97zS%|_!d(*V&=BPbV6cg z$tm4R)^Upc_$=?&>EXw0wwnF%>S<1@E^cB|aZT@%>9fxQyXJ+n-=AOM!!C9Okc&p+ z{#)o(bTuLX00JK#PX#~gnbOdmi(I68iQBdicI?=jAoB6)D`hEK`Ib5!1YhdERo9;M?7)J<*aYadyBxm7@VRHmQVXZu;%x- z)go%8C7Uamr8wp{81nb{Pv-=@KBQ`m^)j9`-*xI+FyC3hLh5@{wCd|=Jl>oZ4 zR>sT5g(&vRjc#fe#iT5GKI#o4^S(7ffo%GH=y9s23p`h9Z`9HmWGMLi}WgTh&> zg&R%W+^_jhv~NpS(@#)|W)3`E9!%6Nd(pG&{FXx_Z#P6gnYDXLyU?ETa8Ygdq8OXY z=Q0iUhBC~&r6zmV#VO@3&Ae9Kk%s33U8~k^YxL8KeUPBvb|{Yd>Y>1;*NFK$?KRK) zPCm3yYWs8kxD_Dg-}rhM98!l#!Bo^rf3-ebt-Gb|Fi_6xo0={l6=nz~HSCer9EzqK zhHDJ@nUt?{EUqx{3J{O*7iw86LKcKfK+slXtGq)S_$;GvDW&(!&ZEZWIY~|69+LO#)x>>wDkEkD|FnEuV?05K=UVt(dFlwx3E) znG)Wa&%1B7v*wz2)c%RO`QeSI1a^GNu(#vZp#D!&yPFp#vI<7;if^s8L{A?iA}8m*%Hr=&s6bFE%}$Rj43)GJ^TRF7IQaSKGs`e{(1|ES`Rg{TQHsiC1&Y z37scXIh-fjV$$i|?+zV2Q=_I&`;$AUpUuH?)FC635TR({N&KEOQd2p0c#ksC!Qz;g zKGDP822w;Go+KqZ3lH!i3US!tdmfF$!4Oa*Wum8*hrJug)dPL(U9SosmLGrL?#0*Lyk&rNg zfB)p5*g-&?4hUrg;ewDh1#yllE+G033UvH8!16~C454gqV?zXCicnnQ2@rn>M|iqf zSP|hEU>V(rAY2gyFam!L;ReMrIuUJ1-?JX}wss&7H~<*q=?KRHEp6xLW=C{^LGA2a ztwB#t7M^x+gg?>46^?Lq0n;Y=fE>xrg9!3AAQ}?ydU#s7dJy3RK%)hSiUj0=(TE@i&>0xe9^^nnXNYfL zVDN*5A;4IkU~JF<9N`UOpg=Gr!qDEDWCs`pL9L(TDj93{<2&QO!tvAUGM3_hT1v1` zP8PPFfSZ4_s0yA8CKJ$rWM~A4-V+|-x3frKH9eq=!mDu|A?w^DGibs_@7 z5&9&e^D!u+W#P-15gdiVqJIuh{Q)IU!0~`t{~r?U|FsVc5K1f(4aee)v=|9|x5R^Og@furzu=1! z25fmKkH=zx@a#>9bN23!6i4iL$ZAZU76=m7kOJD%7@_>JkRDcBg{tHRP zFG&V`FaIGFkU;#B0AY-0AQ+Hvd`~cjT!51>JRUF($iPa8@Let-nfN&|WCD_l<@8+^ zuz;;#{eMf#YrQNNkZk-&kX$Uw9qy+v(4SGa3LFJNOFc;Af`J&h!@vqH=PclWTa2=Q z6hHcV@Vzx0Y!dJa1{5)X?2Jmw$_xG75;9ah5K6`?hV}=$h4Ekq2g+|b$1ovq#^-I& zl|2xY z%a#vlmF03^3>h8|V73<=pjSrePxg%ga=}Y)u*sm8uw_$Ljm2SPhue@6~=XF+VH}fEfSM1l>laZ?D z6(q-clW(7s0GkD&XixIgA$lmeI=cZZ1Hqyd{Zas_7`H9J5t%2P5a@>;hl{fP;Vb4XwiW$Q)p4AkhFN|6Jd&nhUI}sjTYQRmv;2nrUOoDuX{E zE~6R5J~R?~3FR0p{$FtvXoDXt#mEr;f~jb%3{0Fh&`J(ywgNNtYg>C_MJOgwu_L@6PpC#jRHaqd_uJG z59Bjuw7zC)V<9E+EZJF`+Feb_3?pB5oHJ9VJWZg_P17cZ=#(W|ZzHv#fO6{$*-4A` zd5BJ<(xyjg{VmiN8FVUza_(KK^hx@{SIX1tlqd1@g*n=2AGPKQ z<+?xF_c-m-drIngai#F=>)@onPRc3(@ ze=(AuI(Bf<;J}2Q+_<(ZZKP-P5Me}JdRP_zQ5iR+h^0Kc`9U7@Za<3t?MuH5f;RlN z7b{7vdi)l>=M8fAYbg^6vbSN^?p?ij$#$`wZ97C;MMPWR+na^9Q}S*#ZW9Jug!Xxm z_O6{yUZgy{O8+`dK4nV#_=fCbM2YjEO@AVL9i=>urH%B`7w4(fIkYdI=rkI&>^_}B zrq4}L;{7No=jd~j6hr= zX2=(8Xz$vnrS~Y2?qoL|N6ayI!yw$l7?9MT;PAN zA)rn8hl-G4DC!?7;ve1lkBSJdF=W}Yib?goD7WI~o!H2_Z1FvD8ButyI=e$Ex@p)} zsVGJDdU1Y5@7*vSXGIm&Ts`+UoHA3D^CkyM?9ycG?>VFel*zQKOIP?dsrwXtc^Eax zcPBK;rAYUUx?!#FTXn-$nMU>b*r-W^K@HJ9?m=$RK5MHRQQ5($R&}3T_q$OpCA$6U zhE2S-X|lDvR%rsxx9rjadVL$z4ZD0h)D5e|ozn#DW6^DK*tT%wA>OeNfr>CvHJ?&A zDKX3{oOG49ns-bJ8^Svlfsa79ZNP`}jS1n`@Qp=ah0tw+7ABTyvfVQM>htkA_3HVR z>iJc^1M2fPP@U@Yxvkae-IaE>X^xf+xHd1$eSxtTsD4yipY2n0TcQQUC5>FO@K)WY z*gYs}5;>?T+IOz%y=qaaYSF>)MAf3kDhJWNfT}&JMV#U3szv=(rm97Gs#c<=ri0r= zP2H_-Momf%o)qo-T9vO_G!Vw3S|l94q*^o^P7foQ@s8ODn1+)!gl!HZ&4zDLEou%= zQY{iv)e`M98uSzGdsa07`VH#|Bb^9O4kxjMi>Ma)t15~130HB5_E}oRL``Bu`--a& zsztlPq{B!9eDl0xfw;H4WBhzxVI*6DgM4Gz7$Ia^pM!~W+6khGRa$Xv!GL;x`D?)F zW?u8O;tD(Gv>w;A;<|#@>V_2sI;EYk-G+rd7YHYv4*#lJivVcVWN z?8Ua>FxLde_VEgYk@kcSs}|*lONW!5g%MPX`l?t(`-JYa6j!Qu_ZNT;`o!n<2|N-Q zGv=rBk9F}s2_sR%uY{2_!=@zrZZB%jPeP|GViIhRgps-jp62AwD+l7)rB6msr1fG07h&C|DLxcgB_4^-i zaBu+Q+Q9!QgH9@-Ke)^HG7#|sf%RX@!2Js_-QUY_0E+!qCJp?@Kgy(6s0WTbf2{|e zfN$}883qFK(BD7wAyF7`Z1Y+$Kj3rzUJr$qfsptw{h=@@@CEm`G8F1h9v}%2 z`1!Rh5(Aw1-^$QXq{6RdNE8;r)W5VvqVT{^{Jl&XoFe>QMu2eGuWiv7$hY{d3`0{}=;!&VTfW!9zaDuWbRF|7eSa;=O*W2T9-`WjM%p`L#cA-0(+R918MmeyxYY z;r>`loXnr&f^&pFc#20u?%uEc;jw@64^M!cxL@l@Lx&W Date: Sat, 8 Aug 2026 15:21:00 +0530 Subject: [PATCH 35/41] Make pred_std optional in the metric contract mse and mae documented pred_std as unused but implemented themselves through wmse and wmae with torch.ones_like(pred_std), so the argument was load-bearing for its shape and callers holding no std had to fabricate one just to be ignored. Compute them directly instead and default pred_std to None across every metric, with the std-dependent ones raising a clear ValueError rather than failing inside the maths. Metrics now also declare the requirement through requires_pred_std, so DeterministicForecaster resolves the per_var_std fallback only for a scoring rule that uses it and an unweighted loss needs no config. --- neural_lam/metrics.py | 149 ++++++++++++++---- .../models/forecasters/deterministic.py | 63 ++++---- neural_lam/models/modules/deterministic.py | 14 +- neural_lam/models/modules/probabilistic.py | 8 - tests/test_prediction_model_classes.py | 19 +++ tests/test_probabilistic_forecaster.py | 3 - tests/test_probabilistic_objectives.py | 37 ++++- 7 files changed, 214 insertions(+), 79 deletions(-) diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 1eb1d526..1296a084 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -35,6 +35,38 @@ def get_metric(metric_name: str) -> Callable[..., torch.Tensor]: return DEFINED_METRICS[metric_name_lower] +def _require_pred_std( + pred_std: Optional[torch.Tensor], metric_name: str +) -> torch.Tensor: + """ + Return ``pred_std``, raising if a std-dependent metric was given none. + + Parameters + ---------- + pred_std : torch.Tensor or None + The standard deviation the metric was called with. + metric_name : str + Name of the calling metric, used in the error message. + + Returns + ------- + torch.Tensor + ``pred_std`` unchanged. + + Raises + ------ + ValueError + If ``pred_std`` is ``None``. + """ + if pred_std is None: + raise ValueError( + f"{metric_name} scores a predicted distribution and so requires " + "pred_std, but got None. Only the unweighted metrics (mse, mae) " + "can be computed without one." + ) + return pred_std + + def mask_and_reduce_metric( metric_entry_vals: torch.Tensor, mask: Optional[torch.Tensor], @@ -88,7 +120,7 @@ def mask_and_reduce_metric( def wmse( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -105,9 +137,10 @@ def wmse( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor + pred_std : torch.Tensor or None, optional Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation used as per-entry weight. + standard deviation used as per-entry weight. Required here; ``None`` + raises. Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -122,7 +155,13 @@ def wmse( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None``. """ + pred_std = _require_pred_std(pred_std, "wmse") entry_mse = torch.nn.functional.mse_loss( pred, target, reduction="none" ) # (..., num_grid_nodes, num_variables) @@ -141,7 +180,7 @@ def wmse( def mse( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -158,10 +197,10 @@ def mse( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor - Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation (unused; ``pred_std`` is replaced by ones - internally). + pred_std : torch.Tensor or None, optional + Unused. Accepted so that every metric in ``DEFINED_METRICS`` shares + one signature and callers can stay agnostic about which they got. + Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -177,16 +216,19 @@ def mse( ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ - # Replace pred_std with constant ones - return wmse( - pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars + entry_mse = torch.nn.functional.mse_loss( + pred, target, reduction="none" + ) # (..., num_grid_nodes, num_variables) + + return mask_and_reduce_metric( + entry_mse, mask=mask, average_grid=average_grid, sum_vars=sum_vars ) def wmae( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -203,9 +245,10 @@ def wmae( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor + pred_std : torch.Tensor or None, optional Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation used as per-entry weight. + standard deviation used as per-entry weight. Required here; ``None`` + raises. Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -220,7 +263,13 @@ def wmae( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None``. """ + pred_std = _require_pred_std(pred_std, "wmae") entry_mae = torch.nn.functional.l1_loss( pred, target, reduction="none" ) # (..., num_grid_nodes, num_variables) @@ -239,7 +288,7 @@ def wmae( def mae( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -256,10 +305,10 @@ def mae( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor - Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation (unused; ``pred_std`` is replaced by ones - internally). + pred_std : torch.Tensor or None, optional + Unused. Accepted so that every metric in ``DEFINED_METRICS`` shares + one signature and callers can stay agnostic about which they got. + Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -275,16 +324,19 @@ def mae( ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ - # Replace pred_std with constant ones - return wmae( - pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars + entry_mae = torch.nn.functional.l1_loss( + pred, target, reduction="none" + ) # (..., num_grid_nodes, num_variables) + + return mask_and_reduce_metric( + entry_mae, mask=mask, average_grid=average_grid, sum_vars=sum_vars ) def nll( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -301,9 +353,10 @@ def nll( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor + pred_std : torch.Tensor or None, optional Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation of the Gaussian. + standard deviation of the Gaussian. Required here; ``None`` raises. + Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -318,7 +371,13 @@ def nll( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None``. """ + pred_std = _require_pred_std(pred_std, "nll") # Broadcast pred_std if shaped (num_variables,) via distribution internals dist = torch.distributions.Normal( pred, pred_std @@ -333,7 +392,7 @@ def nll( def crps_gauss( pred: torch.Tensor, target: torch.Tensor, - pred_std: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -351,9 +410,10 @@ def crps_gauss( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor + pred_std : torch.Tensor or None, optional Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation of the Gaussian. + standard deviation of the Gaussian. Required here; ``None`` raises. + Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -368,7 +428,13 @@ def crps_gauss( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None``. """ + pred_std = _require_pred_std(pred_std, "crps_gauss") std_normal = torch.distributions.Normal( torch.zeros((), device=pred.device), torch.ones((), device=pred.device) ) @@ -395,3 +461,30 @@ def crps_gauss( "nll": nll, "crps_gauss": crps_gauss, } + +# The metrics that weight by, or parameterize a distribution with, pred_std, +# i.e. exactly those calling _require_pred_std. Kept in step with the guards +# by test_pred_std_requirement_matches_declaration. +_STD_DEPENDENT_METRICS = frozenset({wmse, wmae, nll, crps_gauss}) + + +def requires_pred_std(metric: Callable[..., torch.Tensor]) -> bool: + """ + Return whether ``metric`` needs a ``pred_std`` to be computed. + + Lets a caller holding a metric decide whether it has to come up with a + standard deviation at all, rather than assuming every metric uses one. + + Parameters + ---------- + metric : callable + A metric from ``DEFINED_METRICS``, e.g. as returned by + ``get_metric``. + + Returns + ------- + bool + True if calling ``metric`` without ``pred_std`` raises + ``ValueError``. + """ + return metric in _STD_DEPENDENT_METRICS diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 711c5932..706ea7f2 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -44,10 +44,11 @@ def __init__( config : NeuralLAMConfig or None, optional Configuration used to compute the constant per-variable std substituted for ``pred_std`` when the forecast carries no std of - its own. Required in that case for ``score`` and - ``compute_training_loss`` to work (they raise ``ValueError`` via - ``_resolve_pred_std`` otherwise); forecasters used purely for - inference can omit it. Default ``None``. + its own. Needed only when ``loss`` is a scoring rule that uses a + std; without it in that case ``score`` and + ``compute_training_loss`` raise ``ValueError`` via + ``_resolve_pred_std``. Forecasters used purely for inference can + always omit it. Default ``None``. loss : str, optional The scoring rule (from ``neural_lam.metrics``) applied by ``compute_training_loss``, stored as ``self.loss``. Default @@ -116,8 +117,8 @@ def compute_training_loss( Raises ------ ValueError - If the forecast carries no std of its own and no - ``per_var_std`` fallback is available; see + If ``self.loss`` needs a std, the forecast carries none of its + own and no ``per_var_std`` fallback is available; see ``_resolve_pred_std``. """ prediction, pred_std = self( @@ -133,9 +134,9 @@ def compute_training_loss( def _resolve_pred_std( self, pred_std: Optional[torch.Tensor] - ) -> torch.Tensor: + ) -> Optional[torch.Tensor]: """ - Return ``pred_std``, or the constant ``per_var_std`` fallback. + Return the std ``self.loss`` should be applied with. Parameters ---------- @@ -145,23 +146,30 @@ def _resolve_pred_std( Returns ------- - torch.Tensor - ``pred_std`` unchanged when given; otherwise ``self.per_var_std``. + torch.Tensor or None + ``pred_std`` unchanged when given; otherwise + ``self.per_var_std``, or ``None`` if ``self.loss`` ignores the + std anyway. Raises ------ ValueError - If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is - available (this forecaster was constructed without ``config``). + If ``pred_std`` is ``None``, ``self.loss`` needs one, and no + ``per_var_std`` fallback is available (this forecaster was + constructed without ``config``). """ if pred_std is not None: return pred_std + if not metrics.requires_pred_std(self.loss): + return None if self.per_var_std is None: raise ValueError( - "No pred_std available for scoring: the forecast carries no " - "std 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." + "No pred_std available for scoring: this forecaster's " + "scoring rule needs one, the forecast carries no std and " + "there is no per_var_std fallback (it was constructed " + "without config). Pass config to the constructor, use a " + "predictor that outputs its own std, or score with an " + "unweighted metric." ) return self.per_var_std @@ -178,9 +186,10 @@ def score( Apply this forecaster's scoring rule to an already-produced forecast. Resolves ``pred_std`` via ``_resolve_pred_std`` (substituting - ``self.per_var_std`` when ``None``), then applies ``self.loss``. Used - to report the loss on a forecast the caller already has, at a - reduction of its choosing, without recomputing it. + ``self.per_var_std`` when ``None`` and ``self.loss`` needs one), then + applies ``self.loss``. Used to report the loss on a forecast the + caller already has, at a reduction of its choosing, without + recomputing it. Only the loss goes through here. Metrics unrelated to the training objective depend on nothing but the shapes of the three tensors @@ -199,8 +208,8 @@ def score( Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. Predicted standard deviation for ``prediction``; ``None`` when the forecast carries no std, in which case ``self.per_var_std`` - is substituted (see ``_resolve_pred_std`` for when this raises - instead). + is substituted if ``self.loss`` needs a std at all (see + ``_resolve_pred_std`` for when this raises instead). mask : torch.Tensor or None, optional Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``self.loss``. average_grid : bool, optional @@ -217,8 +226,8 @@ def score( Raises ------ ValueError - If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is - available; see ``_resolve_pred_std``. + If ``pred_std`` is ``None``, ``self.loss`` needs one and no + ``per_var_std`` fallback is available; see ``_resolve_pred_std``. """ pred_std = self._resolve_pred_std(pred_std) return self.loss( @@ -259,10 +268,10 @@ 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. Required in that case for ``score`` and - ``compute_training_loss`` to work (they raise ``ValueError`` - otherwise); forecasters used purely for inference (``forward``) - can omit it. + its own. Needed only when ``loss`` is a scoring rule that uses a + std; without it in that case ``score`` and + ``compute_training_loss`` raise ``ValueError``. Forecasters used + purely for inference (``forward``) can always omit it. loss : str, default "wmse" The scoring rule (from ``neural_lam.metrics``) applied by ``compute_training_loss``. diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 5480c86b..465fc553 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -209,15 +209,10 @@ def validation_step(self, batch, batch_idx): self._log_step_loss(time_step_loss, mean_loss, "val", batch_size) # Reported independently of the training objective, so computed here - # rather than through the forecaster. metrics.mse ignores the std - # argument, but requires one - std_placeholder = torch.ones( - target_states.shape[-1], device=target_states.device - ) + # rather than through the forecaster entry_mses = metrics.mse( prediction, target_states, - std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) @@ -250,16 +245,11 @@ def test_step(self, batch, batch_idx): self._log_step_loss(time_step_loss, mean_loss, "test", batch_size) # Reported independently of the training objective, so computed here - # rather than through the forecaster. Both ignore the std argument, - # but require one - std_placeholder = torch.ones( - target_states.shape[-1], device=target_states.device - ) + # rather than through the forecaster for metric_name in ("mse", "mae"): batch_metric_vals = metrics.get_metric(metric_name)( prediction, target_states, - std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 2149963e..19a150b7 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -3,9 +3,6 @@ # Standard library import warnings -# Third-party -import torch - # Local from ... import metrics from ...config import NeuralLAMConfig @@ -160,15 +157,10 @@ def _ensemble_step(self, batch): 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 - ) entry_mses = metrics.mse( ensemble_mean, target_states, - std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) # (B, pred_steps, num_state_vars) diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index bcadb1e9..1981c61d 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -156,6 +156,25 @@ def test_ar_forecaster_without_config_raises_on_use_not_construction(): ) +def test_ar_forecaster_without_config_scores_with_unweighted_loss(): + """The std fallback is only needed by scoring rules that use a std, so + an unweighted loss must score without a config, rather than demanding a + per_var_std it would immediately ignore.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + forecaster = DeterministicARForecaster(predictor, datastore, loss="mse") + 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) + + scored = forecaster.score(prediction, target, None) + torch.testing.assert_close(scored, torch.full((B,), float(d_state))) + + class MeanAbsObjective(Forecaster): """Test-only objective class taking a constructor argument of its own. diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 870a4c1d..f6b2d957 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -122,8 +122,6 @@ def test_ar_forecaster_training_loss_matches_direct_score(): 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] - forecaster.per_var_std = torch.ones(d_state) batch_loss, loss_components = forecaster.compute_training_loss( init_states, @@ -137,7 +135,6 @@ def test_ar_forecaster_training_loss_matches_direct_score(): score_metric( prediction, target_states, - forecaster.per_var_std, mask=interior_mask_bool, ) ) diff --git a/tests/test_probabilistic_objectives.py b/tests/test_probabilistic_objectives.py index 49075f08..6e4de197 100644 --- a/tests/test_probabilistic_objectives.py +++ b/tests/test_probabilistic_objectives.py @@ -1,8 +1,19 @@ # Third-party +import pytest import torch # First-party -from neural_lam.metrics import crps_gauss, nll, wmae, wmse +from neural_lam.metrics import ( + DEFINED_METRICS, + crps_gauss, + get_metric, + mae, + mse, + nll, + requires_pred_std, + wmae, + wmse, +) def _single_residual_case(residual=1.0): @@ -66,6 +77,30 @@ def test_crps_gauss_prefers_calibrated_std_to_extreme_scales(): assert torch.all(residual_scale_loss < large_loss) +@pytest.mark.parametrize("metric", (mse, mae)) +def test_unweighted_metrics_ignore_pred_std(metric): + pred = torch.tensor([[[[0.0, 1.0], [2.0, 3.0]]]], dtype=torch.float32) + target = torch.tensor([[[[1.0, 0.0], [1.0, 5.0]]]], dtype=torch.float32) + pred_std = torch.tensor([1.5, 0.5], dtype=torch.float32) + + torch.testing.assert_close( + metric(pred, target), metric(pred, target, pred_std) + ) + + +@pytest.mark.parametrize("metric_name", sorted(DEFINED_METRICS)) +def test_pred_std_requirement_matches_declaration(metric_name): + """Every metric raises without a pred_std iff it declares needing one.""" + metric = get_metric(metric_name) + pred, target = _single_residual_case() + + if requires_pred_std(metric): + with pytest.raises(ValueError, match="requires pred_std"): + metric(pred, target) + else: + metric(pred, target) + + def test_probabilistic_losses_support_pred_std_broadcasting(): pred = torch.tensor( [[[[0.0, 1.0], [2.0, 3.0]]]], From ce5d6014cfd43b2b6255155d39549986227dd9fb Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Mon, 10 Aug 2026 16:26:10 +0530 Subject: [PATCH 36/41] Remove accidentally committed evaluation artifacts --- test_ens_rmse.csv | 3 --- test_ens_rmse.pdf | Bin 16342 -> 0 bytes 2 files changed, 3 deletions(-) delete mode 100644 test_ens_rmse.csv delete mode 100644 test_ens_rmse.pdf diff --git a/test_ens_rmse.csv b/test_ens_rmse.csv deleted file mode 100644 index eba95fc5..00000000 --- a/test_ens_rmse.csv +++ /dev/null @@ -1,3 +0,0 @@ -1.254199600219726562e+01,6.226966381072998047e+00,1.941929012537002563e-01,3.444211006164550781e+00 -1.272980690002441406e+01,5.614918231964111328e+00,1.835764646530151367e-01,3.391820907592773438e+00 -1.302767562866210938e+01,4.653012752532958984e+00,1.753242611885070801e-01,3.230829238891601562e+00 diff --git a/test_ens_rmse.pdf b/test_ens_rmse.pdf deleted file mode 100644 index 812d639d414ecbafcbb2638756fa7d4106bfc8e5..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 16342 zcmb_@30#d&^mi#sSBexZs;fw$?tbe^v@cTHC@HPCeebptg^)_Cgi5IhMWG!n+V>r8 zC@s>WqUCPyJU4znOaJ%h|9;+l#+_%LSJet2m?PiA{?P)L9%diwS^;eENqFMa5Si(2S0ECMs%@;D$4xSK-JZS1jjHc zAdHS!IuNZ$aOveo6zl5FU z{K>h9^;S%8h}wi6rY74q#8L@olcqeI^ITp8mw&wHvBjQ4DG~VfeQ%ikxZK}~idkNX zHvOji%Ej=TFge+@q%=5VY|Lu_?!DPi<3#lN?PKFwYR`Lpx}@Ge_?Rti^M(I&z3{g- zlefhK`+z?ao9y8_dAMx9ql`Z3)!yw6FNPF% zd@y2ainsZaPdoVfW4z7aHR)k9a%|F@tKN%MCB`>2+t2a_kBY3?+{JX1#C29a`91g7 z8>>rCpS@i#!bQB#e1CL@JltDE$TmB9{OcY*>(=0arXZ(|(?gNF4EnPU47{f1H`WFE zE`6hN~zDRwSaQ(3Wylye3bE=Ka2+TmfrBD$$eF-PB8 zGW(;$MT)EySw|zEwkmA&X00qzCKv@sYhAM$a4VVcxP96a*;^$T!R}Y1b}_5fWCQSgx);$n6ZCsnEV`n?LK@dmi* zp0b&8X2FKAJMqhJ#No}Vd@ZuJ?q9gB%myUCD@84$e^nLrbEiH{EMbxYmtzfm@>%X^~S;7dD}*sps+U$a^0u6@Ry za~fvHHQMgj4`S>(&>d%jAN$@&Epu_6Jlwi(&;6(;SKxdcm8v8{+nn;8Loqey+4|z1 zGyRVaOnj5}+!dX3>|$E0j#&5mW64%bXLi^uly6iW++uj66mcM9EgjN|-?#O5{m9S( zkX|52qkrg3hVFpMP5=e-quxsT186hEe<=^GW(AJ3$aSox0cK^J#mn{RR$*(Y^v+b( zf>?WekB-^(>kY%qRvyRg$0rKSP=`mIOx0FR>SLBRr9V+fubDo72X@Jz2$EJ-_oNcijHX z!Q3SsQJ$iGU(zd@L$s^!CP-EXyw1D0=wQQz^tL}L;TRlip!&3_(k?Ah>9eea7Y7GVW_|bd&F#BYh;aJmnCsgbSjM^T z_TX}=l*4UPOOhn@eV=7n<}29hNnhdmF!Ww|w~F_cr#x4(CHQYnJ$@OPV2p2n@*0&w zvCQ+%jlbDu?&@4IMl^h9e9Kv$I=#YgbqT2-5tv1Sk~dzn1_PQMjoB$Qu>9tM00Fo!;Wky@JC z>TO~1MoSKDArJ4fXVjQT+8ZYB@^MK{@aal2Sk!uRy!-i`ZBzZWr>aLdH?m78tUx&g zwtk}=jYt1CF6eUYyZN>{bN(ZLID6`_Wd?d1 z^|ENw=cdXJZXRY#QIFH3Z+{?lDJM6t7i>Gajro}PxB{t^mt>$SW!#Y17Eya?9p8Fs zW&N#ftiJS=vU}ExO3%X|1%8M!(!F!^QKE8i?{?0V*8*oJw%!U>W9r8*yflL2|IafB9%}ltij>gM1BK&=vGZgqk?x8LB*dpprdTclv8bnoVs64y7 z8`;zy@SxyIvm{?y;z`+81m-l~T-92GON#UBww2xb(4bLXp>iZVnC(@YY=PE~%F zyR*@9ik*&o^-f{dcHK`U)Zg>S+&iIDno%KkbqH5;@4(j97}LVK8C927L2bgBr=uUu zUmi?7v#X>EA>GKaDDputIxzIx(+WS?tizoTE8IvuivTjdfSxAcXIK;~+)`|KA+M=_y3)K&>n8-m_48C`ewtzx_e$ z$zgi;v&2tIY=*1yf*uw+y>=iEwmS%pvy;bd5b>Vt@2xU3Zd$`@oj;sXxwpL%d!tm+RN^N7 zc=j6kdFP>7N&q^ZXYveUU%~2%8x~_04~!1Z=RDhfw5O}kKg{V%q^?NtZ*j`jIJe zdpYBos7!_}w)Pi%nztPoObl;%#U{2!2?tk3)cW7F9Y5V7D(d557UuK5pP-PTl&s|7 zj=Hq%tb^>1$sN&x&9Mptk`+<;t0Ojyn!8`vPFUTd?ZKvPD;Uxb<25otsJg5>^8NljpHqm}^ptnls7{`8{M6?}yd3wXZkhS2cq8dJeezuK z;v(nt%~YEemJ;&pep^a3?k|>7OZSIiT9Rj)$Q{6_v??vp9VJ#{Mz_PQ^tmMN?nF)8 z)(jKjh!a|h=5bOwb1_}GD1K=ct#{HAD9xl`LKz#CFb1qq@`P>%@_1KOtRv|4{`BC z_vK`5U{1f^4sXB7clkX6FL?2FW&s&iB9>`&uX#U{h#xK|(Yd~|*r!;Zt1NpJ=E<7m z{S_?*JR!}iPwuvrZLtWoIPcgPR@F3zyr8MJN`*CEqmy4%jPzK$om}64PJKi1#jjWT zlk=ZVPX*u`RKV1zk}aed6B@`t5tfTrW0s?KYS8ELP zemptN?EIWsE5)=Zf4oG3P-$PUnKE1ux88Kq_7@zdugN@}<@ffEA?C%nUwwR>$D90K z*rhCIpL6p4hI#&pQ5Ir1`9<}>!fECPCQY3`nSogW0^|My9$@f!I!GNPN*9U7>mYNn zx=1`yUmZ1s4tnfSXl_X`-=cn$Bc6?mMI^o~ZVl6#S`OR_gRcNzQTV?YT}4NM3*hf! zGAv16M|Yq-F;VC(-Ar1?HlW2Cu55|piThOIw1LlBm>ZqE+b6nnwF8<2(;=GI>Y9gJ z@^sj)doGqcSgd%?VDD$X9;vF$*_RDeE;igeUoI=7`&scVs%}~q~SuL{P+u$y#vC}L=WCx zVXiA+VKm_{;EYOiy;miiU3uOhh3otaovWT6x#~*;HajSvS0^#^k#@e$yQIOAEE!`G z_hNXuum(vBc@X0l84xWg`lw8``_knB6^Xlr#^YSpH3jU1 z`K&sgAkUP6R&L>n7InX;EMpXY9-7Rn&3o|)xij0gG&CeyrbMW>Odgw_b?0!I_1mr7 zv>vy16vpE4X1AjuYqtJcKjJ8j3ZPV+E!#wrRlr5^i@g%{3nh|^Y z)_Gj|_5N6qw%B0((A+tlYvUcci#QeAS<&bFRZd#vgdRLO=-J)BZ^B*dc=ql~sS2Fk zA__;FbmQC-oligbx&}KJ$Wv9FSS7!wT=31;0}jWegR1%;P4OQbBnw{)3oKZH*(;!4 z4CXJu4h}dTEGUNKAgm`KmbeBS+(dx|*L(u=eb;~n5Xg*ck59-RQYwC1psSwvYu_(4O#f{XTR$`W>6 z=UPXv&6|%t#4XV2aGmcs-eWQR^8T9RtT_Z6-mHARd(*wklIOhdoAFAFpF160-X~zV z@Xq^o@Rp7P;^X`huY_$4q6@a)BG-wj_8|l_Vyh+5B-z_{UTTUQ+_dMy(D|iL^4X#{ zGjh-E_dnN~RkWd_bNATH_D248;=J8mJlj$n+eeE7YyBT36Uickc~|>@`g?nF?E;HG z;I)M&pLZmZl-CGP?fOuc{WzxGD#f!Zb?|A3K6&$~T(~vg7wb`TLzzmW{m1fphHBYd z$Yl@i$Zu*buJIn3ORMjVPKG1;^lwttjyRHOYsV~4d#q z0WisFrg6uBr2X8_t?7;$J6Er|6r~bspcNxQ_;^PuL2OLQOz^>%HF8HZHl`JP9tfFPEy3Q36|Rq$*I9p2wbUy`Z@Qz^%`@4e zKhnq9bYJZGeL_}}4V^BDZz?KlxNhw$t8VeUVkNPdXRA#*`DW2I;B&5C1#WLLedC(ZtgWo1!MZ)L!Hu77N6(hr|J2fdMsg@K(G5wnb_#-D(#EiJy$@d7q!fc{5$*jr%xzhGq z0kL8K=By5!$#}2=FJ^wvg8cmk%f|zlOal%K{iSMoAK`EcO7T*N+(q^T{TCg?2?-vt z0sn*xLsE+-p=u(M@g`FHG(9p(q&CDomHZ^-Iez(l1JmkHHS4C2Z$5B6=n(tnK8Y83 zk2wd&bL~gcQ~C0$O$S9!RN{=|zRcO5Z?aCy$|m&RQ_MF=C)*}e=a~Q z3%6+6)?#E%WW6QdsNaxvC)+$^-9fZGQA1)NWs}n?jkw9OksD9whNu%goB`pDpVB)m z1y3Yh+Mw~c*VMxK78mwSD3TiXIY1!p{rQz0t&V6|YqE@3nPMqbb-pn%^9ZI4$lqH~_ zJlcd&NE$uOH}{%Py*gZ|UjP^K;%>-V`JqrmmFW71PVUblRCZoO}ISxE3T0w$WLb&L^K;@F!Ct%>vhukbd3JRXs;dMGGJU6+I?h`G$PJEQ8_0s&_ zg^FDwTY8U_oQ`}G@o0K)0iTq8(4mr$db^=`uTmeJUhS#tq2A{wc5x-$Z(-X4dwt|a za&5iA;!wYX#nQ%@=bB#?Vf$aRTcO=_$c0zkBNBs_49`?iSqkmMS2(=Ft^iI^;Oo|Z znhKRj1un(y?BIMeRqeL@Z8sso>I?m69PbVKlr;+&w^7jXURN?XFNl^;_QP5yOu37D z9&DNw;XnN;4Hc3z?)E0DGQ_$(QZ~FaPLYvycD>ihTDA+9@i`e!lIo*hCpv=P0 z%kQnQb2_=>8Oxb>CU-q)77JTqdsLK#%4@f-*_mo1k{?cMyqTNEnoBHt;9vnShrAnsS2F`3-UuJe-lQrY4hdF$m zLk^3b9SKJ?I=X9mki0ubp5mIPrq;XgDc$MRd3+12P?>w&&NU|2a+Tnw&KJ%530$UY zcUB*X)t25{$5kaf`z?>Zgb#k-(J?etxne$U5A9BnFs}y&BW!%3XO6AEDMY3~K8iJA zgQ&_2Kh3SlfxZ>_mlsNcUCzj8Ul2Ep4tjvCYd^q07CCdJ(dMyh?wA@|c2!q)49|t$ zsN@a8+-jeMwbn1D%W2e8Z>{1HIcD6u%Kw!}(~?n{d8wyEnFQO;PqlWJ5!)_Sj_-~0 zN|!XaCfB|2;zf3|nHA@m&nv9Y3SeIPFW{wC1B#;*&aTudcXxUdd_zd@-TV8_lOJt} z3oUx>y0rgNumvGJBs=0sok+1&=ca&A>@9}j9-sZ1Kb}x&$*Do_j=ZteoXLG|>)1@RWamt})Eg;5 zv4Nsq)!Y~?6#&8%|)+sD8v%YeRfyy4ie*A;b7xXzqJ%`GO&~JyY-n-4S(~xHC||YC5w+NF z#{LA}WPh}0Tw91A;YmoVhxB(wN&QcZ&tK^^k7fE3W;Z ziFp^+2J2<-bQc@3v&g%*uFmV8j!<4dT7;N!Kla?q9%bLkw9gxcJzBLzLT>l3Pvv~G z=B`XR=X;r3fv5XQ9Z!BW$x-eKLdu%>oq7n#9No(oc|~|Wh9ZvVN6x?V zxxD%iw_dQ`xk=@OcZD4ZIV_UNhIv*WJoy&+3!{!~K6xn5dS^{6-@0RZUY0{+heW;3 zZ(ER9=$jr5_1tY}^QLb;D6VYMAZ*lR#>{BTjlPGw)^UpCo{DTp6yN`(O&uX?W{|(# z{(<|5bMBZ-%KXH-qaGo(qg8{75wkTV3kr&F6WVF586%e5(Jw}CP(+ghxdh@L`vyoO zPS2g68QhRAw8nUZ<)Pqh>sq1QyMyh4o2+M!Y*Q1k*EvrPNq+kdygSWd}CvAG}%97zS%|_!d(*V&=BPbV6cg z$tm4R)^Upc_$=?&>EXw0wwnF%>S<1@E^cB|aZT@%>9fxQyXJ+n-=AOM!!C9Okc&p+ z{#)o(bTuLX00JK#PX#~gnbOdmi(I68iQBdicI?=jAoB6)D`hEK`Ib5!1YhdERo9;M?7)J<*aYadyBxm7@VRHmQVXZu;%x- z)go%8C7Uamr8wp{81nb{Pv-=@KBQ`m^)j9`-*xI+FyC3hLh5@{wCd|=Jl>oZ4 zR>sT5g(&vRjc#fe#iT5GKI#o4^S(7ffo%GH=y9s23p`h9Z`9HmWGMLi}WgTh&> zg&R%W+^_jhv~NpS(@#)|W)3`E9!%6Nd(pG&{FXx_Z#P6gnYDXLyU?ETa8Ygdq8OXY z=Q0iUhBC~&r6zmV#VO@3&Ae9Kk%s33U8~k^YxL8KeUPBvb|{Yd>Y>1;*NFK$?KRK) zPCm3yYWs8kxD_Dg-}rhM98!l#!Bo^rf3-ebt-Gb|Fi_6xo0={l6=nz~HSCer9EzqK zhHDJ@nUt?{EUqx{3J{O*7iw86LKcKfK+slXtGq)S_$;GvDW&(!&ZEZWIY~|69+LO#)x>>wDkEkD|FnEuV?05K=UVt(dFlwx3E) znG)Wa&%1B7v*wz2)c%RO`QeSI1a^GNu(#vZp#D!&yPFp#vI<7;if^s8L{A?iA}8m*%Hr=&s6bFE%}$Rj43)GJ^TRF7IQaSKGs`e{(1|ES`Rg{TQHsiC1&Y z37scXIh-fjV$$i|?+zV2Q=_I&`;$AUpUuH?)FC635TR({N&KEOQd2p0c#ksC!Qz;g zKGDP822w;Go+KqZ3lH!i3US!tdmfF$!4Oa*Wum8*hrJug)dPL(U9SosmLGrL?#0*Lyk&rNg zfB)p5*g-&?4hUrg;ewDh1#yllE+G033UvH8!16~C454gqV?zXCicnnQ2@rn>M|iqf zSP|hEU>V(rAY2gyFam!L;ReMrIuUJ1-?JX}wss&7H~<*q=?KRHEp6xLW=C{^LGA2a ztwB#t7M^x+gg?>46^?Lq0n;Y=fE>xrg9!3AAQ}?ydU#s7dJy3RK%)hSiUj0=(TE@i&>0xe9^^nnXNYfL zVDN*5A;4IkU~JF<9N`UOpg=Gr!qDEDWCs`pL9L(TDj93{<2&QO!tvAUGM3_hT1v1` zP8PPFfSZ4_s0yA8CKJ$rWM~A4-V+|-x3frKH9eq=!mDu|A?w^DGibs_@7 z5&9&e^D!u+W#P-15gdiVqJIuh{Q)IU!0~`t{~r?U|FsVc5K1f(4aee)v=|9|x5R^Og@furzu=1! z25fmKkH=zx@a#>9bN23!6i4iL$ZAZU76=m7kOJD%7@_>JkRDcBg{tHRP zFG&V`FaIGFkU;#B0AY-0AQ+Hvd`~cjT!51>JRUF($iPa8@Let-nfN&|WCD_l<@8+^ zuz;;#{eMf#YrQNNkZk-&kX$Uw9qy+v(4SGa3LFJNOFc;Af`J&h!@vqH=PclWTa2=Q z6hHcV@Vzx0Y!dJa1{5)X?2Jmw$_xG75;9ah5K6`?hV}=$h4Ekq2g+|b$1ovq#^-I& zl|2xY z%a#vlmF03^3>h8|V73<=pjSrePxg%ga=}Y)u*sm8uw_$Ljm2SPhue@6~=XF+VH}fEfSM1l>laZ?D z6(q-clW(7s0GkD&XixIgA$lmeI=cZZ1Hqyd{Zas_7`H9J5t%2P5a@>;hl{fP;Vb4XwiW$Q)p4AkhFN|6Jd&nhUI}sjTYQRmv;2nrUOoDuX{E zE~6R5J~R?~3FR0p{$FtvXoDXt#mEr;f~jb%3{0Fh&`J(ywgNNtYg>C_MJOgwu_L@6PpC#jRHaqd_uJG z59Bjuw7zC)V<9E+EZJF`+Feb_3?pB5oHJ9VJWZg_P17cZ=#(W|ZzHv#fO6{$*-4A` zd5BJ<(xyjg{VmiN8FVUza_(KK^hx@{SIX1tlqd1@g*n=2AGPKQ z<+?xF_c-m-drIngai#F=>)@onPRc3(@ ze=(AuI(Bf<;J}2Q+_<(ZZKP-P5Me}JdRP_zQ5iR+h^0Kc`9U7@Za<3t?MuH5f;RlN z7b{7vdi)l>=M8fAYbg^6vbSN^?p?ij$#$`wZ97C;MMPWR+na^9Q}S*#ZW9Jug!Xxm z_O6{yUZgy{O8+`dK4nV#_=fCbM2YjEO@AVL9i=>urH%B`7w4(fIkYdI=rkI&>^_}B zrq4}L;{7No=jd~j6hr= zX2=(8Xz$vnrS~Y2?qoL|N6ayI!yw$l7?9MT;PAN zA)rn8hl-G4DC!?7;ve1lkBSJdF=W}Yib?goD7WI~o!H2_Z1FvD8ButyI=e$Ex@p)} zsVGJDdU1Y5@7*vSXGIm&Ts`+UoHA3D^CkyM?9ycG?>VFel*zQKOIP?dsrwXtc^Eax zcPBK;rAYUUx?!#FTXn-$nMU>b*r-W^K@HJ9?m=$RK5MHRQQ5($R&}3T_q$OpCA$6U zhE2S-X|lDvR%rsxx9rjadVL$z4ZD0h)D5e|ozn#DW6^DK*tT%wA>OeNfr>CvHJ?&A zDKX3{oOG49ns-bJ8^Svlfsa79ZNP`}jS1n`@Qp=ah0tw+7ABTyvfVQM>htkA_3HVR z>iJc^1M2fPP@U@Yxvkae-IaE>X^xf+xHd1$eSxtTsD4yipY2n0TcQQUC5>FO@K)WY z*gYs}5;>?T+IOz%y=qaaYSF>)MAf3kDhJWNfT}&JMV#U3szv=(rm97Gs#c<=ri0r= zP2H_-Momf%o)qo-T9vO_G!Vw3S|l94q*^o^P7foQ@s8ODn1+)!gl!HZ&4zDLEou%= zQY{iv)e`M98uSzGdsa07`VH#|Bb^9O4kxjMi>Ma)t15~130HB5_E}oRL``Bu`--a& zsztlPq{B!9eDl0xfw;H4WBhzxVI*6DgM4Gz7$Ia^pM!~W+6khGRa$Xv!GL;x`D?)F zW?u8O;tD(Gv>w;A;<|#@>V_2sI;EYk-G+rd7YHYv4*#lJivVcVWN z?8Ua>FxLde_VEgYk@kcSs}|*lONW!5g%MPX`l?t(`-JYa6j!Qu_ZNT;`o!n<2|N-Q zGv=rBk9F}s2_sR%uY{2_!=@zrZZB%jPeP|GViIhRgps-jp62AwD+l7)rB6msr1fG07h&C|DLxcgB_4^-i zaBu+Q+Q9!QgH9@-Ke)^HG7#|sf%RX@!2Js_-QUY_0E+!qCJp?@Kgy(6s0WTbf2{|e zfN$}883qFK(BD7wAyF7`Z1Y+$Kj3rzUJr$qfsptw{h=@@@CEm`G8F1h9v}%2 z`1!Rh5(Aw1-^$QXq{6RdNE8;r)W5VvqVT{^{Jl&XoFe>QMu2eGuWiv7$hY{d3`0{}=;!&VTfW!9zaDuWbRF|7eSa;=O*W2T9-`WjM%p`L#cA-0(+R918MmeyxYY z;r>`loXnr&f^&pFc#20u?%uEc;jw@64^M!cxL@l@Lx&W Date: Mon, 10 Aug 2026 16:26:10 +0530 Subject: [PATCH 37/41] Rename Forecaster.score to compute_loss_from_forecast --- .../models/forecasters/deterministic.py | 20 ++++++++++--------- neural_lam/models/modules/base.py | 3 ++- neural_lam/models/modules/deterministic.py | 19 ++++++++++-------- tests/test_prediction_model_classes.py | 12 ++++++----- tests/test_probabilistic_forecaster.py | 2 +- 5 files changed, 32 insertions(+), 24 deletions(-) diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py index 706ea7f2..ddb31e1a 100644 --- a/neural_lam/models/forecasters/deterministic.py +++ b/neural_lam/models/forecasters/deterministic.py @@ -22,9 +22,9 @@ class DeterministicForecaster(Forecaster): single forecast. ``compute_training_loss`` produces one forecast and scores it, and - ``score`` applies the same metric to an already-produced forecast for - reporting. ``forward`` is left abstract; see - ``DeterministicARForecaster`` for the auto-regressive combination. + ``compute_loss_from_forecast`` applies the same metric to an + already-produced forecast for reporting. ``forward`` is left abstract; + see ``DeterministicARForecaster`` for the auto-regressive combination. """ def __init__( @@ -45,7 +45,7 @@ def __init__( Configuration used to compute the constant per-variable std substituted for ``pred_std`` when the forecast carries no std of its own. Needed only when ``loss`` is a scoring rule that uses a - std; without it in that case ``score`` and + std; without it in that case ``compute_loss_from_forecast`` and ``compute_training_loss`` raise ``ValueError`` via ``_resolve_pred_std``. Forecasters used purely for inference can always omit it. Default ``None``. @@ -79,7 +79,8 @@ def compute_training_loss( Produces one forecast over every predicted step, scores it against the target states on interior nodes and averages over batch and time. Callers that already hold a forecast should score that one - with ``score`` rather than producing another here. + with ``compute_loss_from_forecast`` rather than producing another + here. Parameters ---------- @@ -124,7 +125,7 @@ def compute_training_loss( prediction, pred_std = self( init_states, forcing_features, target_states ) - step_losses = self.score( + step_losses = self.compute_loss_from_forecast( prediction, target_states, pred_std, @@ -173,7 +174,7 @@ def _resolve_pred_std( ) return self.per_var_std - def score( + def compute_loss_from_forecast( self, prediction: torch.Tensor, target_states: torch.Tensor, @@ -246,7 +247,8 @@ class DeterministicARForecaster(ARForecaster, DeterministicForecaster): Combines the two orthogonal halves: ``ARForecaster`` supplies the auto-regressive ``forward``, ``DeterministicForecaster`` supplies the - single-forecast training objective and the reporting ``score``. + single-forecast training objective and the reporting + ``compute_loss_from_forecast``. """ def __init__( @@ -269,7 +271,7 @@ def __init__( Configuration used to compute the constant per-variable std substituted for ``pred_std`` when ``predictor`` does not output its own. Needed only when ``loss`` is a scoring rule that uses a - std; without it in that case ``score`` and + std; without it in that case ``compute_loss_from_forecast`` and ``compute_training_loss`` raise ``ValueError``. Forecasters used purely for inference (``forward``) can always omit it. loss : str, default "wmse" diff --git a/neural_lam/models/modules/base.py b/neural_lam/models/modules/base.py index c9043393..368c8c41 100644 --- a/neural_lam/models/modules/base.py +++ b/neural_lam/models/modules/base.py @@ -65,7 +65,8 @@ def __init__( forecaster : Forecaster 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 + scoring (``compute_loss_from_forecast``); 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 diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py index 465fc553..e2cbcb88 100644 --- a/neural_lam/models/modules/deterministic.py +++ b/neural_lam/models/modules/deterministic.py @@ -27,14 +27,16 @@ class DeterministicForecasterModule(BaseForecasterModule): an ensemble. Training is shared with that module unchanged (see ``BaseForecasterModule.training_step``). - The reported loss comes from ``forecaster.score``, since only the - forecaster knows its objective. The reported metrics (mse, mae) are + The reported loss comes from ``forecaster.compute_loss_from_forecast``, + since only the forecaster knows its objective. The reported metrics + (mse, mae) are computed here from ``neural_lam.metrics``: they are fixed regardless of what the forecaster trains on, so routing them through it would add a layer without adding meaning. """ - # Narrowed from Forecaster: this module calls forecaster.score() + # Narrowed from Forecaster: this module calls + # forecaster.compute_loss_from_forecast() forecaster: DeterministicForecaster def __init__( @@ -58,9 +60,10 @@ def __init__( Parameters ---------- forecaster : DeterministicForecaster - The forecaster to evaluate. Must supply ``score``, i.e. carry - the deterministic objective, since validation and testing score - a single prediction through it. + The forecaster to evaluate. Must supply + ``compute_loss_from_forecast``, i.e. carry the deterministic + objective, since validation and testing score a single + prediction through it. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore @@ -179,7 +182,7 @@ def _compute_prediction_and_loss( prediction, target_states, pred_std, _ = self.common_step(batch) time_step_loss = torch.mean( - self.forecaster.score( + self.forecaster.compute_loss_from_forecast( prediction, target_states, pred_std, @@ -255,7 +258,7 @@ def test_step(self, batch, batch_idx): ) self.test_metrics[metric_name].append(batch_metric_vals) - spatial_loss = self.forecaster.score( + spatial_loss = self.forecaster.compute_loss_from_forecast( 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 1981c61d..3f5e2818 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -78,7 +78,7 @@ def test_ar_forecaster_unroll(): assert torch.all(prediction[:, :, 1:, :] == 5.0) -def test_ar_forecaster_score(): +def test_ar_forecaster_compute_loss_from_forecast(): datastore = init_datastore_example("mdp") config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( @@ -98,7 +98,9 @@ def test_ar_forecaster_score(): # 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) + scored = forecaster.compute_loss_from_forecast( + prediction, target, None, mask=mask + ) expected = forecaster.loss( prediction, target, forecaster.per_var_std, mask=mask ) @@ -106,7 +108,7 @@ def test_ar_forecaster_score(): # 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( + scored_explicit = forecaster.compute_loss_from_forecast( prediction, target, explicit_std, mask=mask ) expected_explicit = forecaster.loss( @@ -135,7 +137,7 @@ def test_ar_forecaster_without_config_raises_on_use_not_construction(): target = torch.ones(B, num_grid_nodes, d_state) with pytest.raises(ValueError, match="per_var_std fallback"): - forecaster.score(prediction, target, None) + forecaster.compute_loss_from_forecast(prediction, target, None) num_past_forcing_steps = 1 num_future_forcing_steps = 1 @@ -171,7 +173,7 @@ def test_ar_forecaster_without_config_scores_with_unweighted_loss(): prediction = torch.zeros(B, num_grid_nodes, d_state) target = torch.ones(B, num_grid_nodes, d_state) - scored = forecaster.score(prediction, target, None) + scored = forecaster.compute_loss_from_forecast(prediction, target, None) torch.testing.assert_close(scored, torch.full((B,), float(d_state))) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index f6b2d957..15285e35 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -389,7 +389,7 @@ def test_deterministic_training_step_logs_per_step_losses(): init_states, forcing_features, target_states ) step_losses = torch.mean( - forecaster.score( + forecaster.compute_loss_from_forecast( prediction, target_states, pred_std, From 494bf625aac2d72ca8861b24ec9844839eafa9e4 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Mon, 10 Aug 2026 16:28:56 +0530 Subject: [PATCH 38/41] Update neural_lam/models/modules/probabilistic.py Co-authored-by: Joel Oskarsson --- neural_lam/models/modules/probabilistic.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 19a150b7..257757d2 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -218,6 +218,9 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ + # Note that we here do two forward passes: One for computing loss + # and one for computing ensemble metrics. Required as computing loss + # might not involve making a forecast the same way as during inference. self._log_objective(batch) entry_mses = self._ensemble_step(batch) self.val_metrics["ens_mse"].append(entry_mses) From 5814f64edbd3975666ca1338a2d057db498ffc98 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Mon, 10 Aug 2026 16:30:03 +0530 Subject: [PATCH 39/41] Update neural_lam/models/modules/probabilistic.py Co-authored-by: Joel Oskarsson --- neural_lam/models/modules/probabilistic.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 257757d2..8b9d1fcb 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -240,6 +240,9 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ + # Note that we here do two forward passes: One for computing loss + # and one for computing ensemble metrics. Required as computing loss + # might not involve making a forecast the same way as during inference. entry_mses = self._ensemble_step(batch) self.test_metrics["ens_mse"].append(entry_mses) From 8e69ff51e4f41404ac4038c5d9869d1144086fc5 Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Mon, 10 Aug 2026 16:37:02 +0530 Subject: [PATCH 40/41] Note the two forward passes only where they happen --- neural_lam/models/modules/probabilistic.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py index 8b9d1fcb..0270e2e4 100644 --- a/neural_lam/models/modules/probabilistic.py +++ b/neural_lam/models/modules/probabilistic.py @@ -218,8 +218,8 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - # Note that we here do two forward passes: One for computing loss - # and one for computing ensemble metrics. Required as computing loss + # Note that we here do two forward passes: One for computing loss + # and one for computing ensemble metrics. Required as computing loss # might not involve making a forecast the same way as during inference. self._log_objective(batch) entry_mses = self._ensemble_step(batch) @@ -240,9 +240,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - # Note that we here do two forward passes: One for computing loss - # and one for computing ensemble metrics. Required as computing loss - # might not involve making a forecast the same way as during inference. entry_mses = self._ensemble_step(batch) self.test_metrics["ens_mse"].append(entry_mses) From 2ae76d0c475fbbe17d15f2852b4b87b73a92acde Mon Sep 17 00:00:00 2001 From: Jeevant Singh Date: Mon, 10 Aug 2026 18:03:23 +0530 Subject: [PATCH 41/41] Require pred_std in the metrics that score a distribution --- neural_lam/metrics.py | 112 ++++++------------------- tests/test_probabilistic_objectives.py | 4 +- 2 files changed, 27 insertions(+), 89 deletions(-) diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 1296a084..af1e3d1b 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -1,6 +1,7 @@ """Evaluation metrics shared across training and validation routines.""" # Standard library +import inspect from collections.abc import Callable from typing import Optional @@ -35,38 +36,6 @@ def get_metric(metric_name: str) -> Callable[..., torch.Tensor]: return DEFINED_METRICS[metric_name_lower] -def _require_pred_std( - pred_std: Optional[torch.Tensor], metric_name: str -) -> torch.Tensor: - """ - Return ``pred_std``, raising if a std-dependent metric was given none. - - Parameters - ---------- - pred_std : torch.Tensor or None - The standard deviation the metric was called with. - metric_name : str - Name of the calling metric, used in the error message. - - Returns - ------- - torch.Tensor - ``pred_std`` unchanged. - - Raises - ------ - ValueError - If ``pred_std`` is ``None``. - """ - if pred_std is None: - raise ValueError( - f"{metric_name} scores a predicted distribution and so requires " - "pred_std, but got None. Only the unweighted metrics (mse, mae) " - "can be computed without one." - ) - return pred_std - - def mask_and_reduce_metric( metric_entry_vals: torch.Tensor, mask: Optional[torch.Tensor], @@ -120,7 +89,7 @@ def mask_and_reduce_metric( def wmse( pred: torch.Tensor, target: torch.Tensor, - pred_std: Optional[torch.Tensor] = None, + pred_std: torch.Tensor, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -137,10 +106,9 @@ def wmse( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor or None, optional + pred_std : torch.Tensor Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation used as per-entry weight. Required here; ``None`` - raises. Default ``None``. + standard deviation used as per-entry weight. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -155,13 +123,7 @@ def wmse( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. - - Raises - ------ - ValueError - If ``pred_std`` is ``None``. """ - pred_std = _require_pred_std(pred_std, "wmse") entry_mse = torch.nn.functional.mse_loss( pred, target, reduction="none" ) # (..., num_grid_nodes, num_variables) @@ -198,9 +160,9 @@ def mse( Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor or None, optional - Unused. Accepted so that every metric in ``DEFINED_METRICS`` shares - one signature and callers can stay agnostic about which they got. - Default ``None``. + Unused. Accepted so that a caller holding any metric can pass one + the same way; optional here, unlike in the metrics that score a + distribution. Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -228,7 +190,7 @@ def mse( def wmae( pred: torch.Tensor, target: torch.Tensor, - pred_std: Optional[torch.Tensor] = None, + pred_std: torch.Tensor, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -245,10 +207,9 @@ def wmae( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor or None, optional + pred_std : torch.Tensor Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation used as per-entry weight. Required here; ``None`` - raises. Default ``None``. + standard deviation used as per-entry weight. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -263,13 +224,7 @@ def wmae( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. - - Raises - ------ - ValueError - If ``pred_std`` is ``None``. """ - pred_std = _require_pred_std(pred_std, "wmae") entry_mae = torch.nn.functional.l1_loss( pred, target, reduction="none" ) # (..., num_grid_nodes, num_variables) @@ -306,9 +261,9 @@ def mae( Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor or None, optional - Unused. Accepted so that every metric in ``DEFINED_METRICS`` shares - one signature and callers can stay agnostic about which they got. - Default ``None``. + Unused. Accepted so that a caller holding any metric can pass one + the same way; optional here, unlike in the metrics that score a + distribution. Default ``None``. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -336,7 +291,7 @@ def mae( def nll( pred: torch.Tensor, target: torch.Tensor, - pred_std: Optional[torch.Tensor] = None, + pred_std: torch.Tensor, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -353,10 +308,9 @@ def nll( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor or None, optional + pred_std : torch.Tensor Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation of the Gaussian. Required here; ``None`` raises. - Default ``None``. + standard deviation of the Gaussian. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -371,13 +325,7 @@ def nll( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. - - Raises - ------ - ValueError - If ``pred_std`` is ``None``. """ - pred_std = _require_pred_std(pred_std, "nll") # Broadcast pred_std if shaped (num_variables,) via distribution internals dist = torch.distributions.Normal( pred, pred_std @@ -392,7 +340,7 @@ def nll( def crps_gauss( pred: torch.Tensor, target: torch.Tensor, - pred_std: Optional[torch.Tensor] = None, + pred_std: torch.Tensor, mask: Optional[torch.Tensor] = None, average_grid: bool = True, sum_vars: bool = True, @@ -410,10 +358,9 @@ def crps_gauss( target : torch.Tensor Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. - pred_std : torch.Tensor or None, optional + pred_std : torch.Tensor Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted - standard deviation of the Gaussian. Required here; ``None`` raises. - Default ``None``. + standard deviation of the Gaussian. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. @@ -428,13 +375,7 @@ def crps_gauss( Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. - - Raises - ------ - ValueError - If ``pred_std`` is ``None``. """ - pred_std = _require_pred_std(pred_std, "crps_gauss") std_normal = torch.distributions.Normal( torch.zeros((), device=pred.device), torch.ones((), device=pred.device) ) @@ -453,7 +394,7 @@ def crps_gauss( ) -DEFINED_METRICS = { +DEFINED_METRICS: dict[str, Callable[..., torch.Tensor]] = { "mse": mse, "mae": mae, "wmse": wmse, @@ -462,11 +403,6 @@ def crps_gauss( "crps_gauss": crps_gauss, } -# The metrics that weight by, or parameterize a distribution with, pred_std, -# i.e. exactly those calling _require_pred_std. Kept in step with the guards -# by test_pred_std_requirement_matches_declaration. -_STD_DEPENDENT_METRICS = frozenset({wmse, wmae, nll, crps_gauss}) - def requires_pred_std(metric: Callable[..., torch.Tensor]) -> bool: """ @@ -474,6 +410,8 @@ def requires_pred_std(metric: Callable[..., torch.Tensor]) -> bool: Lets a caller holding a metric decide whether it has to come up with a standard deviation at all, rather than assuming every metric uses one. + Read off the signature, so declaring ``pred_std`` without a default is + the only place a metric states that it scores a distribution. Parameters ---------- @@ -484,7 +422,7 @@ def requires_pred_std(metric: Callable[..., torch.Tensor]) -> bool: Returns ------- bool - True if calling ``metric`` without ``pred_std`` raises - ``ValueError``. + True if ``metric`` takes ``pred_std`` as a required argument. """ - return metric in _STD_DEPENDENT_METRICS + pred_std_param = inspect.signature(metric).parameters["pred_std"] + return pred_std_param.default is inspect.Parameter.empty diff --git a/tests/test_probabilistic_objectives.py b/tests/test_probabilistic_objectives.py index 6e4de197..df13d0fd 100644 --- a/tests/test_probabilistic_objectives.py +++ b/tests/test_probabilistic_objectives.py @@ -90,12 +90,12 @@ def test_unweighted_metrics_ignore_pred_std(metric): @pytest.mark.parametrize("metric_name", sorted(DEFINED_METRICS)) def test_pred_std_requirement_matches_declaration(metric_name): - """Every metric raises without a pred_std iff it declares needing one.""" + """A metric is callable without a pred_std iff it declares one optional.""" metric = get_metric(metric_name) pred, target = _single_residual_case() if requires_pred_std(metric): - with pytest.raises(ValueError, match="requires pred_std"): + with pytest.raises(TypeError, match="pred_std"): metric(pred, target) else: metric(pred, target)