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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 62dcb764..2a0c35ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add `--train_steps_to_log` CLI option to log training loss for individual unroll steps, and deduplicate common prediction and loss computation steps across loops [\#674](https://github.com/mllam/neural-lam/issues/674) @GiGiKoneti +- 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 + - 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..be27cfec 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,7 +3,13 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster -from .module import ForecasterModule +from .forecasters.probabilistic import ( + ProbabilisticARForecaster, + ProbabilisticForecaster, +) +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/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..e7add53c 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,12 +1,16 @@ """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 @@ -18,7 +22,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. @@ -29,6 +37,17 @@ 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 @@ -44,6 +63,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: """ @@ -103,7 +147,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``. """ @@ -137,10 +181,179 @@ 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: 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 4d957916..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,7 +76,116 @@ 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 - ``ForecasterModule``. 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 + 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. + + 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. + + 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. + 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 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. + """ + + @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/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..57526d1c --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,230 @@ +"""Forecasters producing probabilistic (ensemble) forecasts.""" + +# Standard library +from abc import abstractmethod + +# Third-party +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 + + +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. + + 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 + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int, + ) -> 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 + Number of ensemble members ``S`` to sample. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + 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``. + """ + + +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. + + ``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__( + 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, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int, + ) -> 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 + Number of ensemble members ``S`` to sample. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + 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``. + + Raises + ------ + ValueError + If ``num_members`` is less than 1. + """ + if num_members < 1: + raise ValueError( + f"num_members must be at least 1, got {num_members}" + ) + + 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) + # 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, + 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/__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/module.py b/neural_lam/models/modules/base.py similarity index 73% rename from neural_lam/models/module.py rename to neural_lam/models/modules/base.py index b097d417..94823c6f 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/modules/base.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,18 +18,26 @@ from neural_lam.utils import get_integer_time # Local -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 +from ... import vis +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ...weather_dataset import WeatherDataset +from ..forecasters.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 @@ -38,7 +47,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, @@ -50,18 +58,20 @@ def __init__( args=None, ): """ - Initialize the ForecasterModule. + Initialize the BaseForecasterModule. Parameters ---------- forecaster : Forecaster - The forecaster model to use for predictions. + 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 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 @@ -72,6 +82,10 @@ def __init__( Whether to create GIFs of example predictions. 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). 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 @@ -80,21 +94,20 @@ 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`` - correctly. + ``train_steps_to_log``, ``metrics_watch``, + ``var_leads_metrics_watch``) so legacy checkpoints round-trip + through ``load_from_checkpoint`` correctly. """ super().__init__() # Pre-refactor ``ARModel`` checkpoints saved every hyperparameter nested # 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) @@ -145,29 +158,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. @@ -213,19 +203,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": [], - } - 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 @@ -234,9 +211,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 steps_to_log exceeds the actual rollout self._steps_warn_issued = { "val": False, @@ -347,7 +321,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 ---------- @@ -371,6 +345,15 @@ def training_step(self, batch): """ Perform a single training step. + 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. + Parameters ---------- batch : tuple @@ -381,12 +364,26 @@ def training_step(self, batch): torch.Tensor The computed loss for the training step. """ - _, _, _, time_step_loss = self._compute_prediction_and_loss(batch) - batch_loss = torch.mean(time_step_loss) - batch_size = batch[0].shape[0] - - self._log_step_loss(time_step_loss, batch_loss, "train", batch_size) + 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"train_{name}": value for name, value in loss_components.items() + } + log_dict["train_loss"] = batch_loss + self.log_dict( + log_dict, + prog_bar=True, + on_step=True, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) return batch_loss def all_gather_cat(self, tensor_to_gather): @@ -432,55 +429,6 @@ def _warn_skipped_steps(self, pred_steps: int, phase: str) -> None: ) self._steps_warn_issued[phase] = True - 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]: - """ - Compute predicted mean, standard deviation, and step-wise loss. - Also extract and return corresponding target from batch. - - Parameters - ---------- - batch : tuple of torch.Tensor - The batch of data. - - 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) - if pred_std is None: - pred_std = self.per_var_std - assert pred_std is not None - - time_step_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - return ( - prediction, - target_states, - pred_std, - time_step_loss, - ) - def _log_step_loss( self, time_step_loss: torch.Tensor, @@ -515,10 +463,15 @@ def _log_step_loss( batch_size=batch_size, ) + @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 @@ -526,22 +479,6 @@ def validation_step(self, batch, batch_idx): 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) - - 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): """ @@ -565,11 +502,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 @@ -577,65 +518,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, time_step_loss = ( - self._compute_prediction_and_loss(batch) - ) - - if self.forecaster.predicts_std: - mean_pred_std = torch.mean( - pred_std[..., self.interior_mask_bool, :], dim=-2 - ) - self.test_metrics["output_std"].append(mean_pred_std) - - 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) - - hparams = self.hparams - val_steps_to_log = ( - hparams.val_steps_to_log # type: ignore[attr-defined] - ) - - 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.loss( - prediction, target_states, pred_std, average_grid=False - ) - log_spatial_losses = spatial_loss[ - :, - [ - step - 1 - for step in 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): """ @@ -924,82 +806,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): """ @@ -1016,15 +832,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/modules/deterministic.py b/neural_lam/models/modules/deterministic.py new file mode 100644 index 00000000..817abdc9 --- /dev/null +++ b/neural_lam/models/modules/deterministic.py @@ -0,0 +1,241 @@ +"""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 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 via ``forecaster.score``, 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) + + 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, "val", batch_size) + + entry_mses = self.forecaster.score( + prediction, + target_states, + pred_std, + metric=metrics.mse, + 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) + + 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) + + for metric_name in ("mse", "mae"): + 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.score( + 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/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py new file mode 100644 index 00000000..83c7d50f --- /dev/null +++ b/neural_lam/models/modules/probabilistic.py @@ -0,0 +1,192 @@ +"""Lightning module evaluating probabilistic forecasters as ensembles.""" + +# Standard library +import warnings + +# Third-party +import torch + +# Local +from ... import metrics +from ..forecasters.probabilistic import ProbabilisticForecaster +from .base import BaseForecasterModule + + +class ProbabilisticForecasterModule(BaseForecasterModule): + """ + Lightning module for forecasters that sample ensemble forecasts. + + 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 + forecaster: ProbabilisticForecaster + + def __init__(self, *args, eval_ensemble_size: int, **kwargs): + """ + Initialize the module and store the evaluation ensemble size. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + eval_ensemble_size : int + Number of ensemble members sampled during validation and + testing. + **kwargs + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` + (``lr``, ...). + """ + super().__init__(*args, **kwargs) + if 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: dict[str, list] = {"ens_mse": []} + self.test_metrics: dict[str, list] = {"ens_mse": []} + + def _ensemble_step(self, batch, phase: str): + """ + Sample an ensemble and score its mean against the target states. + + 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. + + Parameters + ---------- + batch : tuple + The batch of data. + 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( + 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_steps(len(time_step_rmse), phase) + + 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, + ) + + 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): + """ + 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 + ---------- + batch : tuple + 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) + + 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. + """ + 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/neural_lam/train_model.py b/neural_lam/train_model.py index 2a9775c9..26e70d94 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 @@ -63,8 +63,10 @@ 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) - return ForecasterModule.load_from_checkpoint( + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) + return DeterministicForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, @@ -455,7 +457,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, @@ -474,13 +476,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( + model = DeterministicForecasterModule( 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/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 2e5f3148..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 @@ -50,12 +54,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) - model = ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + model = DeterministicForecasterModule( 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..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 @@ -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( + model = DeterministicForecasterModule( 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..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 @@ -26,8 +31,8 @@ def _build_module(datastore): ) ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) - return ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config) + 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 616d563d..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 @@ -467,13 +471,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( + model = DeterministicForecasterModule( 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, @@ -665,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", @@ -679,12 +684,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) - return ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + return DeterministicForecasterModule( 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..85bb06d1 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -2,12 +2,18 @@ from argparse import Namespace # Third-party +import pytest import pytorch_lightning as pl import torch # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam import metrics +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -71,6 +77,92 @@ 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_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") @@ -81,7 +173,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 @@ -97,13 +189,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( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1e-3, restore_opt=False, n_example_pred=1, @@ -133,10 +224,12 @@ 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( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -193,21 +286,22 @@ 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. + # to DeterministicForecasterModule's defaults during load. saved_loss = "mse" saved_lr = 0.123 saved_create_gif = True saved_val_steps = [2] saved_n_example_pred = 7 - model = ForecasterModule( + forecaster = ARForecaster( + predictor, datastore, config=config, loss=saved_loss + ) + + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=saved_loss, lr=saved_lr, restore_opt=False, n_example_pred=saved_n_example_pred, @@ -269,10 +363,12 @@ 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( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -283,8 +379,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. - assert loaded_model.hparams.loss == saved_loss + # 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 new file mode 100644 index 00000000..d98cc345 --- /dev/null +++ b/tests/test_probabilistic_forecaster.py @@ -0,0 +1,362 @@ +# 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, + DeterministicForecasterModule, + 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 + + +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. + """ + + 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, + forcing_features, + target_states, + interior_mask_bool, + ): + ensemble, per_member_std = self.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=self.train_num_members, + ) + 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 + 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, 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] + forecaster.per_var_std = torch.ones(d_state) + + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + prediction, _ = forecaster(init_states, forcing_features, target_states) + expected_loss = torch.mean( + score_metric( + prediction, + target_states, + forecaster.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 = ConcreteProbabilisticARForecaster(predictor, datastore) + + # 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, per_member_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 per_member_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) + + +def test_probabilistic_training_loss_gradient_flow(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster( + predictor, datastore, loss="mse", train_num_members=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] + 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, + interior_mask_bool=interior_mask_bool, + ) + + 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_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="num_members"): + forecaster.sample_ensemble( + init_states, forcing_features, target_states, num_members=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) + + +def test_module_training_step_delegates_to_forecaster(): + 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 = ARForecaster(predictor, datastore, config=config, loss="mse") + model = DeterministicForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + ) + + 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, + interior_mask_bool=model.interior_mask_bool, + ) + + torch.testing.assert_close(batch_loss, expected_loss) + + +class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster): + """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) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, config=config + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + 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) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = ConcreteProbabilisticARForecaster( + predictor, datastore, config=config + ) + + with pytest.raises(ValueError, match="eval_ensemble_size"): + ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=0, + ) + + +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 = MemberCountRecordingForecaster( + predictor, datastore, config=config + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + 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.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)) diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index eb4c2bc0..270c04cf 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -50,7 +50,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 @@ -84,7 +85,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.modules.deterministic." + "DeterministicForecasterModule.__init__", capture_init, ), pytest.raises(SystemExit), @@ -93,7 +95,7 @@ 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 assert ( "train_steps_to_log" in captured_kwargs diff --git a/tests/test_training.py b/tests/test_training.py index bf1a5884..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 @@ -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( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, @@ -176,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) @@ -207,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)