diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ff85253..37c1784e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,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/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/metrics.py b/neural_lam/metrics.py index 1eb1d526..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 @@ -141,7 +142,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 +159,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 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. @@ -177,9 +178,12 @@ 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 ) @@ -239,7 +243,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 +260,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 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. @@ -275,9 +279,12 @@ 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 ) @@ -387,7 +394,7 @@ def crps_gauss( ) -DEFINED_METRICS = { +DEFINED_METRICS: dict[str, Callable[..., torch.Tensor]] = { "mse": mse, "mae": mae, "wmse": wmse, @@ -395,3 +402,27 @@ def crps_gauss( "nll": nll, "crps_gauss": 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. + Read off the signature, so declaring ``pred_std`` without a default is + the only place a metric states that it scores a distribution. + + Parameters + ---------- + metric : callable + A metric from ``DEFINED_METRICS``, e.g. as returned by + ``get_metric``. + + Returns + ------- + bool + True if ``metric`` takes ``pred_std`` as a required argument. + """ + pred_std_param = inspect.signature(metric).parameters["pred_std"] + return pred_std_param.default is inspect.Parameter.empty diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index e986a872..4c997040 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,7 +3,17 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster -from .module import ForecasterModule +from .forecasters.deterministic import ( + DeterministicARForecaster, + DeterministicForecaster, +) +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_efm import GraphEFM, GraphEFMMultiScale diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 7ea9f6fd..9590b33f 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -5,3 +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 a8135f62..2aaad08a 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 Any # Third-party import torch @@ -13,12 +14,21 @@ 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 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__( - self, predictor: StepPredictor, datastore: BaseDatastore + self, + predictor: StepPredictor, + datastore: BaseDatastore, + **kwargs: Any, ) -> None: """ Initialize the ARForecaster. @@ -29,8 +39,11 @@ def __init__( The predictor to use for each step. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. + **kwargs : Any + Constructor arguments of the objective class this mix-in is + combined with, forwarded to it unchanged. """ - super().__init__() + super().__init__(datastore=datastore, **kwargs) self.predictor = predictor # Register boundary/interior masks on the forecaster, not the predictor @@ -103,7 +116,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,7 +150,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: diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4d957916..30f3e46b 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -7,14 +7,34 @@ 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 combine a way of producing forecasts (the + ``ARForecaster`` mix-in) with a training objective + (``DeterministicForecaster``). """ + def __init__(self, datastore: BaseDatastore) -> None: + """ + Initialize the forecaster. + + Parameters + ---------- + datastore : BaseDatastore + The datastore this forecaster is built for, providing grid + metadata, boundary masks and standardization statistics. + """ + super().__init__() + self.datastore = datastore + @property @abstractmethod def predicts_std(self) -> bool: @@ -75,7 +95,66 @@ 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``, substituting a + fallback std is left to whatever consumes the forecast, not to + the caller of ``forward``. 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. 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 + ---------- + 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. """ diff --git a/neural_lam/models/forecasters/deterministic.py b/neural_lam/models/forecasters/deterministic.py new file mode 100644 index 00000000..ddb31e1a --- /dev/null +++ b/neural_lam/models/forecasters/deterministic.py @@ -0,0 +1,286 @@ +"""Forecasters trained by scoring a single deterministic forecast.""" + +# Standard library +from typing import 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. + + ``compute_training_loss`` produces one forecast and scores it, and + ``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__( + self, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + ) -> None: + """ + Set up the scoring rule and the constant ``pred_std`` fallback. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing the state standardization statistics + 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 + its own. Needed only when ``loss`` is a scoring rule that uses a + 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``. + loss : str, optional + The scoring rule (from ``neural_lam.metrics``) applied by + ``compute_training_loss``, stored as ``self.loss``. Default + ``"wmse"``. + """ + super().__init__(datastore=datastore) + self.loss = metrics.get_metric(loss) + + # 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 + else None + ) + self.register_buffer("per_var_std", per_var_std, persistent=False) + + 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 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 ``compute_loss_from_forecast`` rather than producing another + here. + + 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 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). + 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 ``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( + init_states, forcing_features, target_states + ) + step_losses = self.compute_loss_from_forecast( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ) + return torch.mean(step_losses), {} + + def _resolve_pred_std( + self, pred_std: Optional[torch.Tensor] + ) -> Optional[torch.Tensor]: + """ + Return the std ``self.loss`` should be applied with. + + Parameters + ---------- + pred_std : torch.Tensor or None + Predicted standard deviation as returned by ``forward``, + possibly ``None``. + + Returns + ------- + 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``, ``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: 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 + + def compute_loss_from_forecast( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + 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`` 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 + below, so callers compute those directly from + ``neural_lam.metrics``. + + 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 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 + Forwarded to ``self.loss``. Default ``True``. + sum_vars : bool, optional + Forwarded to ``self.loss``. Default ``True``. + + Returns + ------- + torch.Tensor + The scoring rule's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + + Raises + ------ + ValueError + 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( + prediction, + target_states, + pred_std, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) + + +class DeterministicARForecaster(ARForecaster, DeterministicForecaster): + """ + 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 + ``compute_loss_from_forecast``. + """ + + 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. Needed only when ``loss`` is a scoring rule that uses a + 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" + The scoring rule (from ``neural_lam.metrics``) applied by + ``compute_training_loss``. + """ + super().__init__( + predictor=predictor, + datastore=datastore, + config=config, + loss=loss, + ) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..04ab75d9 --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,183 @@ +"""Forecasters producing probabilistic (ensemble) forecasts.""" + +# Standard library +from abc import abstractmethod + +# Third-party +import torch + +# Local +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. + + 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 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. + + 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 + 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 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 72% rename from neural_lam/models/module.py rename to neural_lam/models/modules/base.py index b097d417..368c8c41 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,21 @@ 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 (``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 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 +83,11 @@ 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 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 @@ -80,21 +96,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) @@ -122,12 +137,36 @@ 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"]) + 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 self.matched_metrics: set = set() @@ -145,29 +184,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 +229,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 +237,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 +347,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 +371,19 @@ 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. + + 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 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 + ``DeterministicForecasterModule.training_step``. + Parameters ---------- batch : tuple @@ -381,12 +394,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 +459,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 +493,21 @@ 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 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 ---------- batch : tuple @@ -526,22 +515,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 +538,20 @@ 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 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 ---------- batch : tuple @@ -577,65 +559,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 +847,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 +873,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..e2cbcb88 --- /dev/null +++ b/neural_lam/models/modules/deterministic.py @@ -0,0 +1,365 @@ +"""Lightning module evaluating forecasters through a single deterministic +forecast 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 ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ..forecasters.deterministic import DeterministicForecaster +from .base import BaseForecasterModule + + +class DeterministicForecasterModule(BaseForecasterModule): + """ + Lightning module for a single deterministic forecast per batch. + + 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.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.compute_loss_from_forecast() + forecaster: DeterministicForecaster + + 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 + ---------- + forecaster : DeterministicForecaster + 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 + 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__( + 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": [], + } + 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 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 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 + ---------- + batch : tuple + The batch of data. + + Returns + ------- + 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) + self._log_step_loss( + time_step_loss, batch_loss, "train", batch[0].shape[0] + ) + return batch_loss + + 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. + + Shared by ``training_step``, ``validation_step`` and ``test_step``. + 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. + + 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) + + time_step_loss = torch.mean( + self.forecaster.compute_loss_from_forecast( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + 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) + + # Reported independently of the training objective, so computed here + # rather than through the forecaster + entry_mses = metrics.mse( + prediction, + target_states, + 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, time_step_loss = ( + self._compute_prediction_and_loss(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) + + 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) + + # Reported independently of the training objective, so computed here + # rather than through the forecaster + for metric_name in ("mse", "mae"): + batch_metric_vals = metrics.get_metric(metric_name)( + prediction, + target_states, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.test_metrics[metric_name].append(batch_metric_vals) + + spatial_loss = self.forecaster.compute_loss_from_forecast( + 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..0270e2e4 --- /dev/null +++ b/neural_lam/models/modules/probabilistic.py @@ -0,0 +1,270 @@ +"""Lightning module evaluating probabilistic forecasters as ensembles.""" + +# Standard library +import warnings + +# Local +from ... import metrics +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +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 its mean scored per lead time and + 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 + forecaster: ProbabilisticForecaster + + 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 + ---------- + 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. 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 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 + ``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__( + 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, " + 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": []} + + 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; ``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. The + objective itself is logged by ``_log_objective``. + + Parameters + ---------- + batch : tuple + The batch of data. + + 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) + + entry_mses = metrics.mse( + ensemble_mean, + target_states, + mask=self.interior_mask_bool, + sum_vars=False, + ) # (B, pred_steps, num_state_vars) + + return entry_mses + + def _log_objective(self, batch) -> None: + """ + Log the forecaster's own training objective as ``val_mean_loss``. + + 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. + """ + 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"val_{name}": value for name, value in loss_components.items() + } + log_dict["val_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. + + 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 + ---------- + batch : tuple + The batch of data. + 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) + + 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) + self.test_metrics["ens_mse"].append(entry_mses) + + def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + + 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.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 0bcb706f..79d55f10 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, ForecasterModule +from .models import ( + MODELS, + DeterministicARForecaster, + DeterministicForecasterModule, +) from .weather_dataset import WeatherDataModule @@ -40,8 +44,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 +67,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 = DeterministicARForecaster( + predictor, datastore, config=config, loss=args.loss + ) + return DeterministicForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, @@ -460,7 +466,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, @@ -479,13 +485,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 = DeterministicARForecaster( + 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..40ee58f5 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 ( + DeterministicARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from tests.dummy_datastore import DummyDatastore @@ -50,12 +54,13 @@ 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 = DeterministicARForecaster( + 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..91788d2b 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 @@ -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,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 = DeterministicARForecaster( + 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..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) + 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 8d516bfb..4e0b8c31 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 ( + DeterministicARForecaster, + 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 = DeterministicARForecaster(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..f6e3d495 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 ( + DeterministicARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -450,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"] @@ -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 = DeterministicARForecaster( + 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, @@ -526,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 @@ -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,13 @@ 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 = DeterministicARForecaster( + 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..3f5e2818 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -2,12 +2,19 @@ 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.models import ( + ARForecaster, + DeterministicARForecaster, + DeterministicForecasterModule, + Forecaster, + StepPredictor, +) from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -40,7 +47,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) @@ -71,6 +78,157 @@ def test_ar_forecaster_unroll(): assert torch.all(prediction[:, :, 1:, :] == 5.0) +def test_ar_forecaster_compute_loss_from_forecast(): + 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 = 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") + 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.compute_loss_from_forecast( + prediction, target, None, mask=mask + ) + expected = forecaster.loss( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored, expected) + + # 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.compute_loss_from_forecast( + 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 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 = DeterministicARForecaster(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.compute_loss_from_forecast(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_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.compute_loss_from_forecast(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. + + Stands in for a future objective (CRPS, a variational loss, ...) to + check that adding one requires nothing beyond declaring its arguments. + """ + + def __init__(self, datastore, scale: float = 1.0): + super().__init__(datastore=datastore) + 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(ARForecaster, MeanAbsObjective): + """Auto-regressive forecast production plus the mean-abs objective.""" + + +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) + + forecaster = MeanAbsARForecaster( + predictor=predictor, datastore=datastore, scale=2.0 + ) + + assert forecaster.scale == 2.0 + assert forecaster.predictor is predictor + assert forecaster.boundary_mask.shape[1] == predictor.num_grid_nodes + assert forecaster.datastore is datastore + + +def test_unclaimed_constructor_argument_raises(): + """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) + + 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") @@ -81,7 +239,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 +255,14 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = DeterministicARForecaster( + 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 +292,12 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = DeterministicARForecaster( + 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 +354,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 = DeterministicARForecaster( + 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 +431,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 = DeterministicARForecaster( + 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 +447,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 @@ -355,7 +519,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") @@ -389,8 +553,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 new file mode 100644 index 00000000..15285e35 --- /dev/null +++ b/tests/test_probabilistic_forecaster.py @@ -0,0 +1,590 @@ +# 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.loss_weighting import get_per_var_std +from neural_lam.models import ( + BaseForecasterModule, + DeterministicARForecaster, + 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`` 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, + 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( + 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 = DeterministicARForecaster(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) + + 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, + 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_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. + + 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) + + 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, + ) + + 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) + + +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 + prediction, pred_std = forecaster( + init_states, forcing_features, target_states + ) + step_losses = torch.mean( + forecaster.compute_loss_from_forecast( + 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]) + torch.testing.assert_close(batch_loss, torch.mean(step_losses)) + + +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) + + +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) + + 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)) + + +def test_probabilistic_module_logs_forecaster_objective(): + """Validation reports the forecaster's own training objective, giving + ModelCheckpoint a val_mean_loss to monitor. Testing does not: nothing + monitors it there, so the extra forward pass would buy nothing.""" + 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 = ComponentReportingForecaster( + predictor, datastore, config=config, train_num_members=2 + ) + 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 = ( + init_states, + target_states, + forcing_features, + torch.zeros(B, pred_steps), + ) + + captured = {} + model.log_dict = lambda log_dict, **kwargs: captured.update(log_dict) + + torch.manual_seed(42) + model.validation_step(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["val_mean_loss"] is forecaster.last_batch_loss + assert captured["val_kl"] is forecaster.last_components["kl"] + + # 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] + + # Testing draws only the evaluation ensemble, and logs no loss + captured.clear() + forecaster.num_members_seen.clear() + model.test_step(batch, 0) + + assert forecaster.num_members_seen == [3] + assert captured == {} diff --git a/tests/test_probabilistic_objectives.py b/tests/test_probabilistic_objectives.py index 49075f08..df13d0fd 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): + """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(TypeError, match="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]]]], diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index eb4c2bc0..9642acfe 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 @@ -82,9 +83,10 @@ 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.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..e660eef4 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,9 +105,9 @@ def run_simple_training( ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # 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,13 +123,14 @@ 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 = DeterministicARForecaster( + 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 +177,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 +208,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)