diff --git a/.gitignore b/.gitignore index a2c5bb26..37242198 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,7 @@ tags # Coc configuration directory .vim .vscode +.idea # macos .DS_Store diff --git a/CHANGELOG.md b/CHANGELOG.md index 62dcb764..2a0c35ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add `--train_steps_to_log` CLI option to log training loss for individual unroll steps, and deduplicate common prediction and loss computation steps across loops [\#674](https://github.com/mllam/neural-lam/issues/674) @GiGiKoneti +- Add a general probabilistic forecasting interface (`ProbabilisticForecaster`, + `ProbabilisticARForecaster`, `ProbabilisticForecasterModule`) and move + ownership of the training objective, scoring rule and per-variable std + from `ForecasterModule` onto the `Forecaster`. + [\#685](https://github.com/mllam/neural-lam/issues/685) + @Sir-Sloth-The-Lazy + - Add `PropagationNet` GNN layer that incentivises directional message propagation from sender to receiver nodes, and expose it alongside `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`, diff --git a/neural_lam/config.py b/neural_lam/config.py index 1da43fff..37f39714 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -3,7 +3,7 @@ # Standard library import dataclasses from pathlib import Path -from typing import Dict, Union +from typing import Dict, Optional, Union # Third-party import dataclass_wizard @@ -116,6 +116,49 @@ class TrainingConfig: ) +@dataclasses.dataclass +class ProbabilisticConfig: + """ + Configuration for the probabilistic (Graph-EFM) models. + + Only used when training or evaluating a probabilistic model + (``--model graph_efm``/``graph_efm_ms``); ignored otherwise. Every field + has a default, so the ``probabilistic`` config section may be omitted + entirely for deterministic models. + + Attributes + ---------- + latent_dim : int, optional + Dimensionality of the latent variable at each latent-carrying mesh + node. Defaults to the model's ``hidden_dim`` when None. + prior_layers : int + Number of on-mesh GNN layers in the prior. + encoder_layers : int + Number of on-mesh GNN layers in the variational encoder. + decoder_layers : int + Number of on-mesh GNN layers in the latent decoder. + learn_prior : bool + If True, the prior is a learned encoder conditioned on the previous + state; if False, a constant ``Normal(0, 1)`` prior is used. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. + kl_beta : float + Weight of the KL term in the ELBO. When 0, the prior and KL are not + computed (pure auto-encoder training). + eval_ensemble_size : int + Number of ensemble members sampled during validation and testing. + """ + + latent_dim: Optional[int] = None + prior_layers: int = 2 + encoder_layers: int = 2 + decoder_layers: int = 4 + learn_prior: bool = True + prior_dist: str = "isotropic" + kl_beta: float = 1.0 + eval_ensemble_size: int = 5 + + @dataclasses.dataclass class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): """ @@ -133,10 +176,16 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): training : TrainingConfig Configuration for training the model, including loss function and feature-weighting strategy. Defaults to ``TrainingConfig()``. + probabilistic : ProbabilisticConfig + Configuration for the probabilistic (Graph-EFM) models. Defaults to + ``ProbabilisticConfig()`` and is ignored by deterministic models. """ datastore: DatastoreSelection training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig) + probabilistic: ProbabilisticConfig = dataclasses.field( + default_factory=ProbabilisticConfig + ) class _(dataclass_wizard.JSONWizard.Meta): """ diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 7cad45d7..24a508c9 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -18,7 +18,7 @@ from numpy import ndarray # Local -from ..utils import log_on_rank_zero +from ..utils.logging import log_on_rank_zero from .base import BaseRegularGridDatastore, CartesianGridShape diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..dadcee46 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,16 +3,32 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster -from .module import ForecasterModule +from .forecasters.graph_efm import GraphEFMForecaster +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 from .step_predictors.graph.graph_lam import GraphLAM from .step_predictors.graph.hi_lam import HiLAM from .step_predictors.graph.hi_lam_parallel import HiLAMParallel from .step_predictors.graph.hierarchical import BaseHiGraphModel +# Graph-EFM models are probabilistic: train_model.py builds them with a +# config-aware, probabilistic assembly path (GraphEFMForecaster wrapped in a +# ProbabilisticForecasterModule), distinct from the deterministic models +# above. ``PROBABILISTIC_MODELS`` marks which entries take that path. MODELS = { "graph_lam": GraphLAM, "hi_lam": HiLAM, "hi_lam_parallel": HiLAMParallel, + "graph_efm": GraphEFM, + "graph_efm_ms": GraphEFMMultiScale, } + +PROBABILISTIC_MODELS = {"graph_efm", "graph_efm_ms"} diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 7ea9f6fd..3c2093de 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 .graph_efm import GraphEFMForecaster +from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index a8135f62..e7add53c 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,12 +1,16 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" # Standard library +from typing import Callable, Optional # Third-party import torch # Local +from ... import metrics +from ...config import NeuralLAMConfig from ...datastore import BaseDatastore +from ...loss_weighting import get_state_feature_weighting from ..step_predictors.base import StepPredictor from .base import Forecaster @@ -18,7 +22,11 @@ class ARForecaster(Forecaster): """ def __init__( - self, predictor: StepPredictor, datastore: BaseDatastore + self, + predictor: StepPredictor, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", ) -> None: """ Initialize the ARForecaster. @@ -29,6 +37,17 @@ def __init__( The predictor to use for each step. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. """ super().__init__() self.predictor = predictor @@ -44,6 +63,31 @@ def __init__( "interior_mask", 1.0 - self.boundary_mask, persistent=False ) + self.loss = metrics.get_metric(loss) + + # Store per_var_std here if the predictor does not output its own std + if not self.predicts_std and config is not None: + da_state_stats = datastore.get_standardization_dataarray( + category="state" + ) + state_feature_weights = get_state_feature_weighting( + config=config, datastore=datastore + ) + diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + feature_weights_t = torch.tensor( + state_feature_weights, dtype=torch.float32 + ) + self.register_buffer( + "per_var_std", + diff_std / torch.sqrt(feature_weights_t), + persistent=False, + ) + else: + self.per_var_std = None + @property def predicts_std(self) -> bool: """ @@ -103,7 +147,7 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the wrapped predictor outputs an std, otherwise ``None`` (in which - case ``ForecasterModule`` substitutes the constant + case ``DeterministicForecasterModule`` substitutes the constant per-variable std). Dims: same as ``prediction``. """ @@ -137,10 +181,179 @@ def forward( prediction = torch.stack(prediction_list, dim=1) # If predictor outputs std, stack it; otherwise return None so - # ForecasterModule can substitute the constant per_var_std + # callers can substitute the constant per_var_std if pred_std_list: pred_std = torch.stack(pred_std_list, dim=1) else: pred_std = None return prediction, pred_std + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score the deterministic rollout with ``self.loss``. + + Unrolls a single forecast over the full rollout, scores it against + the target states on interior nodes and averages over batch and + time. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during the rollout. + Dims: same as the prediction. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior + nodes are scored. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the rollout, averaged over + batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; the deterministic objective has no separate components. + + Raises + ------ + ValueError + If the predictor does not output its own std and no + ``per_var_std`` fallback is available; see + ``_resolve_pred_std``. + """ + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + pred_std = self._resolve_pred_std(pred_std) + + batch_loss = torch.mean( + self.loss( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} + + def _resolve_pred_std( + self, pred_std: Optional[torch.Tensor] + ) -> torch.Tensor: + """ + Return ``pred_std``, or the constant ``per_var_std`` fallback. + + Parameters + ---------- + pred_std : torch.Tensor or None + Predicted standard deviation as returned by ``forward``, + possibly ``None``. + + Returns + ------- + torch.Tensor + ``pred_std`` unchanged when given; otherwise ``self.per_var_std``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available (``predictor.predicts_std`` is False and this + forecaster was constructed without ``config``). + """ + if pred_std is not None: + return pred_std + if self.per_var_std is None: + raise ValueError( + "No pred_std available for scoring: predictor.predicts_std " + "is False and this forecaster has no per_var_std fallback " + "(it was constructed without config). Pass config to the " + "constructor, or use a predictor that outputs its own std." + ) + return self.per_var_std + + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Resolves ``pred_std`` via ``_resolve_pred_std`` (substituting + ``self.per_var_std`` when ``None``), then applies ``metric`` + (defaulting to ``self.loss``, the configured scoring rule). + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``; ``None`` when + the wrapped predictor does not output one, in which case + ``self.per_var_std`` is substituted (see ``_resolve_pred_std`` + for when this raises instead). + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to ``self.loss``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available; see ``_resolve_pred_std``. + """ + pred_std = self._resolve_pred_std(pred_std) + metric_fn = self.loss if metric is None else metric + return metric_fn( + prediction, + target_states, + pred_std, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4d957916..da63869b 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,6 +2,7 @@ # Standard library from abc import ABC, abstractmethod +from typing import Callable, Optional # Third-party import torch @@ -75,7 +76,116 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when ``predicts_std`` is True, otherwise ``None``. Per-feature - predicted standard deviation; when ``None``, the constant - per-variable std is substituted upstream by - ``ForecasterModule``. Dims: same as ``prediction``. + predicted standard deviation; when ``None``, the forecaster's + own constant per-variable std fallback is substituted by + ``compute_training_loss``/``score``, not by the caller. Dims: + same as ``prediction``. + """ + + @abstractmethod + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the training objective for one batch. + + The forecaster owns its complete training objective: which forecasts + to produce from the batch, which loss terms to compute from them and + how to combine those terms into a single scalar, using its own + ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The + wrapping ``BaseForecasterModule`` only injects the interior mask, + logs the returned components and optimizes the returned loss. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during forecasting. + Dims: same as the prediction. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior + nodes are scored. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The full training loss for the batch, to take gradients + of. + loss_components : dict of {str: torch.Tensor} + Scalar loss-related quantities to log alongside the loss, keyed + by component name. The wrapping module prefixes the names with + the training phase. Empty when the objective has no separate + components worth logging. + """ + + @abstractmethod + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Wrapping ``BaseForecasterModule`` subclasses use this for + validation/test logging and diagnostics instead of computing a loss + themselves: the forecaster owns both its scoring rule and its + ``pred_std`` fallback, so it is the only place that knows how to + turn a raw ``pred_std`` (possibly ``None``) into a valid one and + apply a metric to it. + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``, as returned + alongside it by ``forward``. When ``None``, implementations + substitute their own constant per-variable std fallback. + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to the forecaster's + own configured scoring rule when ``None``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). """ diff --git a/neural_lam/models/forecasters/graph_efm.py b/neural_lam/models/forecasters/graph_efm.py new file mode 100644 index 00000000..84a71b69 --- /dev/null +++ b/neural_lam/models/forecasters/graph_efm.py @@ -0,0 +1,208 @@ +"""Probabilistic forecaster training a Graph-EFM predictor via its ELBO.""" + +# Third-party +import torch + +# Local +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ..step_predictors.graph.graph_efm import BaseGraphEFM +from .probabilistic import ProbabilisticARForecaster + + +class GraphEFMForecaster(ProbabilisticARForecaster): + """ + Auto-regressive ensemble forecaster for Graph-EFM step predictors. + + Wraps a :class:`BaseGraphEFM` predictor (hierarchical ``GraphEFM`` or + flat ``GraphEFMMultiScale``), a latent-variable model consisting of a + conditional prior, a variational encoder and a latent decoder. Forecast + sampling (``forward``, ``sample_ensemble``) is inherited from + :class:`ProbabilisticARForecaster`: each predictor call samples the + prior and decodes, so unrolling produces one stochastic trajectory and + stacking several gives an ensemble. + + This class supplies the training objective the base class leaves + abstract: the evidence lower bound (ELBO). ``compute_training_loss`` + runs a variational rollout in which, at each step, the latent is drawn + from the encoder (conditioned on the target), and accumulates a + reconstruction likelihood term and a KL term between the encoder and the + prior. The scoring rule ``self.loss``, the constant per-variable std + fallback ``self.per_var_std`` and the boundary/interior masks all live + on the forecaster (set up by :class:`ARForecaster`); the predictor only + provides the network building blocks. + """ + + def __init__( + self, + predictor: BaseGraphEFM, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + kl_beta: float = 1.0, + ) -> None: + """ + Initialize the GraphEFMForecaster. + + Parameters + ---------- + predictor : BaseGraphEFM + The Graph-EFM step predictor to use for each step. Each + ``forward`` call samples the prior and decodes; the encoder, + prior and decoder sub-models are also used directly by + ``compute_training_loss`` to assemble the ELBO. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``ARForecaster.per_var_std``). Required for + training when the predictor's ``output_std`` is False, since the + likelihood term then needs the fallback std. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used for the + reconstruction likelihood term and stored as ``self.loss``. + kl_beta : float, default 1.0 + Weight of the KL term in the ELBO. When ``0`` the prior and KL + are not computed at all (pure auto-encoder training); the prior + network then receives no gradient. + """ + super().__init__(predictor, datastore, config=config, loss=loss) + self.kl_beta = kl_beta + + 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 ELBO training objective for one batch. + + Unrolls a variational rollout over the full forecast: at each step + the grid, graph and target are embedded, the latent is drawn from + the encoder (variational posterior), the decoder reconstructs the + next-state mean, and a reconstruction likelihood term is + accumulated. When ``kl_beta`` is positive a KL term between the + encoder and the prior is accumulated too. Both terms are summed over + the rollout and averaged over the batch; the loss is + ``-likelihood + kl_beta * kl``. The rollout advances on the + predicted mean with boundary nodes overwritten by the true state. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used as the encoder conditioning + and reconstruction target and to overwrite boundary nodes during + the rollout. Dims: same as the prediction. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``self.loss`` so only interior + nodes contribute to the likelihood. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The negative ELBO for the batch, to take gradients of. + loss_components : dict of {str: torch.Tensor} + Scalar ELBO diagnostics to log: ``"elbo_likelihood"`` always, + plus ``"elbo_kl"`` and ``"elbo"`` when ``kl_beta > 0``. + + Raises + ------ + ValueError + If the predictor does not output its own std and no + ``per_var_std`` fallback is available (this forecaster was + constructed without ``config``); see ``_resolve_pred_std``. + """ + predictor = self.predictor + + prev_prev_state = init_states[:, 0] + prev_state = init_states[:, 1] + pred_steps = forcing_features.shape[1] + compute_kl = self.kl_beta > 0 + + # The graph embedding depends only on static features, so it is + # constant across the rollout and computed once here. + graph_emb = predictor.embedd_graph(init_states.shape[0]) + + likelihood_terms = [] + kl_terms = [] + + for i in range(pred_steps): + forcing = forcing_features[:, i] + target_state = target_states[:, i] + + # Posterior latent (conditioned on target), reconstruction, and + # -- when a KL term is needed -- the prior, in one predictor call. + prior_dist, posterior_dist, pred_mean, pred_std = ( + predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=compute_kl, + ) + ) + pred_std = self._resolve_pred_std(pred_std) + + # Reconstruction likelihood, summed over interior grid and vars + entry_likelihoods = -self.loss( + pred_mean, + target_state, + pred_std, + mask=interior_mask_bool, + average_grid=False, + sum_vars=False, + ) # (B, num_interior_grid_nodes, num_state_vars) + likelihood_terms.append(torch.sum(entry_likelihoods, dim=(1, 2))) + + if compute_kl: + kl_terms.append( + torch.sum( + torch.distributions.kl_divergence( + posterior_dist, prior_dist + ), + dim=(1, 2), + ) + ) # (B,) + + # Advance the rollout on the predicted mean, boundary overwritten + new_state = ( + self.boundary_mask * target_state + + self.interior_mask * pred_mean + ) + prev_prev_state = prev_state + prev_state = new_state + + # Sum each term over the rollout, then average over the batch + mean_likelihood = torch.mean( + torch.sum(torch.stack(likelihood_terms, dim=1), dim=1) + ) + loss_components = {"elbo_likelihood": mean_likelihood} + + if compute_kl: + mean_kl = torch.mean(torch.sum(torch.stack(kl_terms, dim=1), dim=1)) + batch_loss = -mean_likelihood + self.kl_beta * mean_kl + loss_components["elbo_kl"] = mean_kl + loss_components["elbo"] = mean_likelihood - mean_kl + else: + batch_loss = -mean_likelihood + + return batch_loss, loss_components diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..57526d1c --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,230 @@ +"""Forecasters producing probabilistic (ensemble) forecasts.""" + +# Standard library +from abc import abstractmethod + +# Third-party +import torch + +# Local +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ..step_predictors.base import StepPredictor +from .autoregressive import ARForecaster +from .base import Forecaster + + +class ProbabilisticForecaster(Forecaster): + """ + Forecaster whose forecasts are samples from a predictive distribution. + + Adds the capability that probabilistic evaluation and ensemble-based + objectives build on: sampling an ensemble of forecasts. How the + members are produced (auto-regressive sampling, diffusion, ...) is + left to subclasses; consumers only rely on the shape of the returned + ensemble. + + When ``sample_ensemble`` returns a ``per_member_std``, it is each + member's own predicted std, not a std describing the spread across + members. The predictive distribution is then a mixture of ``S`` + Gaussians, one per member: ``p(x) = mean_s N(x; ensemble[:, s], + per_member_std[:, s]**2)``, not a single Gaussian. In particular, the + variance of that mixture is not the average of the per-member + variances: it also includes the spread between the member means. + """ + + @abstractmethod + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each + predicted step, identically in every member. Dims: same as one + member. + num_members : int + Number of ensemble members ``S`` to sample. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the forecaster predicts an std, otherwise + ``None``. Dims: same as ``ensemble``. + """ + + +class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): + """ + Auto-regressive forecaster for step predictors that sample their output. + + Each call to the wrapped predictor draws a fresh sample of the next + state, so the inherited ``ARForecaster.forward`` unrolls one sampled + trajectory. This class adds ensemble forecasting on top: unrolling + several trajectories and stacking them along an ensemble dimension. + + ``compute_training_loss`` is intentionally left abstract here (it does + not fall back to ``ARForecaster``'s single-rollout objective, which + would silently train on one stochastic sample). There is no default + objective that fits every stochastic model: scoring the ensemble mean + with a pointwise metric only rewards the mean being right, giving the + model no incentive to keep a calibrated spread, and risks training it + to collapse the ensemble to a point estimate. Concrete subclasses must + define an objective appropriate to how they are meant to be trained + (e.g. an ensemble scoring rule such as CRPS, or a variational + objective). + """ + + def __init__( + self, + predictor: StepPredictor, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + ) -> None: + """ + Initialize the ProbabilisticARForecaster. + + Parameters + ---------- + predictor : StepPredictor + The predictor to use for each step. Each call should draw a + fresh sample of the next state. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. + """ + super().__init__(predictor, datastore, config=config, loss=loss) + + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Unrolls ``num_members`` independent forecasts, each sampling fresh + randomness at every step, and stacks them along a new ensemble + dimension after the batch dimension. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each AR + step, identically in every member. Dims: same as one member. + num_members : int + Number of ensemble members ``S`` to sample. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the wrapped predictor outputs an std, otherwise + ``None``. Dims: same as ``ensemble``. + + Raises + ------ + ValueError + If ``num_members`` is less than 1. + """ + if num_members < 1: + raise ValueError( + f"num_members must be at least 1, got {num_members}" + ) + + member_list = [] + member_std_list = [] + for _ in range(num_members): + prediction, pred_std = self( + init_states, forcing_features, boundary_states + ) + member_list.append(prediction) + if pred_std is not None: + member_std_list.append(pred_std) + + ensemble = torch.stack(member_list, dim=1) + # After stacking, ensemble has shape + # (B, S, pred_steps, num_grid_nodes, num_state_vars) + per_member_std = ( + torch.stack(member_std_list, dim=1) if member_std_list else None + ) + return ensemble, per_member_std + + @abstractmethod + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the training objective for one batch. + + Left abstract; see the class docstring for why there is no default + objective. Concrete subclasses typically call ``sample_ensemble`` + and score the resulting members with an objective appropriate to + the model (see ``Forecaster.compute_training_loss`` for the + signature and general contract). + """ diff --git a/neural_lam/models/latent/__init__.py b/neural_lam/models/latent/__init__.py new file mode 100644 index 00000000..fba7ed42 --- /dev/null +++ b/neural_lam/models/latent/__init__.py @@ -0,0 +1,22 @@ +"""Latent encoder and decoder modules for latent-variable models such as +GraphEFM, mapping between grid representations and distributions over +latent variables on mesh nodes.""" + +# Local +from .base_decoder import BaseGraphLatentDecoder +from .base_encoder import BaseLatentEncoder +from .constant_encoder import ConstantLatentEncoder +from .graph_decoder import GraphLatentDecoder +from .graph_encoder import GraphLatentEncoder +from .hi_graph_decoder import HiGraphLatentDecoder +from .hi_graph_encoder import HiGraphLatentEncoder + +__all__ = [ + "BaseGraphLatentDecoder", + "BaseLatentEncoder", + "ConstantLatentEncoder", + "GraphLatentDecoder", + "GraphLatentEncoder", + "HiGraphLatentDecoder", + "HiGraphLatentEncoder", +] diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py new file mode 100644 index 00000000..dbd27e58 --- /dev/null +++ b/neural_lam/models/latent/base_decoder.py @@ -0,0 +1,171 @@ +"""Abstract base class for graph-based latent decoders.""" + +# Third-party +from torch import nn + +# First-party +from neural_lam import utils + + +class BaseGraphLatentDecoder(nn.Module): + """ + Abstract decoder mapping a grid representation plus a latent sample on + mesh to the next-state increment (and optionally std) on the grid. + + Subclasses implement :meth:`combine_with_latent`, which fuses the latent + representation with the grid representation. The resulting features are + mapped to either ``num_state_vars`` outputs (mean increment only) or + ``2 * num_state_vars`` outputs (mean increment, std) depending on + ``output_std``. + """ + + def __init__( + self, + hidden_dim, + latent_dim, + num_state_vars, + hidden_layers=1, + output_std=True, + ): + """ + Set up the latent embedder, grid-residual MLP and output param map. + + Parameters + ---------- + hidden_dim : int + Dimensionality of internal node and edge representations. + Latent samples are embedded to this dimensionality before + being fused with the grid representation. + latent_dim : int + Dimensionality of the latent variable at each mesh node, i.e. + the feature dimension of the ``latent_samples`` given to + ``forward``. + num_state_vars : int + Number of state variables predicted at each grid node, i.e. + the feature dimension of the predicted mean (and std). + hidden_layers : int + Number of hidden layers in the internal MLPs (latent embedder, + grid-residual MLP and output parameter map). + output_std : bool + If True, the decoder outputs both mean and std of the + next-state distribution (the output parameter map produces + ``2 * num_state_vars`` features per grid node); if False, only + the mean. + """ + super().__init__() + + self.grid_update_mlp = utils.make_mlp( + [hidden_dim] * (hidden_layers + 2) + ) + + self.latent_embedder = utils.make_mlp( + [latent_dim] + [hidden_dim] * (hidden_layers + 1) + ) + + self.output_std = output_std + if self.output_std: + output_dim = 2 * num_state_vars + else: + output_dim = num_state_vars + + self.param_map = utils.make_mlp( + [hidden_dim] * (hidden_layers + 1) + [output_dim], layer_norm=False + ) + + def combine_with_latent( + self, original_grid_rep, latent_rep, residual_grid_rep, graph_emb + ): + """ + Fuse grid and latent representations and return a grid-shaped output. + + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input + features, used as the sender representation when encoding the + grid onto the mesh. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes, d_h)``. Latent sample embedded to + the internal dimensionality ``d_h``. Where this enters the + message passing is up to the concrete decoder. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Residually updated grid + representation, used as the receiver representation when + decoding the mesh back to the grid. This keeps a direct path + from the grid input to the output. + graph_emb : dict + Embedded static graph node and edge features. The required + entries depend on the concrete decoder, but include at least + the ``g2m``, ``m2m`` and ``m2g`` edge embeddings. + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation, incorporating both the grid input and the + latent sample. + """ + raise NotImplementedError("combine_with_latent not implemented") + + def forward(self, grid_rep, latent_samples, graph_emb): + """ + Predict the next-state increment (and optionally std). + + The latent samples are embedded to the internal dimensionality and + fused with the grid representation by ``combine_with_latent``; the + result is mapped to distribution parameters. The predicted mean is an + increment relative to the current state, which the caller adds onto + the current state (and clamps). + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input + features (states, forcing and static features). + latent_samples : torch.Tensor + Shape ``(B, num_mesh_nodes, latent_dim)``. Sample of the + latent variable on the mesh nodes, e.g. drawn from the prior + or the variational distribution. + graph_emb : dict + Embedded static graph node and edge features, forwarded to + ``combine_with_latent``; includes at least the ``g2m``, + ``m2m`` and ``m2g`` edge embeddings. + + Returns + ------- + mean_delta : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted + increment of the next state relative to the current state. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, num_state_vars)`` when + ``output_std`` is True, otherwise None. Predicted std of the + next state, obtained from the output parameter map through a + softplus. + """ + latent_emb = self.latent_embedder(latent_samples) + + # Residually update the grid rep with a node-wise MLP. This gives a + # direct path from grid input to output that bypasses the mesh, used + # as the receiver (base) representation that mesh information is + # added onto in the final mesh-to-grid step of combine_with_latent. + residual_grid_rep = grid_rep + self.grid_update_mlp(grid_rep) + + combined_grid_rep = self.combine_with_latent( + grid_rep, latent_emb, residual_grid_rep, graph_emb + ) + + state_params = self.param_map(combined_grid_rep) + + if self.output_std: + # When outputting std the param map produces 2 * num_state_vars + # features per grid node. Split these along the feature dim into + # the first half (mean delta) and second half (unconstrained std + # parameters), then map the latter through a softplus to get + # positive std values. + mean_delta, std_raw = state_params.chunk(2, dim=-1) + pred_std = nn.functional.softplus(std_raw) + else: + mean_delta = state_params + pred_std = None + + return mean_delta, pred_std diff --git a/neural_lam/models/latent/base_encoder.py b/neural_lam/models/latent/base_encoder.py new file mode 100644 index 00000000..cad46d4f --- /dev/null +++ b/neural_lam/models/latent/base_encoder.py @@ -0,0 +1,94 @@ +"""Abstract base class for latent encoders.""" + +# Third-party +import torch +from torch import distributions as tdists +from torch import nn + + +class BaseLatentEncoder(nn.Module): + """ + Abstract encoder mapping an input grid representation to a Gaussian + distribution over a latent variable defined on mesh nodes. + + Subclasses implement :meth:`compute_dist_params`, which returns the raw + parameters used to build the output distribution. With + ``output_dist="isotropic"`` only the mean is produced (unit variance); + with ``output_dist="diagonal"`` both mean and a positive std are output. + """ + + def __init__(self, latent_dim, output_dist="isotropic"): + """ + Set up output dimensionality for the chosen distribution type. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + output_dist : str + Type of output distribution: ``"isotropic"`` (mean only, unit + variance) or ``"diagonal"`` (mean and per-dimension std). + """ + super().__init__() + + self.output_dist = output_dist + if output_dist == "isotropic": + self.output_dim = latent_dim + elif output_dist == "diagonal": + self.output_dim = 2 * latent_dim + # Small floor to prevent the encoder from collapsing to std 0 + self.latent_std_eps = 1e-4 + else: + raise ValueError( + f"Unknown encoder output distribution: {output_dist}" + ) + + def compute_dist_params(self, grid_rep, **kwargs): + """ + Compute raw distribution parameters from the grid representation. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation. + **kwargs + Additional inputs used by concrete encoders (e.g. graph + embeddings). + + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of + the latent distribution. + """ + raise NotImplementedError("compute_dist_params not implemented") + + def forward(self, grid_rep, **kwargs): + """ + Compute the Gaussian distribution over the latent variable. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation. + **kwargs + Additional inputs forwarded to :meth:`compute_dist_params`. + + Returns + ------- + torch.distributions.Normal + Distribution over the latent variable, with batch shape + ``(B, num_mesh_nodes, latent_dim)``. + """ + latent_dist_params = self.compute_dist_params(grid_rep, **kwargs) + + if self.output_dist == "diagonal": + latent_mean, latent_std_raw = latent_dist_params.chunk(2, dim=-1) + latent_std = self.latent_std_eps + nn.functional.softplus( + latent_std_raw + ) + else: + latent_mean = latent_dist_params + latent_std = torch.ones_like(latent_mean) + + return tdists.Normal(latent_mean, latent_std) diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py new file mode 100644 index 00000000..f1bf3be5 --- /dev/null +++ b/neural_lam/models/latent/constant_encoder.py @@ -0,0 +1,62 @@ +"""Constant (input-independent) latent encoder.""" + +# Third-party +import torch + +# Local +from .base_encoder import BaseLatentEncoder + + +class ConstantLatentEncoder(BaseLatentEncoder): + """ + Latent encoder that returns a constant (input-independent) distribution. + + ``compute_dist_params`` returns a tensor of zeros, so the resulting + Normal is ``Normal(mean=0, std=1)`` for ``output_dist="isotropic"`` and + ``Normal(mean=0, std=softplus(0)+eps)`` for ``output_dist="diagonal"``. + """ + + def __init__(self, latent_dim, num_mesh_nodes, output_dist="isotropic"): + """ + Store the number of mesh nodes to produce parameters for. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + num_mesh_nodes : int + Number of mesh nodes the latent variable is defined on. + output_dist : str + Type of output distribution: ``"isotropic"`` or ``"diagonal"``. + """ + super().__init__(latent_dim, output_dist) + self.num_mesh_nodes = num_mesh_nodes + + def compute_dist_params(self, grid_rep, **kwargs): + """ + Compute raw distribution parameters from the grid representation. + + For this constant encoder the parameters are all zeros, independent + of the values in ``grid_rep``. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation, + used only to determine batch size and device. + **kwargs + Ignored; accepted for compatibility with the base class + interface. + + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of + the latent distribution, all zeros. + """ + return torch.zeros( + grid_rep.shape[0], + self.num_mesh_nodes, + self.output_dim, + device=grid_rep.device, + ) diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py new file mode 100644 index 00000000..d1ab64f5 --- /dev/null +++ b/neural_lam/models/latent/graph_decoder.py @@ -0,0 +1,134 @@ +"""Latent decoder for flat (non-hierarchical) graphs.""" + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import get_gnn_class + +# Local +from .base_decoder import BaseGraphLatentDecoder + + +class GraphLatentDecoder(BaseGraphLatentDecoder): + """ + Latent decoder for a flat (non-hierarchical) graph. Encodes grid into + mesh with a g2m GNN (type set by ``g2m_gnn_type``), processes on mesh, + and reads back out to grid with an m2g GNN (type set by ``m2g_gnn_type``). + The grid representation also goes through a residual MLP that is added + back to the mesh-to-grid output. + """ + + def __init__( + self, + g2m_edge_index, + m2m_edge_index, + m2g_edge_index, + hidden_dim, + latent_dim, + num_state_vars, + m2m_layers, + hidden_layers=1, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="InteractionNet", + output_std=True, + ): + """ + Set up the g2m, on-mesh and m2g GNNs. + + Parameters + ---------- + g2m_edge_index : torch.Tensor + Shape ``(2, M_g2m)``. Edge index of grid-to-mesh edges. + m2m_edge_index : torch.Tensor + Shape ``(2, M_m2m)``. Edge index of mesh-to-mesh edges. + m2g_edge_index : torch.Tensor + Shape ``(2, M_m2g)``. Edge index of mesh-to-grid edges. + hidden_dim : int + Dimensionality of internal node and edge representations. + latent_dim : int + Dimensionality of the latent variable at each mesh node. + num_state_vars : int + Number of state variables predicted at each grid node. + m2m_layers : int + Number of on-mesh (m2m) GNN layers; 0 disables on-mesh + processing. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step (key in + ``gnn_layers.GNN_TYPES``). + m2g_gnn_type : str + GNN type for the mesh-to-grid step (key in + ``gnn_layers.GNN_TYPES``). + output_std : bool + If True, the decoder outputs both mean and std of the next-state + distribution; if False, only the mean. + """ + super().__init__( + hidden_dim, latent_dim, num_state_vars, hidden_layers, output_std + ) + + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + # None if m2m_layers == 0, in which case no on-mesh processing is + # done in combine_with_latent + self.m2m_gnns = ( + utils.make_gnn_seq( + m2m_edge_index, m2m_layers, hidden_layers, hidden_dim + ) + if m2m_layers > 0 + else None + ) + + self.m2g_gnn = get_gnn_class(m2g_gnn_type)( + m2g_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + def combine_with_latent( + self, original_grid_rep, latent_rep, residual_grid_rep, graph_emb + ): + """ + Fuse grid and latent reps via g2m -> m2m -> m2g. + + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input + features, used as the sender representation in the + grid-to-mesh step. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes, d_h)``. Latent sample embedded to + ``d_h``, used as the initial mesh node representation (the + receiver in the grid-to-mesh step), so all mesh processing + starts from the latent. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Residually updated grid + representation, used as the receiver representation in the + mesh-to-grid step. + graph_emb : dict + Embedded static graph node and edge features, with at least + the entries ``g2m``: ``(B, M_g2m, d_h)``, + ``m2m``: ``(B, M_m2m, d_h)`` and ``m2g``: ``(B, M_m2g, d_h)``. + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation, incorporating both the grid input and the + latent sample. + """ + mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) + + if self.m2m_gnns is not None: + mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) + + grid_rep = self.m2g_gnn(mesh_rep, residual_grid_rep, graph_emb["m2g"]) + + return grid_rep diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py new file mode 100644 index 00000000..9082f999 --- /dev/null +++ b/neural_lam/models/latent/graph_encoder.py @@ -0,0 +1,103 @@ +"""Latent encoder for flat (non-hierarchical) graphs.""" + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import get_gnn_class + +# Local +from .base_encoder import BaseLatentEncoder + + +class GraphLatentEncoder(BaseLatentEncoder): + """ + Latent encoder that maps grid features to mesh and outputs a Gaussian + distribution over a latent variable on mesh nodes. Uses a flat + (non-hierarchical) graph: one g2m GNN (type set by ``g2m_gnn_type``) + followed by a stack of on-mesh (m2m) InteractionNet layers. + """ + + def __init__( + self, + latent_dim, + g2m_edge_index, + m2m_edge_index, + hidden_dim, + m2m_layers, + hidden_layers=1, + g2m_gnn_type="InteractionNet", + output_dist="isotropic", + ): + """ + Set up the g2m GNN, on-mesh processing stack and latent param map. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + g2m_edge_index : torch.Tensor + Shape ``(2, M_g2m)``. Edge index of grid-to-mesh edges. + m2m_edge_index : torch.Tensor + Shape ``(2, M_m2m)``. Edge index of mesh-to-mesh edges. + hidden_dim : int + Dimensionality of internal node and edge representations. + m2m_layers : int + Number of on-mesh (m2m) GNN layers; 0 disables on-mesh + processing. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step (key in + ``gnn_layers.GNN_TYPES``). + output_dist : str + Type of output distribution: ``"isotropic"`` or ``"diagonal"``. + """ + super().__init__(latent_dim, output_dist) + + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + # None if m2m_layers == 0, in which case no on-mesh processing is + # done in compute_dist_params + self.m2m_gnns = ( + utils.make_gnn_seq( + m2m_edge_index, m2m_layers, hidden_layers, hidden_dim + ) + if m2m_layers > 0 + else None + ) + + self.latent_param_map = utils.make_mlp( + [hidden_dim] * (hidden_layers + 1) + [self.output_dim], + layer_norm=False, + ) + + # pylint: disable-next=arguments-differ + def compute_dist_params(self, grid_rep, graph_emb, **kwargs): + """ + Compute distribution parameters on mesh from grid features. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation. + graph_emb : dict + Embedded graph node and edge features, with at least entries + ``mesh``: ``(B, num_mesh_nodes, d_h)``, + ``g2m``: ``(B, M_g2m, d_h)`` and ``m2m``: ``(B, M_m2m, d_h)``. + **kwargs + Ignored. + + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of + the latent distribution. + """ + mesh_rep = self.g2m_gnn(grid_rep, graph_emb["mesh"], graph_emb["g2m"]) + if self.m2m_gnns is not None: + mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) + return self.latent_param_map(mesh_rep) diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py new file mode 100644 index 00000000..83a9fede --- /dev/null +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -0,0 +1,270 @@ +"""Latent decoder for hierarchical graphs.""" + +# Third-party +from torch import nn + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import ( + InteractionNet, + PropagationNet, + get_gnn_class, +) + +# Local +from .base_decoder import BaseGraphLatentDecoder + + +class HiGraphLatentDecoder(BaseGraphLatentDecoder): + """ + Latent decoder for a hierarchical mesh. The grid representation is + encoded into the bottom mesh level; the message-passing then propagates + *up* through the hierarchy (mixing in the latent at the top level), then + *down* through the hierarchy with residual connections back to the + intra-level reps from the upward pass, and finally maps back to grid + via an m2g GNN (type set by ``m2g_gnn_type``). The g2m GNN type is set + by ``g2m_gnn_type``; mesh-up edges always use InteractionNets and + mesh-down edges always use PropagationNets. + """ + + def __init__( + self, + g2m_edge_index, + m2m_edge_index, + m2g_edge_index, + mesh_up_edge_index, + mesh_down_edge_index, + hidden_dim, + latent_dim, + num_state_vars, + intra_level_layers, + hidden_layers=1, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="InteractionNet", + output_std=True, + ): + """ + Set up the g2m, m2g, mesh-up/-down and intra-level GNNs. + + Parameters + ---------- + g2m_edge_index : torch.Tensor + Shape ``(2, M_g2m)``. Edge index of grid-to-mesh edges. + m2m_edge_index : BufferList + Per-level edge indices of intra-level mesh edges, each of shape + ``(2, M_m2m[l])``. + m2g_edge_index : torch.Tensor + Shape ``(2, M_m2g)``. Edge index of mesh-to-grid edges. + mesh_up_edge_index : BufferList + Per-level edge indices of upward inter-level mesh edges, each of + shape ``(2, M_up[l])``. + mesh_down_edge_index : BufferList + Per-level edge indices of downward inter-level mesh edges, each + of shape ``(2, M_down[l])``. + hidden_dim : int + Dimensionality of internal node and edge representations. + latent_dim : int + Dimensionality of the latent variable at each mesh node. + num_state_vars : int + Number of state variables predicted at each grid node. + intra_level_layers : int + Number of intra-level GNN layers at each mesh level; 0 disables + intra-level processing. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step (key in + ``gnn_layers.GNN_TYPES``). + m2g_gnn_type : str + GNN type for the mesh-to-grid step (key in + ``gnn_layers.GNN_TYPES``). Inter-level edges are not + configurable; mesh-up edges always use ``InteractionNet`` and + mesh-down edges always use ``PropagationNet``. + output_std : bool + If True, the decoder outputs both mean and std of the next-state + distribution; if False, only the mean. + """ + super().__init__( + hidden_dim, latent_dim, num_state_vars, hidden_layers, output_std + ) + + # Hierarchical decoder needs at least 2 mesh levels; with a single + # level the up/down passes are empty and the latent would be + # silently ignored. Use GraphLatentDecoder instead. + if len(m2m_edge_index) < 2: + raise ValueError( + "HiGraphLatentDecoder requires at least 2 mesh levels " + f"(got {len(m2m_edge_index)}). Use GraphLatentDecoder for " + "flat graphs." + ) + + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + self.m2g_gnn = get_gnn_class(m2g_gnn_type)( + m2g_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + # Mesh-up edges must use InteractionNet: with a PropagationNet the + # latent rep at the top level would be overwritten rather than + # residually updated, leaving Z unused at initialization. + self.mesh_up_gnns = nn.ModuleList( + [ + InteractionNet( + edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + for edge_index in mesh_up_edge_index + ] + ) + # Mesh-down edges must use PropagationNet: each downward step has to + # push the latent information from the level above into the lower + # level, so that Z reaches the grid output. + self.mesh_down_gnns = nn.ModuleList( + [ + PropagationNet( + edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + for edge_index in mesh_down_edge_index + ] + ) + + # None if intra_level_layers == 0, in which case no intra-level + # processing is done in combine_with_latent + self.intra_up_gnns = ( + nn.ModuleList( + [ + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + for edge_index in m2m_edge_index + ] + ) + if intra_level_layers > 0 + else None + ) + self.intra_down_gnns = ( + nn.ModuleList( + [ + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + for edge_index in list(m2m_edge_index)[:-1] + # Top level (L) does not need a down intra-level GNN + ] + ) + if intra_level_layers > 0 + else None + ) + + def combine_with_latent( + self, original_grid_rep, latent_rep, residual_grid_rep, graph_emb + ): + """ + Hierarchical up-then-down fusion of grid and latent reps. + + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input + features, used as the sender representation in the + grid-to-mesh step that initializes the bottom mesh level. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes[L], d_h)``. Latent sample embedded + to ``d_h``, defined on the top mesh level ``L``. It is used as + the receiver representation of the last upward step, fusing + the latent in at the top of the hierarchy. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Residually updated grid + representation, used as the receiver representation in the + mesh-to-grid step. + graph_emb : dict + Embedded static graph node and edge features, with at least + the entries + ``mesh``: list of ``(B, num_mesh_nodes[l], d_h)``, + ``g2m``: ``(B, M_g2m, d_h)``, + ``m2m``: list of ``(B, M_m2m[l], d_h)``, + ``mesh_up``: list of ``(B, M_up[l], d_h)``, + ``mesh_down``: list of ``(B, M_down[l], d_h)`` and + ``m2g``: ``(B, M_m2g, d_h)``, where ``l`` indexes the mesh + levels from bottom (0) to top (``L``). + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation, incorporating both the grid input and the + latent sample. + """ + current_mesh_rep = self.g2m_gnn( + original_grid_rep, graph_emb["mesh"][0], graph_emb["g2m"] + ) + + # Upward pass: intra-level processing, then up to the next level. + # On the last upward step, the latent replaces the level-L mesh rep + # so the latent is fused in at the top of the hierarchy. + mesh_level_reps = [] + m2m_level_reps = [] + for level, (up_gnn, mesh_up_level_rep, mesh_level_rep) in enumerate( + zip( + self.mesh_up_gnns, + graph_emb["mesh_up"], + graph_emb["mesh"][1:-1] + [latent_rep], + ) + ): + new_mesh_rep = current_mesh_rep + new_m2m_rep = None + if self.intra_up_gnns is not None: + new_mesh_rep, new_m2m_rep = self.intra_up_gnns[level]( + new_mesh_rep, graph_emb["m2m"][level] + ) + + # Saved for residual connections in the downward pass + mesh_level_reps.append(new_mesh_rep) + m2m_level_reps.append(new_m2m_rep) + + current_mesh_rep = up_gnn( + new_mesh_rep, mesh_level_rep, mesh_up_level_rep + ) + + # Top level processing + if self.intra_up_gnns is not None: + current_mesh_rep, _ = self.intra_up_gnns[-1]( + current_mesh_rep, graph_emb["m2m"][-1] + ) + + # Downward pass: down GNN, then intra-level processing. Residual + # connections feed back the intra-level reps from the upward pass. + for level in reversed(range(len(self.mesh_down_gnns))): + current_mesh_rep = self.mesh_down_gnns[level]( + current_mesh_rep, + mesh_level_reps[level], + graph_emb["mesh_down"][level], + ) + if self.intra_down_gnns is not None: + current_mesh_rep, _ = self.intra_down_gnns[level]( + current_mesh_rep, m2m_level_reps[level] + ) + + grid_rep = self.m2g_gnn( + current_mesh_rep, residual_grid_rep, graph_emb["m2g"] + ) + + return grid_rep diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py new file mode 100644 index 00000000..e0148b38 --- /dev/null +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -0,0 +1,173 @@ +"""Latent encoder for hierarchical graphs.""" + +# Third-party +from torch import nn + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import PropagationNet, get_gnn_class + +# Local +from .base_encoder import BaseLatentEncoder + + +class HiGraphLatentEncoder(BaseLatentEncoder): + """ + Latent encoder for a hierarchical mesh: grid -> bottom mesh level via a + g2m GNN (type set by ``g2m_gnn_type``), then propagates upward through + mesh levels using mesh-up PropagationNets, with optional intra-level + processing at each level. The latent distribution is read out from the + top mesh level. + """ + + def __init__( + self, + latent_dim, + g2m_edge_index, + m2m_edge_index, + mesh_up_edge_index, + hidden_dim, + intra_level_layers, + hidden_layers=1, + g2m_gnn_type="InteractionNet", + output_dist="isotropic", + ): + """ + Set up the g2m, mesh-up and intra-level GNNs and latent param map. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + g2m_edge_index : torch.Tensor + Shape ``(2, M_g2m)``. Edge index of grid-to-mesh edges. + m2m_edge_index : BufferList + Per-level edge indices of intra-level mesh edges, each of shape + ``(2, M_m2m[l])``. + mesh_up_edge_index : BufferList + Per-level edge indices of upward inter-level mesh edges, each of + shape ``(2, M_up[l])``. + hidden_dim : int + Dimensionality of internal node and edge representations. + intra_level_layers : int + Number of intra-level GNN layers at each mesh level; 0 disables + intra-level processing. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step (key in + ``gnn_layers.GNN_TYPES``). Mesh-up edges are not configurable; + they always use ``PropagationNet``. + output_dist : str + Type of output distribution: ``"isotropic"`` or ``"diagonal"``. + """ + super().__init__(latent_dim, output_dist) + + # Hierarchical encoder needs at least 2 mesh levels; with a single + # level there is no upward propagation and the latent readout would + # collapse to a flat encoder. Use GraphLatentEncoder instead. + if len(m2m_edge_index) < 2: + raise ValueError( + "HiGraphLatentEncoder requires at least 2 mesh levels " + f"(got {len(m2m_edge_index)}). Use GraphLatentEncoder for " + "flat graphs." + ) + + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + # Mesh-up edges must use PropagationNet: each upward step has to push + # information into nodes of the next level even when those start from + # their static embedding, so that grid information reaches the latent + # readout at the top level. + self.mesh_up_gnns = nn.ModuleList( + [ + PropagationNet( + edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + for edge_index in mesh_up_edge_index + ] + ) + + # None if intra_level_layers == 0, in which case no intra-level + # processing is done in compute_dist_params + self.intra_level_gnns = ( + nn.ModuleList( + [ + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + for edge_index in m2m_edge_index + ] + ) + if intra_level_layers > 0 + else None + ) + + self.latent_param_map = utils.make_mlp( + [hidden_dim] * (hidden_layers + 1) + [self.output_dim], + layer_norm=False, + ) + + # pylint: disable-next=arguments-differ + def compute_dist_params(self, grid_rep, graph_emb, **kwargs): + """ + Compute distribution parameters on the top mesh level. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation. + graph_emb : dict + Embedded graph node and edge features, with at least entries + ``mesh``: list of ``(B, num_mesh_nodes[l], d_h)``, + ``g2m``: ``(B, M_g2m, d_h)``, + ``m2m``: list of ``(B, M_m2m[l], d_h)`` and + ``mesh_up``: list of ``(B, M_up[l], d_h)``. + **kwargs + Ignored. + + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes[L], output_dim)``. Raw parameters + of the latent distribution on the top mesh level ``L``. + """ + current_mesh_rep = self.g2m_gnn( + grid_rep, graph_emb["mesh"][0], graph_emb["g2m"] + ) + + # Same-level processing on level 0 + if self.intra_level_gnns is not None: + current_mesh_rep, _ = self.intra_level_gnns[0]( + current_mesh_rep, graph_emb["m2m"][0] + ) + + # Walk up levels 1..L + for level, (up_gnn, mesh_up_level_rep, mesh_level_rep) in enumerate( + zip( + self.mesh_up_gnns, + graph_emb["mesh_up"], + graph_emb["mesh"][1:], + ), + start=1, + ): + current_mesh_rep = up_gnn( + current_mesh_rep, mesh_level_rep, mesh_up_level_rep + ) + if self.intra_level_gnns is not None: + current_mesh_rep, _ = self.intra_level_gnns[level]( + current_mesh_rep, graph_emb["m2m"][level] + ) + + return self.latent_param_map(current_mesh_rep) diff --git a/neural_lam/models/modules/__init__.py b/neural_lam/models/modules/__init__.py new file mode 100644 index 00000000..b2e70700 --- /dev/null +++ b/neural_lam/models/modules/__init__.py @@ -0,0 +1,8 @@ +""" +Lightning modules wrapping forecasters for training and evaluation. +""" + +# Local +from .base import BaseForecasterModule +from .deterministic import DeterministicForecasterModule +from .probabilistic import ProbabilisticForecasterModule diff --git a/neural_lam/models/module.py b/neural_lam/models/modules/base.py similarity index 73% rename from neural_lam/models/module.py rename to neural_lam/models/modules/base.py index b097d417..94823c6f 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/modules/base.py @@ -1,9 +1,10 @@ -"""Lightning module handling training, validation and testing loops.""" +"""Abstract Lightning module shared by deterministic and probabilistic +forecaster modules.""" # Standard library import os import warnings -from typing import Any +from abc import ABC, abstractmethod # Third-party import matplotlib.pyplot as plt @@ -17,18 +18,26 @@ from neural_lam.utils import get_integer_time # Local -from .. import metrics, vis -from ..config import NeuralLAMConfig -from ..datastore import BaseDatastore -from ..loss_weighting import get_state_feature_weighting -from ..weather_dataset import WeatherDataset -from .forecasters.base import Forecaster +from ... import vis +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ...weather_dataset import WeatherDataset +from ..forecasters.base import Forecaster -class ForecasterModule(pl.LightningModule): +class BaseForecasterModule(pl.LightningModule, ABC): """ - Lightning module handling training, validation and testing loops. - Wraps a Forecaster instance which performs the actual prediction. + Abstract Lightning module wrapping a ``Forecaster``. + + Owns everything that does not depend on whether the wrapped forecaster + produces a single deterministic forecast or samples an ensemble: + batch standardization, the training loop, optimizer configuration, + checkpoint compatibility, and the plotting/aggregation helpers used by + validation and testing. ``validation_step``, ``test_step`` and + ``on_test_epoch_end`` differ enough between the two evaluation modes + that they are left abstract; concrete subclasses implement them + independently (see ``DeterministicForecasterModule`` and + ``ProbabilisticForecasterModule``) rather than overriding one another. """ # pylint: disable=arguments-differ @@ -38,7 +47,6 @@ def __init__( forecaster: Forecaster, config: NeuralLAMConfig, datastore: BaseDatastore, - loss: str = "wmse", lr: float = 1e-3, restore_opt: bool = False, n_example_pred: int = 1, @@ -50,18 +58,20 @@ def __init__( args=None, ): """ - Initialize the ForecasterModule. + Initialize the BaseForecasterModule. Parameters ---------- forecaster : Forecaster - The forecaster model to use for predictions. + The forecaster model to use for predictions. Owns the training + objective (``compute_training_loss``) and validation/test + scoring (``score``); this module and its subclasses never + compute a loss themselves, only inject shared inputs (e.g. the + interior mask) and log what the forecaster returns. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore Datastore providing grid metadata and data access. - loss : str, default "wmse" - The loss function to use. lr : float, default 1e-3 Learning rate for the optimizer. restore_opt : bool, default False @@ -72,6 +82,10 @@ def __init__( Whether to create GIFs of example predictions. val_steps_to_log : list of int, optional Specific rollout steps to log during validation/testing. + train_steps_to_log : list of int, optional + Specific rollout steps to log during training. Accepted for CLI + and checkpoint compatibility; ``training_step`` only ever logs + the aggregate ``train_loss`` (see its docstring for why). metrics_watch : list of str, optional List of metrics to watch and log specifically. var_leads_metrics_watch : dict of {int: list of int}, optional @@ -80,21 +94,20 @@ def __init__( args : argparse.Namespace, optional Pre-refactor ``ARModel`` checkpoint hyperparameters. When provided, attributes on ``args`` take precedence over the - corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``, + corresponding explicit kwargs (``lr``, ``restore_opt``, ``n_example_pred``, ``create_gif``, ``val_steps_to_log``, - ``metrics_watch``, ``var_leads_metrics_watch``) so legacy - checkpoints round-trip through ``load_from_checkpoint`` - correctly. + ``train_steps_to_log``, ``metrics_watch``, + ``var_leads_metrics_watch``) so legacy checkpoints round-trip + through ``load_from_checkpoint`` correctly. """ super().__init__() # Pre-refactor ``ARModel`` checkpoints saved every hyperparameter nested # inside an argparse Namespace under the single key 'args'. When # Lightning calls __init__ during load_from_checkpoint it would # otherwise drop 'args' (not in the new signature) and silently fall - # back to defaults for loss/lr/create_gif/etc. Unpack the namespace - # here so legacy checkpoints round-trip correctly. + # back to defaults for lr/create_gif/etc. Unpack the namespace here + # so legacy checkpoints round-trip correctly. if args is not None: - loss = getattr(args, "loss", loss) lr = getattr(args, "lr", lr) restore_opt = getattr(args, "restore_opt", restore_opt) n_example_pred = getattr(args, "n_example_pred", n_example_pred) @@ -145,29 +158,6 @@ def __init__( persistent=False, ) - # Store per_var_std here if predictor does not output std - if not self.forecaster.predicts_std: - da_state_stats = datastore.get_standardization_dataarray( - category="state" - ) - state_feature_weights = get_state_feature_weighting( - config=config, datastore=datastore - ) - diff_std = torch.tensor( - da_state_stats.state_diff_std_standardized.values, - dtype=torch.float32, - ) - feature_weights_t = torch.tensor( - state_feature_weights, dtype=torch.float32 - ) - self.register_buffer( - "per_var_std", - diff_std / torch.sqrt(feature_weights_t), - persistent=False, - ) - else: - self.per_var_std = None - # Standardization statistics used to normalize each batch on-device in # `on_after_batch_transfer`. WeatherDataset returns unstandardized # data, so state and forcing are normalized here rather than on CPU. @@ -213,19 +203,6 @@ def __init__( self.forcing_mean = None self.forcing_std = None - # Instantiate loss function - self.loss = metrics.get_metric(loss) - - self.val_metrics: dict[str, list] = { - "mse": [], - } - self.test_metrics: dict[str, list] = { - "mse": [], - "mae": [], - } - if self.forecaster.predicts_std: - self.test_metrics["output_std"] = [] # Treat as metric - # For making restoring of optimizer state optional self.restore_opt = restore_opt @@ -234,9 +211,6 @@ def __init__( self.create_gif = create_gif self.plotted_examples = 0 - # For storing spatial loss maps during evaluation - self.spatial_loss_maps: list[Any] = [] - # Warn once per phase if steps_to_log exceeds the actual rollout self._steps_warn_issued = { "val": False, @@ -347,7 +321,7 @@ def on_after_batch_transfer(self, batch, dataloader_idx): def common_step(self, batch): """ - Perform a common prediction step for training, validation, and testing. + Perform a common prediction step for validation and testing. Parameters ---------- @@ -371,6 +345,15 @@ def training_step(self, batch): """ Perform a single training step. + The training objective is fully assembled by the wrapped forecaster, + which owns its own scoring rule; this method injects the interior + mask, then logs the loss and any loss components the forecaster + returns. Unlike validation/test, no per-step ``train_loss_unroll{i}`` + breakdown is logged here even when ``train_steps_to_log`` is set: + ``compute_training_loss`` intentionally returns only a scalar (a + forecaster's objective, e.g. an ELBO, may not decompose per rollout + step at all), so there is no per-step tensor to select from. + Parameters ---------- batch : tuple @@ -381,12 +364,26 @@ def training_step(self, batch): torch.Tensor The computed loss for the training step. """ - _, _, _, time_step_loss = self._compute_prediction_and_loss(batch) - batch_loss = torch.mean(time_step_loss) - batch_size = batch[0].shape[0] - - self._log_step_loss(time_step_loss, batch_loss, "train", batch_size) + init_states, target_states, forcing_features, _ = batch + batch_loss, loss_components = self.forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=self.interior_mask_bool, + ) + log_dict = { + f"train_{name}": value for name, value in loss_components.items() + } + log_dict["train_loss"] = batch_loss + self.log_dict( + log_dict, + prog_bar=True, + on_step=True, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) return batch_loss def all_gather_cat(self, tensor_to_gather): @@ -432,55 +429,6 @@ def _warn_skipped_steps(self, pred_steps: int, phase: str) -> None: ) self._steps_warn_issued[phase] = True - def _compute_prediction_and_loss( - self, - batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Compute predicted mean, standard deviation, and step-wise loss. - Also extract and return corresponding target from batch. - - Parameters - ---------- - batch : tuple of torch.Tensor - The batch of data. - - Returns - ------- - prediction : torch.Tensor - Model predictions, shape - `(B, pred_steps, num_grid_nodes, num_state_vars)`. - target_states : torch.Tensor - Target states, shape - `(B, pred_steps, num_grid_nodes, num_state_vars)`. - pred_std : torch.Tensor - Predicted or pre-defined standard deviation, shape - `(B, pred_steps, num_grid_nodes, num_state_vars)` or - `(num_state_vars,)`. - time_step_loss : torch.Tensor - Loss for each unroll step, shape `(pred_steps,)`. - """ - prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.per_var_std - assert pred_std is not None - - time_step_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - return ( - prediction, - target_states, - pred_std, - time_step_loss, - ) - def _log_step_loss( self, time_step_loss: torch.Tensor, @@ -515,10 +463,15 @@ def _log_step_loss( batch_size=batch_size, ) + @abstractmethod def validation_step(self, batch, batch_idx): """ Perform a single validation step. + Concrete subclasses must both score the batch and populate + ``self.val_metrics`` for epoch-end aggregation by + ``on_validation_epoch_end``. + Parameters ---------- batch : tuple @@ -526,22 +479,6 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, time_step_loss = ( - self._compute_prediction_and_loss(batch) - ) - mean_loss = torch.mean(time_step_loss) - batch_size = batch[0].shape[0] - - self._log_step_loss(time_step_loss, mean_loss, "val", batch_size) - - entry_mses = metrics.mse( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.val_metrics["mse"].append(entry_mses) def on_validation_epoch_end(self): """ @@ -565,11 +502,15 @@ def on_validation_epoch_end(self): for metric_list in self.val_metrics.values(): metric_list.clear() - # pylint: disable-next=unused-argument + @abstractmethod def test_step(self, batch, batch_idx): """ Perform a single test step. + Concrete subclasses must both score the batch and populate + ``self.test_metrics`` for epoch-end aggregation by + ``on_test_epoch_end``. + Parameters ---------- batch : tuple @@ -577,65 +518,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, time_step_loss = ( - self._compute_prediction_and_loss(batch) - ) - - if self.forecaster.predicts_std: - mean_pred_std = torch.mean( - pred_std[..., self.interior_mask_bool, :], dim=-2 - ) - self.test_metrics["output_std"].append(mean_pred_std) - - mean_loss = torch.mean(time_step_loss) - batch_size = batch[0].shape[0] - - self._log_step_loss(time_step_loss, mean_loss, "test", batch_size) - - hparams = self.hparams - val_steps_to_log = ( - hparams.val_steps_to_log # type: ignore[attr-defined] - ) - - for metric_name in ("mse", "mae"): - metric_func = metrics.get_metric(metric_name) - batch_metric_vals = metric_func( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.test_metrics[metric_name].append(batch_metric_vals) - - spatial_loss = self.loss( - prediction, target_states, pred_std, average_grid=False - ) - log_spatial_losses = spatial_loss[ - :, - [ - step - 1 - for step in val_steps_to_log - if step <= spatial_loss.shape[1] - ], - ] - self.spatial_loss_maps.append(log_spatial_losses) - - if ( - self.trainer.is_global_zero - and self.plotted_examples < self.n_example_pred - ): - n_additional_examples = min( - prediction.shape[0], - self.n_example_pred - self.plotted_examples, - ) - - self.plot_examples( - batch, - n_additional_examples, - prediction=prediction, - split="test", - ) def plot_examples(self, batch, n_examples, split, prediction): """ @@ -924,82 +806,16 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix): plt.close("all") + @abstractmethod def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - Aggregates and plots test metrics and spatial loss maps. - """ - self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") - - spatial_loss_tensor = self.all_gather_cat( - torch.cat(self.spatial_loss_maps, dim=0) - ) - if self.trainer.is_global_zero: - mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) - - loss_map_figs = [ - vis.plot_spatial_error( - error=loss_map, - datastore=self.datastore, - title=f"Test loss, t={t_i} " - f"({(self.time_step_int * t_i)} {self.time_step_unit})", - ) - for t_i, loss_map in zip( - self.hparams.val_steps_to_log, mean_spatial_loss - ) - ] - - for i, fig in enumerate(loss_map_figs): - key = "test_loss" - if not isinstance(self.logger, pl.loggers.WandbLogger): - key = f"{key}_{i}" - if hasattr(self.logger, "log_image"): - self.logger.log_image(key=key, images=[fig]) - - pdf_loss_map_figs = [ - vis.plot_spatial_error(error=loss_map, datastore=self.datastore) - for loss_map in mean_spatial_loss - ] - pdf_loss_maps_dir = os.path.join( - self.logger.save_dir, "spatial_loss_maps" - ) - os.makedirs(pdf_loss_maps_dir, exist_ok=True) - for t_i, fig in zip( - self.hparams.val_steps_to_log, pdf_loss_map_figs - ): - fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) - - torch.save( - mean_spatial_loss.cpu(), - os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), - ) - - if self.hparams.metrics_watch: - unmatched = ( - set(self.hparams.metrics_watch) - self.matched_metrics - ) - if unmatched: - warnings.warn( - "The following metrics in --metrics_watch " - "were not found during test phase: " - f"{sorted(unmatched)}. Ensure the metric prefix " - "matches the evaluation mode (expected 'test_')." - ) - - self.matched_metrics = set() - self.spatial_loss_maps.clear() - # Clear stored test metrics so repeated `trainer.test()` calls on - # the same model instance start from a clean slate (otherwise the - # tensors accumulate and skew the aggregated metrics). - for metric_list in self.test_metrics.values(): - metric_list.clear() - - # Reset the example-plot counter so example prediction plots are - # generated again on every `trainer.test()` call, not just the - # first one (the guard `plotted_examples < n_example_pred` would - # otherwise stay permanently False). - self.plotted_examples = 0 + Concrete subclasses must at least aggregate and plot + ``self.test_metrics`` (typically via ``aggregate_and_plot_metrics``) + and reset any epoch-scoped state they accumulate during + ``test_step``. + """ def on_load_checkpoint(self, checkpoint): """ @@ -1016,15 +832,20 @@ def on_load_checkpoint(self, checkpoint): # 1. Broad namespace remap: for pre-refactor checkpoints # The old ``ARModel`` was a flat LightningModule. Everything that # belonged to the predictor needs to be moved to - # 'forecaster.predictor.' + # 'forecaster.predictor.', while 'per_var_std' (now owned by the + # forecaster itself) moves to 'forecaster.per_var_std' and + # 'interior_mask_bool' (still owned by the module) stays as-is. old_keys = list(loaded_state_dict.keys()) for key in old_keys: - if not key.startswith("forecaster.") and key not in ( - "interior_mask_bool", - "per_var_std", - ): - new_key = f"forecaster.predictor.{key}" - loaded_state_dict[new_key] = loaded_state_dict.pop(key) + if key.startswith("forecaster.") or key == "interior_mask_bool": + continue + if key == "per_var_std": + loaded_state_dict["forecaster.per_var_std"] = ( + loaded_state_dict.pop(key) + ) + continue + new_key = f"forecaster.predictor.{key}" + loaded_state_dict[new_key] = loaded_state_dict.pop(key) # 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp # Will be under forecaster.predictor due to the remap above, or diff --git a/neural_lam/models/modules/deterministic.py b/neural_lam/models/modules/deterministic.py new file mode 100644 index 00000000..817abdc9 --- /dev/null +++ b/neural_lam/models/modules/deterministic.py @@ -0,0 +1,241 @@ +"""Lightning module evaluating forecasters through a single deterministic +rollout per batch.""" + +# Standard library +import os +import warnings +from typing import Any + +# Third-party +import pytorch_lightning as pl +import torch + +# Local +from ... import metrics, vis +from .base import BaseForecasterModule + + +class DeterministicForecasterModule(BaseForecasterModule): + """ + Lightning module for a single deterministic forecast per batch. + + Validation and testing score the forecaster's own single-rollout + prediction via ``forecaster.score``, as opposed to + ``ProbabilisticForecasterModule``, which samples and scores an + ensemble. Training is shared with that module unchanged (see + ``BaseForecasterModule.training_step``). + """ + + def __init__(self, *args, **kwargs): + """ + Initialize the module and its deterministic evaluation metrics. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + **kwargs + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` + (``lr``, ...). + """ + super().__init__(*args, **kwargs) + self.val_metrics: dict[str, list] = { + "mse": [], + } + self.test_metrics: dict[str, list] = { + "mse": [], + "mae": [], + } + if self.forecaster.predicts_std: + self.test_metrics["output_std"] = [] # Treat as metric + + # For storing spatial loss maps during evaluation + self.spatial_loss_maps: list[Any] = [] + + def validation_step(self, batch, batch_idx): + """ + Perform a single validation step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + + time_step_loss = torch.mean( + self.forecaster.score( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + batch_size = batch[0].shape[0] + self._log_step_loss(time_step_loss, mean_loss, "val", batch_size) + + entry_mses = self.forecaster.score( + prediction, + target_states, + pred_std, + metric=metrics.mse, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.val_metrics["mse"].append(entry_mses) + + # pylint: disable-next=unused-argument + def test_step(self, batch, batch_idx): + """ + Perform a single test step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + + if pred_std is not None: + mean_pred_std = torch.mean( + pred_std[..., self.interior_mask_bool, :], dim=-2 + ) + self.test_metrics["output_std"].append(mean_pred_std) + + time_step_loss = torch.mean( + self.forecaster.score( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + batch_size = batch[0].shape[0] + self._log_step_loss(time_step_loss, mean_loss, "test", batch_size) + + for metric_name in ("mse", "mae"): + batch_metric_vals = self.forecaster.score( + prediction, + target_states, + pred_std, + metric=metrics.get_metric(metric_name), + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.test_metrics[metric_name].append(batch_metric_vals) + + spatial_loss = self.forecaster.score( + prediction, target_states, pred_std, average_grid=False + ) + log_spatial_losses = spatial_loss[ + :, + [ + step - 1 + for step in self.hparams.val_steps_to_log + if step <= spatial_loss.shape[1] + ], + ] + self.spatial_loss_maps.append(log_spatial_losses) + + if ( + self.trainer.is_global_zero + and self.plotted_examples < self.n_example_pred + ): + n_additional_examples = min( + prediction.shape[0], + self.n_example_pred - self.plotted_examples, + ) + + self.plot_examples( + batch, + n_additional_examples, + prediction=prediction, + split="test", + ) + + def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + Aggregates and plots test metrics and spatial loss maps. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + spatial_loss_tensor = self.all_gather_cat( + torch.cat(self.spatial_loss_maps, dim=0) + ) + if self.trainer.is_global_zero: + mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) + + loss_map_figs = [ + vis.plot_spatial_error( + error=loss_map, + datastore=self.datastore, + title=f"Test loss, t={t_i} " + f"({(self.time_step_int * t_i)} {self.time_step_unit})", + ) + for t_i, loss_map in zip( + self.hparams.val_steps_to_log, mean_spatial_loss + ) + ] + + for i, fig in enumerate(loss_map_figs): + key = "test_loss" + if not isinstance(self.logger, pl.loggers.WandbLogger): + key = f"{key}_{i}" + if hasattr(self.logger, "log_image"): + self.logger.log_image(key=key, images=[fig]) + + pdf_loss_map_figs = [ + vis.plot_spatial_error(error=loss_map, datastore=self.datastore) + for loss_map in mean_spatial_loss + ] + pdf_loss_maps_dir = os.path.join( + self.logger.save_dir, "spatial_loss_maps" + ) + os.makedirs(pdf_loss_maps_dir, exist_ok=True) + for t_i, fig in zip( + self.hparams.val_steps_to_log, pdf_loss_map_figs + ): + fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) + + torch.save( + mean_spatial_loss.cpu(), + os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), + ) + + if self.hparams.metrics_watch: + unmatched = ( + set(self.hparams.metrics_watch) - self.matched_metrics + ) + if unmatched: + warnings.warn( + "The following metrics in --metrics_watch " + "were not found during test phase: " + f"{sorted(unmatched)}. Ensure the metric prefix " + "matches the evaluation mode (expected 'test_')." + ) + + self.matched_metrics = set() + self.spatial_loss_maps.clear() + + # Clear stored test metrics so repeated `trainer.test()` calls on + # the same model instance start from a clean slate (otherwise the + # tensors accumulate and skew the aggregated metrics). + for metric_list in self.test_metrics.values(): + metric_list.clear() + + # Reset the example-plot counter so example prediction plots are + # generated again on every `trainer.test()` call, not just the + # first one (the guard `plotted_examples < n_example_pred` would + # otherwise stay permanently False). + self.plotted_examples = 0 diff --git a/neural_lam/models/modules/probabilistic.py b/neural_lam/models/modules/probabilistic.py new file mode 100644 index 00000000..83c7d50f --- /dev/null +++ b/neural_lam/models/modules/probabilistic.py @@ -0,0 +1,192 @@ +"""Lightning module evaluating probabilistic forecasters as ensembles.""" + +# Standard library +import warnings + +# Third-party +import torch + +# Local +from ... import metrics +from ..forecasters.probabilistic import ProbabilisticForecaster +from .base import BaseForecasterModule + + +class ProbabilisticForecasterModule(BaseForecasterModule): + """ + Lightning module for forecasters that sample ensemble forecasts. + + Training is inherited unchanged from ``BaseForecasterModule``: the + wrapped forecaster assembles its own training loss. Validation and + testing are ensemble based instead of deterministic: an ensemble is + sampled from the forecaster and scored through its ensemble mean + (root-mean-squared error of the ensemble mean). The module only assumes + that the forecaster can sample ensemble forecasts of the correct shape; + it makes no assumption on how the members are produced. + """ + + # The wrapped forecaster must be able to sample ensemble forecasts + forecaster: ProbabilisticForecaster + + def __init__(self, *args, eval_ensemble_size: int, **kwargs): + """ + Initialize the module and store the evaluation ensemble size. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + eval_ensemble_size : int + Number of ensemble members sampled during validation and + testing. + **kwargs + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` + (``lr``, ...). + """ + super().__init__(*args, **kwargs) + if eval_ensemble_size < 1: + raise ValueError( + "eval_ensemble_size must be at least 1, " + f"got {eval_ensemble_size}" + ) + self.eval_ensemble_size = eval_ensemble_size + self.val_metrics: dict[str, list] = {"ens_mse": []} + self.test_metrics: dict[str, list] = {"ens_mse": []} + + def _ensemble_step(self, batch, phase: str): + """ + Sample an ensemble and score its mean against the target states. + + Shared by ``validation_step`` and ``test_step``: samples + ``self.eval_ensemble_size`` members, scores the ensemble mean with + plain (unweighted) MSE on interior nodes, logs the root-mean-squared + error per configured rollout step and averaged over the rollout + under the given phase's prefix. This RMSE is a diagnostic metric, + not the training loss: it always scores the ensemble mean with + plain MSE, regardless of what objective ``compute_training_loss`` + actually trains on, which is not recomputed here. + + Parameters + ---------- + batch : tuple + The batch of data. + phase : str + Logging phase, either ``"val"`` or ``"test"``. + + Returns + ------- + torch.Tensor + Per-variable ensemble-mean MSE, shape + ``(B, pred_steps, num_state_vars)``, for epoch-end aggregation. + """ + init_states, target_states, forcing_features, _ = batch + ensemble, _ = self.forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=self.eval_ensemble_size, + ) + ensemble_mean = ensemble.mean(dim=1) + # metrics.mse ignores the std argument, but requires one + std_placeholder = torch.ones( + target_states.shape[-1], device=target_states.device + ) + + time_step_mse = torch.mean( + metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + ), + dim=0, + ) + time_step_rmse = torch.sqrt(time_step_mse) + mean_rmse = torch.mean(time_step_rmse) + self._warn_skipped_steps(len(time_step_rmse), phase) + + log_dict = { + f"{phase}_ens_rmse_unroll{step}": time_step_rmse[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_rmse) + } + log_dict[f"{phase}_mean_ens_rmse"] = mean_rmse + self.log_dict( + log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + return metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + sum_vars=False, + ) + + def validation_step(self, batch, batch_idx): + """ + Perform a single ensemble validation step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "val") + self.val_metrics["ens_mse"].append(entry_mses) + + def test_step(self, batch, batch_idx): + """ + Perform a single ensemble test step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "test") + self.test_metrics["ens_mse"].append(entry_mses) + + def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + + Aggregates ensemble test metrics. Implements + ``BaseForecasterModule.on_test_epoch_end`` without the spatial loss + maps and example plots that ``DeterministicForecasterModule`` adds, + since ``test_step`` here does not populate them. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + if self.trainer.is_global_zero and self.hparams.metrics_watch: + unmatched = set(self.hparams.metrics_watch) - self.matched_metrics + if unmatched: + warnings.warn( + "The following metrics in --metrics_watch " + "were not found during test phase: " + f"{sorted(unmatched)}. Ensure the metric prefix " + "matches the evaluation mode (expected 'test_')." + ) + + self.matched_metrics = set() + + for metric_list in self.test_metrics.values(): + metric_list.clear() diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 449f4bb6..1544191f 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -100,24 +100,18 @@ def __init__( # Load graph with static features # NOTE: (IMPORTANT!) mesh nodes MUST have the first # num_mesh_nodes indices, - graph_dir_path = datastore.root_path / "graph" / graph_name grid_xy_extent = datastore.get_xy_extent(category="state") grid_xy_max_span = max( grid_xy_extent[1] - grid_xy_extent[0], grid_xy_extent[3] - grid_xy_extent[2], ) - self.hierarchical, graph_ldict = utils.load_graph( - graph_dir_path=graph_dir_path, + self.hierarchical = utils.load_and_register_graph( + self, + datastore, + graph_name, mesh_node_features_scaling=grid_xy_max_span, ) - for name, attr_value in graph_ldict.items(): - # Make BufferLists module members and register tensors as buffers - if isinstance(attr_value, torch.Tensor): - self.register_buffer(name, attr_value, persistent=False) - else: - setattr(self, name, attr_value) - # Specify dimensions of data self.num_mesh_nodes, _ = self.get_num_mesh() utils.log_on_rank_zero( @@ -126,14 +120,10 @@ def __init__( ) # Compute grid_input_dim: total input dimensionality on the grid - num_state_vars = datastore.get_num_data_vars(category="state") - num_forcing_vars = datastore.get_num_data_vars(category="forcing") - grid_static_dim = self.grid_static_features.shape[1] - self.grid_input_dim = ( - 2 * num_state_vars - + grid_static_dim - + num_forcing_vars - * (num_past_forcing_steps + num_future_forcing_steps + 1) + self.grid_input_dim = utils.compute_grid_input_dim( + datastore, + num_past_forcing_steps, + num_future_forcing_steps, ) self.g2m_edges, g2m_dim = self.g2m_features.shape diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py new file mode 100644 index 00000000..55424b9f --- /dev/null +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -0,0 +1,1079 @@ +"""Graph-based Ensemble Forecasting Model (Graph-EFM) single-step +predictors, for hierarchical (GraphEFM) and flat (GraphEFMMultiScale) mesh +graphs.""" + +# Standard library +from typing import Dict, Optional + +# Third-party +import torch +from torch import nn + +# Local +from .... import utils +from ....datastore import BaseDatastore +from ...latent import ( + ConstantLatentEncoder, + GraphLatentDecoder, + GraphLatentEncoder, + HiGraphLatentDecoder, + HiGraphLatentEncoder, +) +from ..base import StepPredictor + + +class BaseGraphEFM(StepPredictor): + """ + Base class for Graph-based Ensemble Forecasting Model single-step + predictors. + + A latent-variable step predictor consisting of a conditional prior, a + variational encoder and a latent decoder, each of which carries its own + grid-to-mesh, on-mesh and mesh-to-grid GNNs. The + encode-process-decode backbone of + ``BaseGraphModel`` therefore does not apply -- this extends + ``StepPredictor`` directly. ``forward`` samples a single step from the + prior. The encoder (variational posterior) and prior are exposed as + ``self.encoder``/``self.prior_model`` for a forecaster to condition on + the target and assemble a training objective from; the predictor itself + does not compute any loss. + + This base class sets up everything that is independent of the mesh + graph type: it loads the graph, calls the subclass's + :meth:`check_graph_type` to verify the loaded graph matches what it + requires, then builds the prior, delegating to the subclass's + :meth:`build_learnable_prior` when learned and to + :attr:`latent_spatial_dim` for the constant prior's node count. + Concrete subclasses build the mesh embedders and the encoder/decoder + latent modules, and implement :meth:`embedd_mesh`, + :meth:`check_graph_type`, :meth:`build_learnable_prior` and + :attr:`latent_spatial_dim`. See :class:`GraphEFM` (hierarchical graph) + and :class:`GraphEFMMultiScale` (flat graph). + """ + + def __init__( + self, + datastore: BaseDatastore, + graph_name: str, + hidden_dim: int = 64, + hidden_layers: int = 1, + latent_dim: Optional[int] = None, + learn_prior: bool = True, + prior_dist: str = "isotropic", + prior_layers: int = 2, + g2m_gnn_type: str = "InteractionNet", + num_past_forcing_steps: int = 1, + num_future_forcing_steps: int = 1, + output_std: bool = False, + output_clamping_lower: Optional[Dict[str, float]] = None, + output_clamping_upper: Optional[Dict[str, float]] = None, + ): + """ + Set up the graph-type independent parts of the predictor. + + Loads the graph, builds the grid embedders, the grid-mesh edge + embedders and the prior. Building the mesh embedders and the + encoder/decoder latent modules is left to the subclass constructor. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing static features, standardization statistics + and variable counts. + graph_name : str + Name of the graph directory (under ``/graph``) to load. + Must be of the graph type required by the concrete subclass. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + latent_dim : int, optional + Dimensionality of the latent variable at each latent-carrying + mesh node (top level for hierarchical graphs, all mesh nodes + for flat graphs); defaults to ``hidden_dim`` when None. The + resolved value is stored as ``self.latent_dim`` for the + subclass to reuse when building its encoder/decoder. + learn_prior : bool + If True, the prior is the graph-type specific learnable encoder + built by :meth:`build_learnable_prior`, conditioned on the + previous state; if False, a constant ``Normal(0, 1)`` prior is + used. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh GNN layers in the learnable prior. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior (key in + ``gnn_layers.GNN_TYPES``). + num_past_forcing_steps : int + Number of past forcing steps included in the input window. + num_future_forcing_steps : int + Number of future forcing steps included in the input window. + output_std : bool + If True, the decoder outputs a per-variable std alongside the + mean; if False, ``forward`` returns ``None`` for the std. + output_clamping_lower : dict of str to float, optional + Lower clamping limits per output variable. + output_clamping_upper : dict of str to float, optional + Upper clamping limits per output variable. + """ + super().__init__( + datastore=datastore, + output_std=output_std, + output_clamping_lower=output_clamping_lower, + output_clamping_upper=output_clamping_upper, + ) + + # Load graph with static features. + grid_xy_extent = datastore.get_xy_extent(category="state") + grid_xy_max_span = max( + grid_xy_extent[1] - grid_xy_extent[0], + grid_xy_extent[3] - grid_xy_extent[2], + ) + self.hierarchical = utils.load_and_register_graph( + self, + datastore, + graph_name, + mesh_node_features_scaling=grid_xy_max_span, + ) + # Delegated to the subclass, which knows whether it requires a + # hierarchical or flat mesh graph. Must run before anything below + # that assumes a specific graph shape (build_learnable_prior, + # latent_spatial_dim), so it cannot wait until the subclass + # constructor resumes after this call returns. + self.check_graph_type(graph_name) + + # Specify dimensions of data + self.num_state_vars = datastore.get_num_data_vars(category="state") + num_state_vars = self.num_state_vars + # grid_dim: total grid input dim. grid_current_dim additionally + # includes the target state, for the encoder input. + self.grid_dim = utils.compute_grid_input_dim( + datastore, + num_past_forcing_steps, + num_future_forcing_steps, + ) + grid_current_dim = self.grid_dim + num_state_vars + g2m_dim = self.g2m_features.shape[1] + m2g_dim = self.m2g_features.shape[1] + + # Define sub-models + # Feature embedders for grid + self.mlp_blueprint_end = [hidden_dim] * (hidden_layers + 1) + self.grid_prev_embedder = utils.make_mlp( + [self.grid_dim] + self.mlp_blueprint_end + ) # For states up to t-1 + self.grid_current_embedder = utils.make_mlp( + [grid_current_dim] + self.mlp_blueprint_end + ) # For states including t + # Embedders for mesh edges + self.g2m_embedder = utils.make_mlp([g2m_dim] + self.mlp_blueprint_end) + self.m2g_embedder = utils.make_mlp([m2g_dim] + self.mlp_blueprint_end) + + # Compute indices and define the clamping functions applied to the + # predicted next-state mean in step_distributions. + self.prepare_clamping_params(datastore) + + # Prior over the latent variable. When learn_prior is True the + # (graph-type specific) learnable prior is delegated to + # build_learnable_prior; otherwise the constant Normal(0, 1) prior, + # identical for every graph type, is built directly here. + self.latent_dim = latent_dim if latent_dim is not None else hidden_dim + if learn_prior: + self.prior_model = self.build_learnable_prior( + latent_dim=self.latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_layers, + ) + else: + self.prior_model = ConstantLatentEncoder( + latent_dim=self.latent_dim, + num_mesh_nodes=self.latent_spatial_dim, + output_dist=prior_dist, + ) + + def check_graph_type(self, graph_name: str) -> None: + """ + Verify the loaded graph (``self.hierarchical``) is of the type this + predictor requires. + + Implemented by the concrete subclass, which is the only place that + knows whether it requires a hierarchical or flat mesh graph. Called + by this base class right after loading the graph, before anything + that assumes a specific graph shape. + + Parameters + ---------- + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If ``self.hierarchical`` does not match what this predictor + requires. + """ + raise NotImplementedError("check_graph_type not implemented") + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes the latent variable lives on. + + Implemented by the concrete subclass, which knows the mesh graph + type: the top mesh level for hierarchical graphs, or every mesh + node for flat graphs. + + Returns + ------- + int + Number of latent-carrying mesh nodes. + """ + raise NotImplementedError("latent_spatial_dim not implemented") + + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the graph-type specific learnable prior encoder. + + Implemented by the concrete subclass, which knows the mesh graph type + and therefore the appropriate latent encoder class. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh GNN layers in the prior. + + Returns + ------- + torch.nn.Module + The learnable prior latent encoder. + """ + raise NotImplementedError("build_learnable_prior not implemented") + + def embedd_grid_with_target( + self, + prev_state, + prev_prev_state, + forcing, + current_state, + ): + """ + Embed the grid representation including the current (target) state. + Used as input to the encoder, which is conditioned also on the target. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + current_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t+1}`` (target). + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid embedding. + """ + batch_size = prev_state.shape[0] + + grid_current_features = torch.cat( + ( + prev_prev_state, + prev_state, + forcing, + self.expand_to_batch(self.grid_static_features, batch_size), + current_state, + ), + dim=-1, + ) # (B, num_grid_nodes, grid_current_dim) + + return self.grid_current_embedder( + grid_current_features + ) # (B, num_grid_nodes, d_h) + + def embedd_mesh(self, batch_size): + """ + Embed static mesh node and intra-mesh edge features. + + Parameters + ---------- + batch_size : int + Batch size to expand the embeddings to. + + Returns + ------- + dict + Mesh-related entries of the graph embedding (``mesh``, ``m2m`` + and, for hierarchical graphs, ``mesh_up`` and ``mesh_down``). + Entries are tensors of shape ``(B, *, d_h)`` for flat graphs + and per-level lists of such tensors for hierarchical graphs. + """ + raise NotImplementedError("embedd_mesh not implemented") + + def embedd_grid(self, prev_state, prev_prev_state, forcing): + """ + Embed the grid representation of the states up to t-1. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid embedding. + """ + batch_size = prev_state.shape[0] + + grid_features = torch.cat( + ( + prev_prev_state, + prev_state, + forcing, + self.expand_to_batch(self.grid_static_features, batch_size), + ), + dim=-1, + ) # (B, num_grid_nodes, grid_dim) + + return self.grid_prev_embedder(grid_features) + # (B, num_grid_nodes, d_h) + + def embedd_graph(self, batch_size): + """ + Embed the static grid-mesh edge and mesh graph features. + + The embedding depends only on static graph features and is therefore + constant across an autoregressive rollout, unlike :meth:`embedd_grid`. + + Parameters + ---------- + batch_size : int + Batch size to expand the embeddings to. + + Returns + ------- + dict + Edge/mesh embeddings, each entry of shape ``(B, *, d_h)`` (``g2m``, + ``m2g`` and the entries added by :meth:`embedd_mesh`). + """ + graph_emb = { + "g2m": self.expand_to_batch( + self.g2m_embedder(self.g2m_features), batch_size + ), # (B, M_g2m, d_h) + "m2g": self.expand_to_batch( + self.m2g_embedder(self.m2g_features), batch_size + ), # (B, M_m2g, d_h) + } + graph_emb.update(self.embedd_mesh(batch_size)) + + return graph_emb + + def step_distributions( + self, + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=None, + compute_prior=True, + ): + """ + Compute the latent distributions and next-state prediction for a step. + + Embeds the grid, then draws the latent either from the variational + posterior (when ``target_state`` is given -- the training path) or + from the prior (inference), and decodes that sample into the + next-state mean (and, when ``output_std``, std) as a residual on + ``prev_state``. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + graph_emb : dict + Static graph embedding as returned by :meth:`embedd_graph`. + target_state : torch.Tensor, optional + Shape ``(B, num_grid_nodes, d_state)``. Target ``X_{t+1}``. When + given, the latent is sampled from the variational posterior + conditioned on it; when None, from the prior. + compute_prior : bool + On the posterior (training) path, whether to also compute the + prior distribution (needed for a KL term). Ignored on the + inference path, where the prior is always computed since it is the + sampling distribution. + + Returns + ------- + prior_dist : torch.distributions.Normal or None + The prior over the latent, or None on the posterior path when + ``compute_prior`` is False. + posterior_dist : torch.distributions.Normal or None + The variational posterior over the latent, or None on the + inference path (``target_state`` is None). + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Decoder mean of + ``X_{t+1}``. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, d_state)`` when ``output_std`` is True, + otherwise None. + """ + grid_prev_emb = self.embedd_grid(prev_state, prev_prev_state, forcing) + + prior_dist = None + posterior_dist = None + if target_state is not None: + grid_current_emb = self.embedd_grid_with_target( + prev_state, prev_prev_state, forcing, target_state + ) + posterior_dist = self.encoder(grid_current_emb, graph_emb=graph_emb) + latent_samples = posterior_dist.rsample() + if compute_prior: + prior_dist = self.prior_model( + grid_prev_emb, graph_emb=graph_emb + ) + else: + prior_dist = self.prior_model(grid_prev_emb, graph_emb=graph_emb) + latent_samples = prior_dist.rsample() + # (B, num_mesh_nodes, d_latent) + + # Decode the latent into a state increment, then add it onto prev_state + # (X_t) and clamp to the valid range (a no-op when no clamping limits + # are configured), as for the deterministic models. + mean_delta, pred_std = self.decoder( + grid_prev_emb, latent_samples, graph_emb + ) + pred_mean = self.get_clamped_new_state(mean_delta, prev_state) + # (B, num_grid_nodes, d_state) + + return prior_dist, posterior_dist, pred_mean, pred_std + + def forward( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Sample one time step prediction from the prior. + + Embeds the graph and grid, samples the latent from the prior, decodes + it and returns the predicted next state. The prediction is stochastic + only through the latent sample; no observation noise is added. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + + Returns + ------- + new_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Predicted ``X_{t+1}`` + (the decoder mean, given the sampled latent). + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, d_state)`` when ``output_std`` is True, + otherwise None. + """ + graph_emb = self.embedd_graph(prev_state.shape[0]) + _, _, pred_mean, pred_std = self.step_distributions( + prev_state, prev_prev_state, forcing, graph_emb, target_state=None + ) + return pred_mean, pred_std + + +class GraphEFM(BaseGraphEFM): + """ + Graph-based Ensemble Forecasting Model on a hierarchical mesh graph. + + The latent variable lives on the top level of the mesh hierarchy. The + prior and variational encoder are ``HiGraphLatentEncoder``s and the + decoder is a ``HiGraphLatentDecoder``. + """ + + def __init__( + self, + datastore: BaseDatastore, + graph_name: str = "hierarchical", + hidden_dim: int = 64, + hidden_layers: int = 1, + latent_dim: Optional[int] = None, + prior_intra_level_layers: int = 2, + encoder_intra_level_layers: int = 2, + decoder_intra_level_layers: int = 4, + learn_prior: bool = True, + prior_dist: str = "isotropic", + num_past_forcing_steps: int = 1, + num_future_forcing_steps: int = 1, + g2m_gnn_type: str = "InteractionNet", + m2g_gnn_type: str = "InteractionNet", + output_std: bool = False, + output_clamping_lower: Optional[Dict[str, float]] = None, + output_clamping_upper: Optional[Dict[str, float]] = None, + ): + """ + Build the mesh embedders and the hierarchical encoder/decoder + latent modules. The prior is built by the base class. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing static features, standardization statistics + and variable counts. + graph_name : str + Name of the graph directory (under ``/graph``) to load. + Must be a hierarchical graph. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + latent_dim : int, optional + Dimensionality of the latent variable at each top-level mesh + node; defaults to ``hidden_dim`` when None. Forwarded to the + base class, which resolves the default and stores it as + ``self.latent_dim``. + prior_intra_level_layers : int + Number of intra-level GNN layers in the (learned) prior. + Forwarded to the base class as ``prior_layers``. + encoder_intra_level_layers : int + Number of intra-level GNN layers in the variational encoder. + decoder_intra_level_layers : int + Number of intra-level GNN layers in the latent decoder. + learn_prior : bool + If True, the prior is a hierarchical graph encoder conditioned + on the previous state; if False, a constant ``Normal(0, 1)`` + prior is used. Forwarded to the base class. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. Forwarded to the base class. + num_past_forcing_steps : int + Number of past forcing steps included in the input window. + num_future_forcing_steps : int + Number of future forcing steps included in the input window. + g2m_gnn_type : str + GNN type for the grid-to-mesh steps of the prior, encoder and + decoder (key in ``gnn_layers.GNN_TYPES``). + m2g_gnn_type : str + GNN type for the mesh-to-grid step of the decoder (key in + ``gnn_layers.GNN_TYPES``). + output_std : bool + If True, the decoder outputs a per-variable std alongside the + mean; if False, ``forward`` returns ``None`` for the std. + output_clamping_lower : dict of str to float, optional + Lower clamping limits per output variable. + output_clamping_upper : dict of str to float, optional + Upper clamping limits per output variable. + """ + super().__init__( + datastore=datastore, + graph_name=graph_name, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + latent_dim=latent_dim, + learn_prior=learn_prior, + prior_dist=prior_dist, + prior_layers=prior_intra_level_layers, + g2m_gnn_type=g2m_gnn_type, + num_past_forcing_steps=num_past_forcing_steps, + num_future_forcing_steps=num_future_forcing_steps, + output_std=output_std, + output_clamping_lower=output_clamping_lower, + output_clamping_upper=output_clamping_upper, + ) + + level_mesh_sizes = [ + mesh_feat.shape[0] for mesh_feat in self.mesh_static_features + ] + num_levels = len(self.mesh_static_features) + utils.log_on_rank_zero("Loaded hierarchical graph with structure:") + for level_index, level_mesh_size in enumerate(level_mesh_sizes): + same_level_edges = self.m2m_features[level_index].shape[0] + utils.log_on_rank_zero( + f"level {level_index} - {level_mesh_size} nodes, " + f"{same_level_edges} same-level edges" + ) + if level_index < (num_levels - 1): + up_edges = self.mesh_up_features[level_index].shape[0] + down_edges = self.mesh_down_features[level_index].shape[0] + utils.log_on_rank_zero(f" {level_index}<->{level_index + 1}") + utils.log_on_rank_zero( + f" - {up_edges} up edges, {down_edges} down edges" + ) + + # Embedders. Assume all levels share static feature dimensionality. + mesh_dim = self.mesh_static_features[0].shape[1] + m2m_dim = self.m2m_features[0].shape[1] + mesh_up_dim = self.mesh_up_features[0].shape[1] + mesh_down_dim = self.mesh_down_features[0].shape[1] + + # Separate mesh node embedders for each level + self.mesh_embedders = nn.ModuleList( + [ + utils.make_mlp([mesh_dim] + self.mlp_blueprint_end) + for _ in range(num_levels) + ] + ) + self.mesh_up_embedders = nn.ModuleList( + [ + utils.make_mlp([mesh_up_dim] + self.mlp_blueprint_end) + for _ in range(num_levels - 1) + ] + ) + self.mesh_down_embedders = nn.ModuleList( + [ + utils.make_mlp([mesh_down_dim] + self.mlp_blueprint_end) + for _ in range(num_levels - 1) + ] + ) + # If not using any intra-level layers, no need to embed m2m + self.embedd_m2m = ( + max( + prior_intra_level_layers, + encoder_intra_level_layers, + decoder_intra_level_layers, + ) + > 0 + ) + if self.embedd_m2m: + self.m2m_embedders = nn.ModuleList( + [ + utils.make_mlp([m2m_dim] + self.mlp_blueprint_end) + for _ in range(num_levels) + ] + ) + + # Encoder (variational posterior) + Decoder + self.encoder = HiGraphLatentEncoder( + latent_dim=self.latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + mesh_up_edge_index=self.mesh_up_edge_index, + hidden_dim=hidden_dim, + intra_level_layers=encoder_intra_level_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist="diagonal", + ) + self.decoder = HiGraphLatentDecoder( + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + m2g_edge_index=self.m2g_edge_index, + mesh_up_edge_index=self.mesh_up_edge_index, + mesh_down_edge_index=self.mesh_down_edge_index, + hidden_dim=hidden_dim, + latent_dim=self.latent_dim, + num_state_vars=self.num_state_vars, + intra_level_layers=decoder_intra_level_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + m2g_gnn_type=m2g_gnn_type, + output_std=bool(output_std), + ) + + def check_graph_type(self, graph_name: str) -> None: + """ + Verify the loaded graph is hierarchical. + + Parameters + ---------- + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If the loaded graph is flat. + """ + if not self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a hierarchical mesh " + f"graph, but graph '{graph_name}' is flat" + ) + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes on the top mesh level, where the latent + variable lives. + + Returns + ------- + int + Number of top-level mesh nodes. + """ + return self.mesh_static_features[-1].shape[0] + + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the hierarchical learnable prior encoder. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each top-level mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of intra-level GNN layers in the prior. + + Returns + ------- + HiGraphLatentEncoder + The learnable prior latent encoder. + """ + return HiGraphLatentEncoder( + latent_dim=latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + mesh_up_edge_index=self.mesh_up_edge_index, + hidden_dim=hidden_dim, + intra_level_layers=prior_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + + def embedd_mesh(self, batch_size): + """ + Embed static mesh node and intra-mesh edge features per level. + + Parameters + ---------- + batch_size : int + Batch size to expand the embeddings to. + + Returns + ------- + dict + Entries ``mesh``, ``m2m``, ``mesh_up`` and ``mesh_down``, each + a list with one ``(B, *, d_h)`` tensor per mesh level (or + inter-level connection). + """ + mesh_emb = { + "mesh": [ + self.expand_to_batch(emb(node_static_features), batch_size) + for emb, node_static_features in zip( + self.mesh_embedders, + self.mesh_static_features, + ) + ], # each (B, num_mesh_nodes[l], d_h) + "mesh_up": [ + self.expand_to_batch(emb(edge_feat), batch_size) + for emb, edge_feat in zip( + self.mesh_up_embedders, self.mesh_up_features + ) + ], + "mesh_down": [ + self.expand_to_batch(emb(edge_feat), batch_size) + for emb, edge_feat in zip( + self.mesh_down_embedders, self.mesh_down_features + ) + ], + } + + if self.embedd_m2m: + mesh_emb["m2m"] = [ + self.expand_to_batch(emb(edge_feat), batch_size) + for emb, edge_feat in zip(self.m2m_embedders, self.m2m_features) + ] + else: + # No intra-level GNNs consume these, so no embedding is produced + mesh_emb["m2m"] = [] + + return mesh_emb + + +class GraphEFMMultiScale(BaseGraphEFM): + """ + Graph-based Ensemble Forecasting Model on a flat mesh graph + (e.g. a multi-scale graph). + + The latent variable lives on the mesh nodes. The prior and variational + encoder are ``GraphLatentEncoder``s and the decoder is a + ``GraphLatentDecoder``. + """ + + def __init__( + self, + datastore: BaseDatastore, + graph_name: str = "multiscale", + hidden_dim: int = 64, + hidden_layers: int = 1, + latent_dim: Optional[int] = None, + prior_m2m_layers: int = 2, + encoder_m2m_layers: int = 2, + decoder_m2m_layers: int = 4, + learn_prior: bool = True, + prior_dist: str = "isotropic", + num_past_forcing_steps: int = 1, + num_future_forcing_steps: int = 1, + g2m_gnn_type: str = "InteractionNet", + m2g_gnn_type: str = "InteractionNet", + output_std: bool = False, + output_clamping_lower: Optional[Dict[str, float]] = None, + output_clamping_upper: Optional[Dict[str, float]] = None, + ): + """ + Build the mesh embedders and the flat-graph encoder/decoder latent + modules. The prior is built by the base class. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing static features, standardization statistics + and variable counts. + graph_name : str + Name of the graph directory (under ``/graph``) to load. + Must be a flat graph. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + latent_dim : int, optional + Dimensionality of the latent variable at each mesh node; + defaults to ``hidden_dim`` when None. Forwarded to the base + class, which resolves the default and stores it as + ``self.latent_dim``. + prior_m2m_layers : int + Number of on-mesh (m2m) GNN layers in the (learned) prior. + Forwarded to the base class as ``prior_layers``. + encoder_m2m_layers : int + Number of on-mesh (m2m) GNN layers in the variational encoder. + decoder_m2m_layers : int + Number of on-mesh (m2m) GNN layers in the latent decoder. + learn_prior : bool + If True, the prior is a graph encoder conditioned on the + previous state; if False, a constant ``Normal(0, 1)`` prior is + used. Forwarded to the base class. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. Forwarded to the base class. + num_past_forcing_steps : int + Number of past forcing steps included in the input window. + num_future_forcing_steps : int + Number of future forcing steps included in the input window. + g2m_gnn_type : str + GNN type for the grid-to-mesh steps of the prior, encoder and + decoder (key in ``gnn_layers.GNN_TYPES``). + m2g_gnn_type : str + GNN type for the mesh-to-grid step of the decoder (key in + ``gnn_layers.GNN_TYPES``). + output_std : bool + If True, the decoder outputs a per-variable std alongside the + mean; if False, ``forward`` returns ``None`` for the std. + output_clamping_lower : dict of str to float, optional + Lower clamping limits per output variable. + output_clamping_upper : dict of str to float, optional + Upper clamping limits per output variable. + """ + super().__init__( + datastore=datastore, + graph_name=graph_name, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + latent_dim=latent_dim, + learn_prior=learn_prior, + prior_dist=prior_dist, + prior_layers=prior_m2m_layers, + g2m_gnn_type=g2m_gnn_type, + num_past_forcing_steps=num_past_forcing_steps, + num_future_forcing_steps=num_future_forcing_steps, + output_std=output_std, + output_clamping_lower=output_clamping_lower, + output_clamping_upper=output_clamping_upper, + ) + + utils.log_on_rank_zero( + f"Loaded graph with " + f"{self.num_grid_nodes + self.latent_spatial_dim} nodes " + f"({self.num_grid_nodes} grid, {self.latent_spatial_dim} mesh)" + ) + + # Embedders + mesh_static_dim = self.mesh_static_features.shape[1] + self.mesh_embedder = utils.make_mlp( + [mesh_static_dim] + self.mlp_blueprint_end + ) + m2m_dim = self.m2m_features.shape[1] + self.m2m_embedder = utils.make_mlp([m2m_dim] + self.mlp_blueprint_end) + + # Encoder (variational posterior) + Decoder + self.encoder = GraphLatentEncoder( + latent_dim=self.latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + hidden_dim=hidden_dim, + m2m_layers=encoder_m2m_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist="diagonal", + ) + self.decoder = GraphLatentDecoder( + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + m2g_edge_index=self.m2g_edge_index, + hidden_dim=hidden_dim, + latent_dim=self.latent_dim, + num_state_vars=self.num_state_vars, + m2m_layers=decoder_m2m_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + m2g_gnn_type=m2g_gnn_type, + output_std=bool(output_std), + ) + + def check_graph_type(self, graph_name: str) -> None: + """ + Verify the loaded graph is flat. + + Parameters + ---------- + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If the loaded graph is hierarchical. + """ + if self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a flat mesh graph, " + f"but graph '{graph_name}' is hierarchical" + ) + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes, where the latent variable lives. + + Returns + ------- + int + Number of mesh nodes. + """ + return len(self.mesh_static_features) + + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the flat-graph learnable prior encoder. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh (m2m) GNN layers in the prior. + + Returns + ------- + GraphLatentEncoder + The learnable prior latent encoder. + """ + return GraphLatentEncoder( + latent_dim=latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + hidden_dim=hidden_dim, + m2m_layers=prior_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + + def embedd_mesh(self, batch_size): + """ + Embed static mesh node and intra-mesh edge features. + + Parameters + ---------- + batch_size : int + Batch size to expand the embeddings to. + + Returns + ------- + dict + Entries ``mesh``: ``(B, num_mesh_nodes, d_h)`` and + ``m2m``: ``(B, M_m2m, d_h)``. + """ + return { + "mesh": self.expand_to_batch( + self.mesh_embedder(self.mesh_static_features), batch_size + ), # (B, num_mesh_nodes, d_h) + "m2m": self.expand_to_batch( + self.m2m_embedder(self.m2m_features), batch_size + ), # (B, M_m2m, d_h) + } diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 2a9775c9..80590688 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -19,7 +19,14 @@ 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, + PROBABILISTIC_MODELS, + ARForecaster, + DeterministicForecasterModule, + GraphEFMForecaster, + ProbabilisticForecasterModule, +) from .weather_dataset import WeatherDataModule @@ -38,37 +45,155 @@ def __init__(self, prog): ) -def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): +def build_predictor(args, config, datastore): """ - Reconstruct a ForecasterModule 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 - can be recovered automatically. + Construct the step predictor for ``args.model``. + + Graph-EFM models (see ``PROBABILISTIC_MODELS``) are latent-variable + predictors whose constructor differs from the deterministic graph + models: they take ``latent_dim``/``learn_prior``/``prior_dist`` and + per-graph-type latent layer counts instead of ``processor_layers`` and + ``mesh_aggr``. This function supplies the right kwargs for each. + + Parameters + ---------- + args : argparse.Namespace + Parsed command-line arguments (see ``main``). + config : NeuralLAMConfig + Loaded neural-lam configuration, for the output-clamping limits. + datastore : BaseDatastore + Datastore providing static features and variable counts. + + Returns + ------- + StepPredictor + The constructed step predictor. """ - ckpt = torch.load(ckpt_path, weights_only=False) - args = ckpt["hyper_parameters"]["args"] predictor_class = MODELS[args.model] - predictor = predictor_class( + common_kwargs = dict( datastore=datastore, graph_name=args.graph, hidden_dim=args.hidden_dim, hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, num_past_forcing_steps=args.num_past_forcing_steps, num_future_forcing_steps=args.num_future_forcing_steps, output_std=args.output_std, output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, + g2m_gnn_type=args.g2m_gnn_type, + m2g_gnn_type=args.m2g_gnn_type, + ) + + if args.model in PROBABILISTIC_MODELS: + prob = config.probabilistic + # Graph-EFM latent layer counts are named per graph type: + # intra-level for the hierarchical model, m2m for the flat one. + if args.model == "graph_efm": + layer_kwargs = dict( + prior_intra_level_layers=prob.prior_layers, + encoder_intra_level_layers=prob.encoder_layers, + decoder_intra_level_layers=prob.decoder_layers, + ) + else: + layer_kwargs = dict( + prior_m2m_layers=prob.prior_layers, + encoder_m2m_layers=prob.encoder_layers, + decoder_m2m_layers=prob.decoder_layers, + ) + return predictor_class( + **common_kwargs, + latent_dim=prob.latent_dim, + learn_prior=prob.learn_prior, + prior_dist=prob.prior_dist, + **layer_kwargs, + ) + + return predictor_class( + **common_kwargs, + processor_layers=args.processor_layers, + mesh_aggr=args.mesh_aggr, + mesh_up_gnn_type=args.mesh_up_gnn_type, + mesh_down_gnn_type=args.mesh_down_gnn_type, ) - forecaster = ARForecaster(predictor, datastore) - return ForecasterModule.load_from_checkpoint( + + +def build_forecaster_module(args, config, datastore, predictor): + """ + Wrap a predictor in its forecaster and pick the Lightning module class. + + Graph-EFM models are trained probabilistically (``GraphEFMForecaster`` + optimizing the ELBO, evaluated as an ensemble by + ``ProbabilisticForecasterModule``); the other models use the + deterministic ``ARForecaster``/``DeterministicForecasterModule`` path. + The best-checkpoint monitor differs accordingly. + + Parameters + ---------- + args : argparse.Namespace + Parsed command-line arguments (see ``main``). + config : NeuralLAMConfig + Loaded neural-lam configuration. + datastore : BaseDatastore + Datastore providing grid metadata and boundary masks. + predictor : StepPredictor + The step predictor to wrap, as built by ``build_predictor``. + + Returns + ------- + forecaster : Forecaster + The forecaster wrapping ``predictor``. + module_class : type + The ``BaseForecasterModule`` subclass to instantiate. + module_kwargs : dict + Extra keyword arguments for ``module_class`` beyond the shared ones + (e.g. ``eval_ensemble_size`` for the probabilistic module). + val_monitor : str + Name of the validation metric to monitor for the best checkpoint. + """ + if args.model in PROBABILISTIC_MODELS: + forecaster = GraphEFMForecaster( + predictor, + datastore, + config=config, + loss=args.loss, + kl_beta=config.probabilistic.kl_beta, + ) + return ( + forecaster, + ProbabilisticForecasterModule, + {"eval_ensemble_size": config.probabilistic.eval_ensemble_size}, + "val_mean_ens_rmse", + ) + + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) + return forecaster, DeterministicForecasterModule, {}, "val_mean_loss" + + +def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): + """ + Reconstruct a forecaster module 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 + can be recovered automatically. Deterministic and Graph-EFM + (probabilistic) checkpoints are both supported; the correct forecaster + and module class are chosen from ``args.model``. + """ + ckpt = torch.load(ckpt_path, weights_only=False) + args = ckpt["hyper_parameters"]["args"] + predictor = build_predictor(args, config, datastore) + forecaster, module_class, module_kwargs, _ = build_forecaster_module( + args, config, datastore, predictor + ) + return module_class.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, weights_only=False, + **module_kwargs, ) @@ -181,31 +306,41 @@ def main(input_args=None): type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for grid-to-mesh encoding", + help="GNN type for grid-to-mesh encoding. Applies to all models, " + "including the probabilistic Graph-EFM model", ) arch_group.add_argument( "--m2g_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for mesh-to-grid decoding", + help="GNN type for mesh-to-grid decoding. Applies to all models, " + "including the probabilistic Graph-EFM model", ) arch_group.add_argument( "--mesh_up_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for upward mesh message passing in hierarchical models", + help="GNN type for upward mesh message passing in hierarchical " + "models. Only affects Hi-LAM; the probabilistic Graph-EFM model " + "hard-codes its mesh-up GNN types", ) arch_group.add_argument( "--mesh_down_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for downward mesh message passing in " - "hierarchical models", + help="GNN type for downward mesh message passing in hierarchical " + "models. Only affects Hi-LAM; the probabilistic Graph-EFM model " + "hard-codes its mesh-down GNN type", ) + # Probabilistic / Graph-EFM hyperparameters (latent_dim, kl_beta, + # eval_ensemble_size, ...) are not CLI flags: they live in the + # ``probabilistic`` section of the neural-lam config (see + # ``ProbabilisticConfig``), read by build_predictor/build_forecaster_module. + # Training options train_group = parser.add_argument_group("Training Options") train_group.add_argument( @@ -454,33 +589,18 @@ def main(input_args=None): except ValueError: raise ValueError("devices should be 'auto' or a list of integers") - # Build predictor and forecaster externally, then inject into - # ForecasterModule - predictor_class = MODELS[args.model] - predictor = predictor_class( - datastore=datastore, - graph_name=args.graph, - hidden_dim=args.hidden_dim, - hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, - num_past_forcing_steps=args.num_past_forcing_steps, - num_future_forcing_steps=args.num_future_forcing_steps, - output_std=args.output_std, - output_clamping_lower=config.training.output_clamping.lower, - output_clamping_upper=config.training.output_clamping.upper, - g2m_gnn_type=args.g2m_gnn_type, - m2g_gnn_type=args.m2g_gnn_type, - mesh_up_gnn_type=args.mesh_up_gnn_type, - mesh_down_gnn_type=args.mesh_down_gnn_type, + # Build predictor and forecaster externally, then inject into the + # forecaster module. Graph-EFM models take a probabilistic assembly path + # (see build_predictor/build_forecaster_module); the others deterministic. + predictor = build_predictor(args, config, datastore) + forecaster, module_class, module_kwargs, val_monitor = ( + build_forecaster_module(args, config, datastore, predictor) ) - forecaster = ARForecaster(predictor, datastore) - model = ForecasterModule( + model = module_class( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, lr=args.lr, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, @@ -490,6 +610,7 @@ def main(input_args=None): metrics_watch=args.metrics_watch, var_leads_metrics_watch=args.var_leads_metrics_watch, args=args, + **module_kwargs, ) if args.eval: @@ -517,8 +638,8 @@ def main(input_args=None): # checkpoint instead of losing all progress since the last validation. val_checkpoint = pl.callbacks.ModelCheckpoint( dirpath=os.path.join(run_dir, "checkpoints"), - filename="min_val_loss", - monitor="val_mean_loss", + filename=f"min_{val_monitor}", + monitor=val_monitor, mode="min", save_top_k=1, save_on_train_epoch_end=False, @@ -532,11 +653,19 @@ def main(input_args=None): save_on_train_epoch_end=True, enable_version_counter=False, ) + # With kl_beta == 0 the Graph-EFM prior network is never used in the loss, + # so multi-device DDP must be told to expect unused parameters. + strategy = ( + "ddp_find_unused_parameters_true" + if args.model in PROBABILISTIC_MODELS + and config.probabilistic.kl_beta == 0 + else "auto" + ) trainer = pl.Trainer( max_epochs=args.epochs, deterministic=True, default_root_dir=run_dir, - strategy="auto", + strategy=strategy, accelerator=device_name, num_nodes=args.num_nodes, devices=devices, diff --git a/neural_lam/utils/__init__.py b/neural_lam/utils/__init__.py index 7ac299d5..96a2ed20 100644 --- a/neural_lam/utils/__init__.py +++ b/neural_lam/utils/__init__.py @@ -3,6 +3,8 @@ # Local from .buffer_list import BufferList from .graph import ( + compute_grid_input_dim, + load_and_register_graph, load_graph, zero_index_edge_index, zero_index_g2m, @@ -13,21 +15,24 @@ log_on_rank_zero, setup_training_logger, ) -from .networks import make_mlp +from .networks import make_gnn_seq, make_mlp from .plot import fractional_plot_bundle, has_working_latex from .tensor import inverse_sigmoid, inverse_softplus from .time import get_integer_time __all__ = [ "BufferList", + "compute_grid_input_dim", "fractional_plot_bundle", "get_integer_time", "has_working_latex", "init_training_logger_metrics", "inverse_sigmoid", "inverse_softplus", + "load_and_register_graph", "load_graph", "log_on_rank_zero", + "make_gnn_seq", "make_mlp", "setup_training_logger", "zero_index_edge_index", diff --git a/neural_lam/utils/graph.py b/neural_lam/utils/graph.py index 08259cd0..009e747f 100644 --- a/neural_lam/utils/graph.py +++ b/neural_lam/utils/graph.py @@ -9,8 +9,10 @@ # Third-party import torch import yaml +from torch import nn # Local +from ..datastore import BaseDatastore from .buffer_list import BufferList LEGACY_GRAPH_SPEC_VERSION = "legacy" @@ -418,3 +420,93 @@ def load_graph_spec_version() -> str: "mesh_down_features": mesh_down_features, "mesh_static_features": mesh_static_features, } + + +def load_and_register_graph( + module: nn.Module, + datastore: BaseDatastore, + graph_name: str, + mesh_node_features_scaling: float, +) -> bool: + """ + Load a graph and register its tensors on ``module``. + + Loads the graph ``graph_name`` from the datastore's graph directory via + :func:`load_graph`, then registers each tensor as a non-persistent + buffer and each non-tensor (e.g. ``BufferList``) as a plain attribute on + ``module``. + + Parameters + ---------- + module : torch.nn.Module + Module to register the graph tensors and attributes on, in place. + datastore : BaseDatastore + Datastore whose ``root_path`` holds the ``graph`` directory. + graph_name : str + Name of the graph directory (under ``/graph``) to load. + mesh_node_features_scaling : float + Scalar used to normalize mesh node coordinate features for graphs in + the current on-disk format; forwarded to :func:`load_graph`. + + Returns + ------- + bool + Whether the loaded graph is hierarchical. + """ + graph_dir_path = datastore.root_path / "graph" / graph_name + hierarchical, graph_ldict = load_graph( + graph_dir_path=graph_dir_path, + mesh_node_features_scaling=mesh_node_features_scaling, + ) + for name, attr_value in graph_ldict.items(): + # Make BufferLists module members and register tensors as buffers + if isinstance(attr_value, torch.Tensor): + module.register_buffer(name, attr_value, persistent=False) + else: + setattr(module, name, attr_value) + return hierarchical + + +def compute_grid_input_dim( + datastore: BaseDatastore, + num_past_forcing_steps: int, + num_future_forcing_steps: int, +) -> int: + """ + Compute the total grid input dimensionality of a graph step predictor. + + The grid input concatenates the two previous states, the grid static + features and the windowed forcing + (past + current + future forcing steps). + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing the number of state, static and forcing variables. + num_past_forcing_steps : int + Number of past forcing steps included in the input window. + num_future_forcing_steps : int + Number of future forcing steps included in the input window. + + Returns + ------- + int + Total grid input dimensionality. + """ + num_state_vars = datastore.get_num_data_vars(category="state") + num_forcing_vars = datastore.get_num_data_vars(category="forcing") + # The static category is optional: when the datastore provides no static + # data array the grid carries no static features, mirroring the empty + # (N, 0) static buffer the step predictor builds in that case. + da_static = datastore.get_dataarray(category="static", split=None) + num_static_vars = ( + 0 + if da_static is None + else datastore.get_num_data_vars(category="static") + ) + return ( + 2 * num_state_vars + + num_static_vars + + num_forcing_vars + * (num_past_forcing_steps + num_future_forcing_steps + 1) + ) diff --git a/neural_lam/utils/networks.py b/neural_lam/utils/networks.py index df2807fd..c78fa687 100644 --- a/neural_lam/utils/networks.py +++ b/neural_lam/utils/networks.py @@ -1,6 +1,7 @@ """Constructors for neural-network building blocks.""" # Third-party +import torch_geometric as pyg from torch import nn @@ -37,3 +38,69 @@ def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: layers.append(nn.LayerNorm(blueprint[-1])) return nn.Sequential(*layers) + + +def make_gnn_seq( + edge_index, + num_gnn_layers, + hidden_layers, + hidden_dim, + gnn_type="InteractionNet", +): + """ + Build a sequential stack of GNN layers that propagates both node and + edge representations. + + All layer types share the ``(send, rec, edge) -> (rec, edge)`` + interface, so the stack can be applied as a single module. + + Parameters + ---------- + edge_index : torch.Tensor + Shape ``(2, M)``. Edge index of the edges that the GNN layers + operate on. + num_gnn_layers : int + Number of stacked GNN layers; must be at least 1. Callers that + want a no-op stage (e.g. zero intra-level layers) should skip + building and applying the stack rather than calling this with 0. + hidden_layers : int + Number of hidden layers in the MLPs of each GNN layer. + hidden_dim : int + Dimensionality of node and edge representations. + gnn_type : str + GNN layer type, any key in ``gnn_layers.GNN_TYPES``. + + Returns + ------- + pyg.nn.Sequential + Sequential module mapping ``(mesh_rep, edge_rep)`` to updated + ``(mesh_rep, edge_rep)``. + + Raises + ------ + ValueError + If ``num_gnn_layers`` is less than 1. + """ + # First-party + from neural_lam.gnn_layers import get_gnn_class + + if num_gnn_layers < 1: + raise ValueError( + "make_gnn_seq requires num_gnn_layers >= 1 " + f"(got {num_gnn_layers}); skip the stage for a no-op." + ) + gnn_class = get_gnn_class(gnn_type) + return pyg.nn.Sequential( + "mesh_rep, edge_rep", + [ + ( + gnn_class( + edge_index, + hidden_dim, + hidden_layers=hidden_layers, + ), + "mesh_rep, mesh_rep, edge_rep -> mesh_rep, edge_rep", + ) + for _ in range(num_gnn_layers) + ], + ) diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 4168396a..c62e1e5a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -471,7 +471,7 @@ def __getitem__( target states, forcing and batch times. The returned data is unstandardized; normalization is applied on-device - in `ForecasterModule.on_after_batch_transfer`. + in `BaseForecasterModule.on_after_batch_transfer`. Parameters ---------- diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148..6644aae6 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -8,7 +8,11 @@ # First-party from neural_lam import config as nlconfig from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from tests.dummy_datastore import DummyDatastore @@ -50,12 +54,11 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) - model = ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, n_example_pred=1, val_steps_to_log=[1], diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 4b35840e..319e3372 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -12,7 +12,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore @@ -235,13 +235,14 @@ def _create_graph(): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 04c99003..789297d0 100644 --- a/tests/test_gnn_layers.py +++ b/tests/test_gnn_layers.py @@ -73,7 +73,7 @@ def _build_model_and_data( output_clamping_upper=config.training.output_clamping.upper, **gnn_kwargs, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config) B = 2 num_grid_nodes = predictor.num_grid_nodes diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index 8d516bfb..dcc8ba63 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -4,7 +4,11 @@ # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -13,7 +17,8 @@ class _MockStepPredictor(StepPredictor): - """Minimal predictor so a ForecasterModule can be built without a graph.""" + """Minimal predictor so a DeterministicForecasterModule can be built + without a graph.""" def forward(self, prev_state, prev_prev_state, forcing): return torch.zeros_like(prev_state), None @@ -26,8 +31,8 @@ def _build_module(datastore): ) ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) - return ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config) + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) @@ -111,7 +116,9 @@ def test_safe_std_clamps_near_zero(): eps = torch.finfo(torch.float32).eps with pytest.warns(UserWarning, match="near-zero std"): - std = ForecasterModule._safe_std([0.0, 1.0, 2.0], eps, "state") + std = DeterministicForecasterModule._safe_std( + [0.0, 1.0, 2.0], eps, "state" + ) assert std[0] == eps assert std[1] == 1.0 diff --git a/tests/test_graph_efm_model.py b/tests/test_graph_efm_model.py new file mode 100644 index 00000000..e90a8036 --- /dev/null +++ b/tests/test_graph_efm_model.py @@ -0,0 +1,469 @@ +"""Integration tests for the full Graph-EFM model. + +Exercises the pieces that turn the Graph-EFM single-step predictors into a +trainable model: the ``GraphEFMForecaster`` ELBO objective, its inherited +ensemble sampling, the ``ProbabilisticForecasterModule`` wrapping, and the +config-aware assembly path in ``train_model`` (``build_predictor`` / +``build_forecaster_module``). Predictors are built on the real example +datastore with a freshly created graph, mirroring +``tests/test_graph_efm_predictor.py``. +""" + +# Standard library +from argparse import Namespace +from pathlib import Path + +# Third-party +import pytest +import torch + +# First-party +from neural_lam import config as nlconfig +from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.models import ( + GraphEFM, + GraphEFMForecaster, + GraphEFMMultiScale, + ProbabilisticForecasterModule, +) +from neural_lam.train_model import build_forecaster_module, build_predictor +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 + + +def _datastore_and_config(graph_name): + """ + Build the example datastore + config and ensure ``graph_name`` exists. + + Parameters + ---------- + graph_name : str + Graph directory name; ``"hierarchical"`` builds a multi-level graph, + anything else a flat one. + + Returns + ------- + datastore : BaseDatastore + The example ``mdp`` datastore. + config : NeuralLAMConfig + A configuration selecting that datastore. + """ + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + + hierarchical = graph_name == "hierarchical" + n_max_levels = 3 if hierarchical else 1 + graph_dir_path = Path(datastore.root_path) / "graph" / graph_name + if not graph_dir_path.exists(): + create_graph_from_datastore( + datastore=datastore, + output_root_path=str(graph_dir_path), + hierarchical=hierarchical, + n_max_levels=n_max_levels, + ) + return datastore, config + + +def _build_predictor(datastore, graph_name, output_std=False): + """ + Construct a small Graph-EFM predictor for ``graph_name``. + + Parameters + ---------- + datastore : BaseDatastore + Datastore to build the predictor on. + graph_name : str + Graph directory name selecting the flat vs hierarchical variant. + output_std : bool, default False + Whether the decoder outputs its own std. + + Returns + ------- + BaseGraphEFM + The constructed predictor. + """ + if graph_name == "hierarchical": + predictor_class = GraphEFM + layer_kwargs = { + "prior_intra_level_layers": 1, + "encoder_intra_level_layers": 1, + "decoder_intra_level_layers": 1, + } + else: + predictor_class = GraphEFMMultiScale + layer_kwargs = { + "prior_m2m_layers": 1, + "encoder_m2m_layers": 1, + "decoder_m2m_layers": 1, + } + return predictor_class( + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + learn_prior=True, + prior_dist="isotropic", + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + output_std=output_std, + **layer_kwargs, + ) + + +def _example_batch(datastore, predictor, batch_size=2, pred_steps=3): + """ + Build a synthetic ``(init_states, forcing_features, target_states)`` batch. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing variable counts. + predictor : BaseGraphEFM + Predictor providing the grid node count. + batch_size : int, default 2 + Number of samples in the batch. + pred_steps : int, default 3 + Rollout length. + + Returns + ------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. + """ + num_grid_nodes = predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + init_states = torch.randn(batch_size, 2, num_grid_nodes, d_state) + forcing_features = torch.randn( + batch_size, pred_steps, num_grid_nodes, d_forcing + ) + target_states = torch.randn(batch_size, pred_steps, num_grid_nodes, d_state) + return init_states, forcing_features, target_states + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_forecaster_forward_and_ensemble_shapes(graph_name): + """The forecaster unrolls a prior-sampled rollout and stacks members into + an ensemble of the documented shape.""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster(predictor, datastore, config=config) + + B, pred_steps, num_members = 2, 3, 4 + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=B, pred_steps=pred_steps + ) + d_state = target_states.shape[-1] + num_grid_nodes = predictor.num_grid_nodes + + prediction, pred_std = forecaster( + init_states, forcing_features, target_states + ) + assert prediction.shape == (B, pred_steps, num_grid_nodes, d_state) + assert pred_std is None # output_std=False predictor + + 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 latent samples + assert not torch.allclose(ensemble[:, 0], ensemble[:, 1]) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_predictor_step_distributions_contract(graph_name): + """step_distributions reuses the shared graph embedding and returns the + prior on the inference path, and the posterior (with the prior gated on + compute_prior) on the training path.""" + datastore, _ = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=2, pred_steps=1 + ) + prev_prev_state, prev_state = init_states[:, 0], init_states[:, 1] + forcing, target_state = forcing_features[:, 0], target_states[:, 0] + d_state = target_state.shape[-1] + num_grid_nodes = predictor.num_grid_nodes + + graph_emb = predictor.embedd_graph(2) + assert {"g2m", "m2g", "mesh"} <= set(graph_emb) + + # Inference path: latent from the prior, no posterior + prior, posterior, pred_mean, pred_std = predictor.step_distributions( + prev_state, prev_prev_state, forcing, graph_emb, target_state=None + ) + assert prior is not None and posterior is None + assert pred_mean.shape == (2, num_grid_nodes, d_state) + assert pred_std is None + + # Training path with KL: both distributions present + prior, posterior, _, _ = predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=True, + ) + assert prior is not None and posterior is not None + + # Training path without KL: prior skipped + prior, posterior, _, _ = predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=False, + ) + assert prior is None and posterior is not None + + +def test_predictor_clamps_predicted_mean(): + """With output clamping configured for a feature, Graph-EFM keeps the + predicted mean for that feature within the configured bounds, clamping it + like the deterministic models do.""" + datastore, _ = _datastore_and_config("1level") + state_names = datastore.get_vars_names(category="state") + lower, upper = -0.5, 0.5 + predictor = GraphEFMMultiScale( + datastore=datastore, + graph_name="1level", + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + prior_m2m_layers=1, + encoder_m2m_layers=1, + decoder_m2m_layers=1, + output_clamping_lower={state_names[0]: lower}, + output_clamping_upper={state_names[0]: upper}, + ) + # The first state feature has a two-sided (sigmoid) clamp registered + assert predictor.clamp_lower_upper_idx.tolist() == [0] + lower_n = (lower - predictor.state_mean[0]) / predictor.state_std[0] + upper_n = (upper - predictor.state_mean[0]) / predictor.state_std[0] + + B = 2 + num_grid_nodes = predictor.num_grid_nodes + d_state = len(state_names) + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + prev_state = torch.randn(B, num_grid_nodes, d_state) + # The current value of the clamped feature must be within its bounds so + # the inverse clamp is finite; the midpoint is a safe choice. + prev_state[..., 0] = (lower_n + upper_n) / 2 + prev_prev_state = torch.randn(B, num_grid_nodes, d_state) + forcing = torch.randn(B, num_grid_nodes, d_forcing) + + pred_mean, _ = predictor(prev_state, prev_prev_state, forcing) + + clamped_feature = pred_mean[..., 0] + assert torch.all(clamped_feature > lower_n) + assert torch.all(clamped_feature < upper_n) + + +def test_hierarchical_zero_intra_level_layers_runs(): + """A hierarchical GraphEFM with no intra-level layers uses an empty m2m + placeholder; forward must still run (regression for m2m handling).""" + datastore, _ = _datastore_and_config("hierarchical") + predictor = GraphEFM( + datastore=datastore, + graph_name="hierarchical", + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + prior_intra_level_layers=0, + encoder_intra_level_layers=0, + decoder_intra_level_layers=0, + ) + assert not predictor.embedd_m2m + + B = 2 + num_grid_nodes = predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + prev_state = torch.randn(B, num_grid_nodes, d_state) + prev_prev_state = torch.randn(B, num_grid_nodes, d_state) + forcing = torch.randn(B, num_grid_nodes, d_forcing) + + pred_mean, _ = predictor(prev_state, prev_prev_state, forcing) + assert pred_mean.shape == (B, num_grid_nodes, d_state) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_elbo_training_loss_gradient_flow(graph_name): + """compute_training_loss returns a finite scalar ELBO with likelihood/KL + components, and gradients flow back into the predictor.""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=1.0 + ) + + init_states, forcing_features, target_states = _example_batch( + datastore, predictor + ) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + + torch.manual_seed(0) + 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 torch.isfinite(batch_loss) + assert set(loss_components) == {"elbo_likelihood", "elbo_kl", "elbo"} + assert (loss_components["elbo_kl"] >= 0).all() + + batch_loss.backward() + grads = [p.grad for p in predictor.parameters() if p.grad is not None] + assert grads, "no gradients reached the predictor" + assert any(torch.any(g != 0) for g in grads) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_elbo_kl_beta_zero_skips_kl(graph_name): + """With kl_beta=0 the loss is the negative likelihood alone and no KL + component is reported (pure auto-encoder training).""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=0.0 + ) + + init_states, forcing_features, target_states = _example_batch( + datastore, predictor + ) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + + torch.manual_seed(0) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + assert set(loss_components) == {"elbo_likelihood"} + torch.testing.assert_close(batch_loss, -loss_components["elbo_likelihood"]) + batch_loss.backward() # still differentiable + + +def test_module_training_and_validation_steps(): + """The ProbabilisticForecasterModule delegates training to the forecaster + ELBO and scores an ensemble mean during validation.""" + datastore, config = _datastore_and_config("1level") + predictor = _build_predictor(datastore, "1level") + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=1.0 + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=2, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(0) + train_loss = model.training_step(batch) + assert train_loss.shape == () + assert torch.isfinite(train_loss) + + model.validation_step(batch, 0) + (entry_mses,) = model.val_metrics["ens_mse"] + d_state = target_states.shape[-1] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) + + +@pytest.mark.parametrize( + "model_name, predictor_class, graph_name", + [ + ("graph_efm", GraphEFM, "hierarchical"), + ("graph_efm_ms", GraphEFMMultiScale, "1level"), + ], +) +def test_train_model_assembly_selects_probabilistic_path( + model_name, predictor_class, graph_name +): + """train_model's build_predictor/build_forecaster_module route the + graph_efm* models through the probabilistic assembly, reading the + Graph-EFM hyperparameters from the ``probabilistic`` config section and + producing the right predictor, forecaster, module class and checkpoint + monitor.""" + datastore, _ = _datastore_and_config(graph_name) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + probabilistic=nlconfig.ProbabilisticConfig( + latent_dim=4, + prior_layers=1, + encoder_layers=1, + decoder_layers=1, + kl_beta=0.5, + eval_ensemble_size=3, + ), + ) + args = Namespace( + model=model_name, + graph=graph_name, + hidden_dim=4, + hidden_layers=1, + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + output_std=False, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="InteractionNet", + loss="mse", + ) + + predictor = build_predictor(args, config, datastore) + assert isinstance(predictor, predictor_class) + # Hyperparameters came from config.probabilistic, not the CLI args + assert predictor.latent_dim == 4 + + forecaster, module_class, module_kwargs, val_monitor = ( + build_forecaster_module(args, config, datastore, predictor) + ) + assert isinstance(forecaster, GraphEFMForecaster) + assert forecaster.kl_beta == 0.5 + assert module_class is ProbabilisticForecasterModule + assert module_kwargs == {"eval_ensemble_size": 3} + assert val_monitor == "val_mean_ens_rmse" diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py new file mode 100644 index 00000000..a18580e0 --- /dev/null +++ b/tests/test_graph_efm_predictor.py @@ -0,0 +1,155 @@ +"""Unit tests for the Graph-EFM single-step probabilistic predictors. + +These mirror the smoke-test pattern used for the deterministic predictors +(see ``tests/test_gnn_layers.py``): build the flat (GraphEFMMultiScale) and +hierarchical (GraphEFM) variants on the real example datastore with a freshly +created graph, then exercise ``forward`` on synthetic tensors. +""" + +# Standard library +from pathlib import Path + +# Third-party +import pytest +import torch + +# First-party +from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.models.step_predictors.graph.graph_efm import ( + GraphEFM, + GraphEFMMultiScale, +) +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 + + +def _datastore_with_graph(graph_name): + """Create the example datastore and ensure ``graph_name`` exists.""" + datastore = init_datastore_example("mdp") + + if graph_name == "hierarchical": + hierarchical = True + n_max_levels = 3 + else: + hierarchical = False + n_max_levels = 1 + + graph_dir_path = Path(datastore.root_path) / "graph" / graph_name + if not graph_dir_path.exists(): + create_graph_from_datastore( + datastore=datastore, + output_root_path=str(graph_dir_path), + hierarchical=hierarchical, + n_max_levels=n_max_levels, + ) + return datastore + + +def _build_predictor(graph_name, output_std=False): + datastore = _datastore_with_graph(graph_name) + if graph_name == "hierarchical": + predictor_class = GraphEFM + layer_kwargs = { + "prior_intra_level_layers": 1, + "encoder_intra_level_layers": 1, + "decoder_intra_level_layers": 1, + } + else: + predictor_class = GraphEFMMultiScale + layer_kwargs = { + "prior_m2m_layers": 1, + "encoder_m2m_layers": 1, + "decoder_m2m_layers": 1, + } + predictor = predictor_class( + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + learn_prior=True, + prior_dist="isotropic", + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + output_std=output_std, + **layer_kwargs, + ) + return predictor, datastore + + +def _make_inputs(predictor, datastore, batch_size=2): + num_grid_nodes = predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + prev_state = torch.randn(batch_size, num_grid_nodes, d_state) + prev_prev_state = torch.randn(batch_size, num_grid_nodes, d_state) + forcing = torch.randn(batch_size, num_grid_nodes, d_forcing) + return prev_state, prev_prev_state, forcing, d_state + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_forward_shapes_and_no_std(graph_name): + """forward returns a (B, num_grid_nodes, d_state) state and None std when + output_std is False, for both flat and hierarchical graphs.""" + predictor, datastore = _build_predictor(graph_name) + prev_state, prev_prev_state, forcing, d_state = _make_inputs( + predictor, datastore + ) + + new_state, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert new_state.shape == (2, predictor.num_grid_nodes, d_state) + assert pred_std is None + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_forward_output_std_returns_std(graph_name): + """With output_std=True the decoder produces a positive std of the same + shape as the state.""" + predictor, datastore = _build_predictor(graph_name, output_std=True) + prev_state, prev_prev_state, forcing, d_state = _make_inputs( + predictor, datastore + ) + + new_state, pred_std = predictor(prev_state, prev_prev_state, forcing) + + expected = (2, predictor.num_grid_nodes, d_state) + assert new_state.shape == expected + assert pred_std is not None + assert pred_std.shape == expected + assert (pred_std > 0).all() + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_forward_member_stochasticity(graph_name): + """Two forward calls with identical inputs differ, because the latent is + resampled from the prior each call (catches an unused-latent regression).""" + predictor, datastore = _build_predictor(graph_name) + prev_state, prev_prev_state, forcing, _ = _make_inputs(predictor, datastore) + + out_a, _ = predictor(prev_state, prev_prev_state, forcing) + out_b, _ = predictor(prev_state, prev_prev_state, forcing) + + assert not torch.allclose(out_a, out_b) + + +@pytest.mark.parametrize( + "predictor_class, graph_name", + [(GraphEFM, "1level"), (GraphEFMMultiScale, "hierarchical")], +) +def test_graph_type_mismatch_raises(predictor_class, graph_name): + """GraphEFM requires a hierarchical graph and GraphEFMMultiScale a flat one; + constructing with the wrong graph type raises ValueError.""" + datastore = _datastore_with_graph(graph_name) + with pytest.raises(ValueError, match="mesh graph"): + predictor_class( + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + ) diff --git a/tests/test_latent_modules.py b/tests/test_latent_modules.py new file mode 100644 index 00000000..405688c7 --- /dev/null +++ b/tests/test_latent_modules.py @@ -0,0 +1,556 @@ +"""Unit tests for the latent encoder/decoder infrastructure. + +These tests exercise the latent modules in isolation with synthetic edge +indices and tensor inputs, so they do not depend on any datastore or graph +fixture. They cover output shapes, distribution properties and that +backpropagation reaches all parameters. +""" + +# Third-party +import pytest +import torch + +# First-party +from neural_lam.models.latent import ( + BaseLatentEncoder, + ConstantLatentEncoder, + GraphLatentDecoder, + GraphLatentEncoder, + HiGraphLatentDecoder, + HiGraphLatentEncoder, +) +from neural_lam.utils import make_gnn_seq + + +def _fully_connected_edge_index(n_send, n_rec): + senders = ( + torch.arange(n_send).unsqueeze(1).expand(n_send, n_rec).reshape(-1) + ) + receivers = ( + torch.arange(n_rec).unsqueeze(0).expand(n_send, n_rec).reshape(-1) + ) + return torch.stack([senders, receivers]) + + +def _assert_every_param_has_grad(module): + """Fail if any trainable parameter on ``module`` received no gradient. + + Catches dead-param regressions if a future change wires in a sub-module + that the forward pass never reaches. + """ + for name, p in module.named_parameters(): + if p.requires_grad: + assert p.grad is not None, f"parameter {name} received no gradient" + + +@pytest.fixture +def flat_dims(): + return { + "batch_size": 2, + "num_grid": 5, + "num_mesh": 3, + "hidden_dim": 8, + "latent_dim": 4, + "num_state_vars": 2, + "hidden_layers": 1, + "m2m_layers": 2, + } + + +@pytest.fixture +def flat_edges(flat_dims): + n_grid = flat_dims["num_grid"] + n_mesh = flat_dims["num_mesh"] + return { + "g2m": _fully_connected_edge_index(n_grid, n_mesh), + "m2m": _fully_connected_edge_index(n_mesh, n_mesh), + "m2g": _fully_connected_edge_index(n_mesh, n_grid), + } + + +@pytest.fixture +def flat_graph_emb(flat_dims, flat_edges): + B = flat_dims["batch_size"] + d_h = flat_dims["hidden_dim"] + return { + "mesh": torch.randn(B, flat_dims["num_mesh"], d_h), + "g2m": torch.randn(B, flat_edges["g2m"].shape[1], d_h), + "m2m": torch.randn(B, flat_edges["m2m"].shape[1], d_h), + "m2g": torch.randn(B, flat_edges["m2g"].shape[1], d_h), + } + + +def test_make_gnn_seq_zero_layers_raises(): + """make_gnn_seq must build a real sequence; the no-op case is the + caller's responsibility, exercised via the zero-intra-layer tests.""" + edge_index = _fully_connected_edge_index(3, 3) + with pytest.raises(ValueError, match="num_gnn_layers >= 1"): + make_gnn_seq( + edge_index, num_gnn_layers=0, hidden_layers=1, hidden_dim=8 + ) + + +def test_make_gnn_seq_positive_layers_runs(): + edge_index = _fully_connected_edge_index(3, 3) + seq = make_gnn_seq( + edge_index, num_gnn_layers=2, hidden_layers=1, hidden_dim=8 + ) + mesh_rep = torch.randn(2, 3, 8) + edge_rep = torch.randn(2, edge_index.shape[1], 8) + out_mesh, out_edge = seq(mesh_rep, edge_rep) + assert out_mesh.shape == mesh_rep.shape + assert out_edge.shape == edge_rep.shape + + +class _IdentityEncoder(BaseLatentEncoder): + """Trivial encoder used to verify BaseLatentEncoder distribution logic.""" + + def __init__(self, latent_dim, num_mesh_nodes, output_dist): + super().__init__(latent_dim, output_dist) + self.num_mesh_nodes = num_mesh_nodes + # Learnable params so we can verify backprop reaches them + self.bias = torch.nn.Parameter(torch.zeros(self.output_dim)) + + def compute_dist_params(self, grid_rep, **kwargs): + B = grid_rep.shape[0] + return self.bias.expand(B, self.num_mesh_nodes, self.output_dim) + + +def test_base_encoder_isotropic_has_unit_std(): + enc = _IdentityEncoder( + latent_dim=4, num_mesh_nodes=3, output_dist="isotropic" + ) + grid_rep = torch.randn(2, 5, 8) + dist = enc(grid_rep) + assert isinstance(dist, torch.distributions.Normal) + assert dist.mean.shape == (2, 3, 4) + assert torch.allclose(dist.stddev, torch.ones_like(dist.stddev)) + + +def test_base_encoder_diagonal_has_positive_std(): + enc = _IdentityEncoder( + latent_dim=4, num_mesh_nodes=3, output_dist="diagonal" + ) + grid_rep = torch.randn(2, 5, 8) + dist = enc(grid_rep) + assert dist.mean.shape == (2, 3, 4) + assert dist.stddev.shape == (2, 3, 4) + # softplus(0) + eps must be strictly positive + assert (dist.stddev > 0).all() + + +def test_base_encoder_rejects_unknown_dist(): + with pytest.raises(ValueError): + _IdentityEncoder(latent_dim=4, num_mesh_nodes=3, output_dist="bogus") + + +def test_constant_encoder_is_input_independent(): + enc = ConstantLatentEncoder( + latent_dim=4, num_mesh_nodes=3, output_dist="isotropic" + ) + a = enc(torch.randn(2, 5, 8)) + b = enc(torch.randn(2, 5, 8) * 100) + assert torch.equal(a.mean, b.mean) + assert torch.equal(a.stddev, b.stddev) + assert a.mean.shape == (2, 3, 4) + # Isotropic output is a mean-0 standard normal + assert torch.equal(a.mean, torch.zeros_like(a.mean)) + assert torch.allclose(a.stddev, torch.ones_like(a.stddev)) + + +def test_graph_encoder_shapes_and_backprop( + flat_dims, flat_edges, flat_graph_emb +): + enc = GraphLatentEncoder( + latent_dim=flat_dims["latent_dim"], + g2m_edge_index=flat_edges["g2m"], + m2m_edge_index=flat_edges["m2m"], + hidden_dim=flat_dims["hidden_dim"], + m2m_layers=flat_dims["m2m_layers"], + hidden_layers=flat_dims["hidden_layers"], + output_dist="diagonal", + ) + grid_rep = torch.randn( + flat_dims["batch_size"], + flat_dims["num_grid"], + flat_dims["hidden_dim"], + ) + dist = enc(grid_rep, graph_emb=flat_graph_emb) + assert dist.mean.shape == ( + flat_dims["batch_size"], + flat_dims["num_mesh"], + flat_dims["latent_dim"], + ) + + dist.rsample().sum().backward() + _assert_every_param_has_grad(enc) + + +def test_graph_decoder_shapes_with_output_std( + flat_dims, flat_edges, flat_graph_emb +): + dec = GraphLatentDecoder( + g2m_edge_index=flat_edges["g2m"], + m2m_edge_index=flat_edges["m2m"], + m2g_edge_index=flat_edges["m2g"], + hidden_dim=flat_dims["hidden_dim"], + latent_dim=flat_dims["latent_dim"], + num_state_vars=flat_dims["num_state_vars"], + m2m_layers=flat_dims["m2m_layers"], + hidden_layers=flat_dims["hidden_layers"], + output_std=True, + ) + B = flat_dims["batch_size"] + grid_rep = torch.randn(B, flat_dims["num_grid"], flat_dims["hidden_dim"]) + latent_samples = torch.randn( + B, flat_dims["num_mesh"], flat_dims["latent_dim"] + ) + mean_delta, pred_std = dec(grid_rep, latent_samples, flat_graph_emb) + + expected_shape = (B, flat_dims["num_grid"], flat_dims["num_state_vars"]) + assert mean_delta.shape == expected_shape + assert pred_std is not None + assert pred_std.shape == expected_shape + assert (pred_std > 0).all() + + (mean_delta.sum() + pred_std.sum()).backward() + _assert_every_param_has_grad(dec) + + +def test_graph_decoder_no_output_std_returns_none( + flat_dims, flat_edges, flat_graph_emb +): + dec = GraphLatentDecoder( + g2m_edge_index=flat_edges["g2m"], + m2m_edge_index=flat_edges["m2m"], + m2g_edge_index=flat_edges["m2g"], + hidden_dim=flat_dims["hidden_dim"], + latent_dim=flat_dims["latent_dim"], + num_state_vars=flat_dims["num_state_vars"], + m2m_layers=flat_dims["m2m_layers"], + hidden_layers=flat_dims["hidden_layers"], + output_std=False, + ) + B = flat_dims["batch_size"] + grid_rep = torch.randn(B, flat_dims["num_grid"], flat_dims["hidden_dim"]) + latent_samples = torch.randn( + B, flat_dims["num_mesh"], flat_dims["latent_dim"] + ) + mean_delta, pred_std = dec(grid_rep, latent_samples, flat_graph_emb) + assert mean_delta.shape == ( + B, + flat_dims["num_grid"], + flat_dims["num_state_vars"], + ) + assert pred_std is None + + +def test_flat_modules_zero_m2m_layers_skip_processing( + flat_dims, flat_edges, flat_graph_emb +): + """m2m_layers=0 builds no on-mesh GNNs and skips on-mesh processing in + the forward pass. Exercise both flat modules.""" + enc = GraphLatentEncoder( + latent_dim=flat_dims["latent_dim"], + g2m_edge_index=flat_edges["g2m"], + m2m_edge_index=flat_edges["m2m"], + hidden_dim=flat_dims["hidden_dim"], + m2m_layers=0, + hidden_layers=flat_dims["hidden_layers"], + ) + assert enc.m2m_gnns is None + + dec = GraphLatentDecoder( + g2m_edge_index=flat_edges["g2m"], + m2m_edge_index=flat_edges["m2m"], + m2g_edge_index=flat_edges["m2g"], + hidden_dim=flat_dims["hidden_dim"], + latent_dim=flat_dims["latent_dim"], + num_state_vars=flat_dims["num_state_vars"], + m2m_layers=0, + hidden_layers=flat_dims["hidden_layers"], + ) + assert dec.m2m_gnns is None + + B = flat_dims["batch_size"] + grid_rep = torch.randn(B, flat_dims["num_grid"], flat_dims["hidden_dim"]) + dist = enc(grid_rep, graph_emb=flat_graph_emb) + assert dist.mean.shape == ( + B, + flat_dims["num_mesh"], + flat_dims["latent_dim"], + ) + + latent_samples = torch.randn( + B, flat_dims["num_mesh"], flat_dims["latent_dim"] + ) + mean_delta, _ = dec(grid_rep, latent_samples, flat_graph_emb) + assert mean_delta.shape == ( + B, + flat_dims["num_grid"], + flat_dims["num_state_vars"], + ) + + +# --- Hierarchical fixtures and tests ---------------------------------------- + + +@pytest.fixture +def hi_dims(): + return { + "batch_size": 2, + "num_grid": 5, + "mesh_per_level": [4, 3], # bottom -> top + "hidden_dim": 8, + "latent_dim": 4, + "num_state_vars": 2, + "hidden_layers": 1, + "intra_level_layers": 1, + } + + +@pytest.fixture +def hi_edges(hi_dims): + bot, top = hi_dims["mesh_per_level"] + n_grid = hi_dims["num_grid"] + return { + "g2m": _fully_connected_edge_index(n_grid, bot), + "m2g": _fully_connected_edge_index(bot, n_grid), + "m2m": [ + _fully_connected_edge_index(bot, bot), + _fully_connected_edge_index(top, top), + ], + "mesh_up": [_fully_connected_edge_index(bot, top)], + "mesh_down": [_fully_connected_edge_index(top, bot)], + } + + +@pytest.fixture +def hi_graph_emb(hi_dims, hi_edges): + B = hi_dims["batch_size"] + d_h = hi_dims["hidden_dim"] + return { + "mesh": [torch.randn(B, n, d_h) for n in hi_dims["mesh_per_level"]], + "g2m": torch.randn(B, hi_edges["g2m"].shape[1], d_h), + "m2g": torch.randn(B, hi_edges["m2g"].shape[1], d_h), + "m2m": [torch.randn(B, e.shape[1], d_h) for e in hi_edges["m2m"]], + "mesh_up": [ + torch.randn(B, e.shape[1], d_h) for e in hi_edges["mesh_up"] + ], + "mesh_down": [ + torch.randn(B, e.shape[1], d_h) for e in hi_edges["mesh_down"] + ], + } + + +def test_hi_graph_encoder_shape_at_top_level(hi_dims, hi_edges, hi_graph_emb): + enc = HiGraphLatentEncoder( + latent_dim=hi_dims["latent_dim"], + g2m_edge_index=hi_edges["g2m"], + m2m_edge_index=hi_edges["m2m"], + mesh_up_edge_index=hi_edges["mesh_up"], + hidden_dim=hi_dims["hidden_dim"], + intra_level_layers=hi_dims["intra_level_layers"], + hidden_layers=hi_dims["hidden_layers"], + output_dist="diagonal", + ) + grid_rep = torch.randn( + hi_dims["batch_size"], hi_dims["num_grid"], hi_dims["hidden_dim"] + ) + dist = enc(grid_rep, graph_emb=hi_graph_emb) + top_n = hi_dims["mesh_per_level"][-1] + assert dist.mean.shape == ( + hi_dims["batch_size"], + top_n, + hi_dims["latent_dim"], + ) + assert (dist.stddev > 0).all() + + +def test_hi_graph_decoder_shape_back_to_grid(hi_dims, hi_edges, hi_graph_emb): + dec = HiGraphLatentDecoder( + g2m_edge_index=hi_edges["g2m"], + m2m_edge_index=hi_edges["m2m"], + m2g_edge_index=hi_edges["m2g"], + mesh_up_edge_index=hi_edges["mesh_up"], + mesh_down_edge_index=hi_edges["mesh_down"], + hidden_dim=hi_dims["hidden_dim"], + latent_dim=hi_dims["latent_dim"], + num_state_vars=hi_dims["num_state_vars"], + intra_level_layers=hi_dims["intra_level_layers"], + hidden_layers=hi_dims["hidden_layers"], + output_std=True, + ) + B = hi_dims["batch_size"] + top_n = hi_dims["mesh_per_level"][-1] + grid_rep = torch.randn(B, hi_dims["num_grid"], hi_dims["hidden_dim"]) + latent_samples = torch.randn(B, top_n, hi_dims["latent_dim"]) + mean_delta, pred_std = dec(grid_rep, latent_samples, hi_graph_emb) + expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + assert mean_delta.shape == expected_shape + assert pred_std.shape == expected_shape + assert (pred_std > 0).all() + + +def _build_hi_inputs(mesh_per_level, num_grid, hidden_dim, batch_size): + """Construct edge indices and graph_emb for an arbitrary mesh hierarchy.""" + bot = mesh_per_level[0] + edges = { + "g2m": _fully_connected_edge_index(num_grid, bot), + "m2g": _fully_connected_edge_index(bot, num_grid), + "m2m": [_fully_connected_edge_index(n, n) for n in mesh_per_level], + "mesh_up": [ + _fully_connected_edge_index(lo, hi) + for lo, hi in zip(mesh_per_level[:-1], mesh_per_level[1:]) + ], + "mesh_down": [ + _fully_connected_edge_index(hi, lo) + for lo, hi in zip(mesh_per_level[:-1], mesh_per_level[1:]) + ], + } + B, d_h = batch_size, hidden_dim + graph_emb = { + "mesh": [torch.randn(B, n, d_h) for n in mesh_per_level], + "g2m": torch.randn(B, edges["g2m"].shape[1], d_h), + "m2g": torch.randn(B, edges["m2g"].shape[1], d_h), + "m2m": [torch.randn(B, e.shape[1], d_h) for e in edges["m2m"]], + "mesh_up": [torch.randn(B, e.shape[1], d_h) for e in edges["mesh_up"]], + "mesh_down": [ + torch.randn(B, e.shape[1], d_h) for e in edges["mesh_down"] + ], + } + return edges, graph_emb + + +def test_hi_graph_decoder_three_levels(): + """Three-level hierarchy exercises non-empty intra_down loop and the full + up/down recursion, which num_levels=2 only partially covers.""" + mesh_per_level = [5, 4, 3] + B, d_h, latent_dim, num_state_vars = 2, 8, 4, 2 + num_grid = 6 + edges, graph_emb = _build_hi_inputs(mesh_per_level, num_grid, d_h, B) + + dec = HiGraphLatentDecoder( + g2m_edge_index=edges["g2m"], + m2m_edge_index=edges["m2m"], + m2g_edge_index=edges["m2g"], + mesh_up_edge_index=edges["mesh_up"], + mesh_down_edge_index=edges["mesh_down"], + hidden_dim=d_h, + latent_dim=latent_dim, + num_state_vars=num_state_vars, + intra_level_layers=1, + hidden_layers=1, + output_std=True, + ) + grid_rep = torch.randn(B, num_grid, d_h) + top_n = mesh_per_level[-1] + latent_samples = torch.randn(B, top_n, latent_dim) + mean_delta, pred_std = dec(grid_rep, latent_samples, graph_emb) + assert mean_delta.shape == (B, num_grid, num_state_vars) + assert pred_std.shape == (B, num_grid, num_state_vars) + + (mean_delta.sum() + pred_std.sum()).backward() + _assert_every_param_has_grad(dec) + + +def test_hi_graph_encoder_three_levels(): + mesh_per_level = [5, 4, 3] + B, d_h, latent_dim = 2, 8, 4 + num_grid = 6 + edges, graph_emb = _build_hi_inputs(mesh_per_level, num_grid, d_h, B) + + enc = HiGraphLatentEncoder( + latent_dim=latent_dim, + g2m_edge_index=edges["g2m"], + m2m_edge_index=edges["m2m"], + mesh_up_edge_index=edges["mesh_up"], + hidden_dim=d_h, + intra_level_layers=1, + hidden_layers=1, + output_dist="diagonal", + ) + grid_rep = torch.randn(B, num_grid, d_h) + dist = enc(grid_rep, graph_emb=graph_emb) + assert dist.mean.shape == (B, mesh_per_level[-1], latent_dim) + + dist.rsample().sum().backward() + _assert_every_param_has_grad(enc) + + +def test_hi_graph_modules_reject_single_level(): + """Hierarchical encoder/decoder must refuse a single-level mesh, + otherwise the latent would be silently ignored.""" + # Single-level mesh: m2m has length 1, mesh_up/mesh_down are empty. + edges, _ = _build_hi_inputs( + mesh_per_level=[4], num_grid=5, hidden_dim=8, batch_size=2 + ) + with pytest.raises(ValueError, match="at least 2 mesh levels"): + HiGraphLatentEncoder( + latent_dim=4, + g2m_edge_index=edges["g2m"], + m2m_edge_index=edges["m2m"], + mesh_up_edge_index=edges["mesh_up"], + hidden_dim=8, + intra_level_layers=1, + ) + with pytest.raises(ValueError, match="at least 2 mesh levels"): + HiGraphLatentDecoder( + g2m_edge_index=edges["g2m"], + m2m_edge_index=edges["m2m"], + m2g_edge_index=edges["m2g"], + mesh_up_edge_index=edges["mesh_up"], + mesh_down_edge_index=edges["mesh_down"], + hidden_dim=8, + latent_dim=4, + num_state_vars=2, + intra_level_layers=1, + ) + + +def test_hi_graph_decoder_zero_intra_layers(hi_dims, hi_edges, hi_graph_emb): + """intra_level_layers=0 builds no intra-level GNNs and skips + intra-level processing in the forward pass. Exercise end-to-end.""" + dec = HiGraphLatentDecoder( + g2m_edge_index=hi_edges["g2m"], + m2m_edge_index=hi_edges["m2m"], + m2g_edge_index=hi_edges["m2g"], + mesh_up_edge_index=hi_edges["mesh_up"], + mesh_down_edge_index=hi_edges["mesh_down"], + hidden_dim=hi_dims["hidden_dim"], + latent_dim=hi_dims["latent_dim"], + num_state_vars=hi_dims["num_state_vars"], + intra_level_layers=0, + hidden_layers=hi_dims["hidden_layers"], + output_std=True, + ) + B = hi_dims["batch_size"] + top_n = hi_dims["mesh_per_level"][-1] + grid_rep = torch.randn(B, hi_dims["num_grid"], hi_dims["hidden_dim"]) + latent_samples = torch.randn(B, top_n, hi_dims["latent_dim"]) + mean_delta, pred_std = dec(grid_rep, latent_samples, hi_graph_emb) + expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + assert mean_delta.shape == expected_shape + assert pred_std.shape == expected_shape + + +def test_hi_graph_encoder_zero_intra_layers(hi_dims, hi_edges, hi_graph_emb): + enc = HiGraphLatentEncoder( + latent_dim=hi_dims["latent_dim"], + g2m_edge_index=hi_edges["g2m"], + m2m_edge_index=hi_edges["m2m"], + mesh_up_edge_index=hi_edges["mesh_up"], + hidden_dim=hi_dims["hidden_dim"], + intra_level_layers=0, + hidden_layers=hi_dims["hidden_layers"], + output_dist="isotropic", + ) + grid_rep = torch.randn( + hi_dims["batch_size"], hi_dims["num_grid"], hi_dims["hidden_dim"] + ) + dist = enc(grid_rep, graph_emb=hi_graph_emb) + assert dist.mean.shape == ( + hi_dims["batch_size"], + hi_dims["mesh_per_level"][-1], + hi_dims["latent_dim"], + ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 616d563d..f55755e5 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -18,7 +18,11 @@ from neural_lam import config as nlconfig from neural_lam import vis from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -467,13 +471,14 @@ class ModelArgs: output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, @@ -665,7 +670,7 @@ class _SimpleLogger: # Shared ModelArgs for metrics_watch regression tests (issue #302). # Kept at module level to avoid copy-paste duplication across tests. def _build_metrics_watch_module(datastore, config): - """Build a ForecasterModule wired for metrics_watch tests.""" + """Build a DeterministicForecasterModule wired for metrics_watch tests.""" predictor = GraphLAM( datastore=datastore, graph_name="1level", @@ -679,12 +684,11 @@ def _build_metrics_watch_module(datastore, config): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) - return ForecasterModule( + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 73e2f905..85bb06d1 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -2,12 +2,18 @@ from argparse import Namespace # Third-party +import pytest import pytorch_lightning as pl import torch # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam import metrics +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -71,6 +77,92 @@ def test_ar_forecaster_unroll(): assert torch.all(prediction[:, :, 1:, :] == 5.0) +def test_ar_forecaster_score(): + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = MockStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + mask = torch.ones(num_grid_nodes, dtype=torch.bool) + + # pred_std=None falls back to forecaster.per_var_std and applies the + # forecaster's own configured scoring rule (self.loss) + scored = forecaster.score(prediction, target, None, mask=mask) + expected = forecaster.loss( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored, expected) + + # An explicit metric overrides self.loss, still substituting the + # per_var_std fallback + scored_mse = forecaster.score( + prediction, target, None, metric=metrics.mse, mask=mask + ) + expected_mse = metrics.mse( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored_mse, expected_mse) + + # An explicit pred_std is used as-is, not overridden by per_var_std + explicit_std = torch.full((d_state,), 2.0) + scored_explicit = forecaster.score( + prediction, target, explicit_std, mask=mask + ) + expected_explicit = forecaster.loss( + prediction, target, explicit_std, mask=mask + ) + assert torch.equal(scored_explicit, expected_explicit) + + +def test_ar_forecaster_without_config_raises_on_use_not_construction(): + """A predictor that doesn't output std plus no config is a valid, + unambiguous state at construction time (the forecaster may only ever + be used for inference), so ARForecaster must not raise there. It + should only raise once scoring is actually attempted and has no + std to use, and the error should come from the forecaster itself, not + a wrapping module.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + # Construction succeeds even though predicts_std=False and config=None + forecaster = ARForecaster(predictor, datastore) + assert forecaster.per_var_std is None + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.score(prediction, target, None) + + num_past_forcing_steps = 1 + num_future_forcing_steps = 1 + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + pred_steps = 3 + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + true_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.compute_training_loss( + init_states, + forcing_features, + true_states, + interior_mask_bool=torch.ones(num_grid_nodes, dtype=torch.bool), + ) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") @@ -81,7 +173,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS @@ -97,13 +189,12 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1e-3, restore_opt=False, n_example_pred=1, @@ -133,10 +224,12 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss="mse" + ) # Load from checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -193,21 +286,22 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) - # Use distinctive non-default values so we can detect silent fallback - # to ForecasterModule's defaults during load. + # to DeterministicForecasterModule's defaults during load. saved_loss = "mse" saved_lr = 0.123 saved_create_gif = True saved_val_steps = [2] saved_n_example_pred = 7 - model = ForecasterModule( + forecaster = ARForecaster( + predictor, datastore, config=config, loss=saved_loss + ) + + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=saved_loss, lr=saved_lr, restore_opt=False, n_example_pred=saved_n_example_pred, @@ -269,10 +363,12 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss=saved_loss + ) # Load from hacked old checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -283,8 +379,8 @@ def test_forecaster_module_old_checkpoint(tmp_path): assert loaded_model.forecaster.predictor.__class__.__name__ == "GraphLAM" # Hyperparameters nested in the legacy 'args' namespace must round-trip - # rather than silently falling back to ForecasterModule defaults. - assert loaded_model.hparams.loss == saved_loss + # rather than silently falling back to DeterministicForecasterModule + # defaults. assert loaded_model.hparams.lr == saved_lr assert loaded_model.hparams.val_steps_to_log == saved_val_steps assert loaded_model.create_gif is saved_create_gif diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py new file mode 100644 index 00000000..d98cc345 --- /dev/null +++ b/tests/test_probabilistic_forecaster.py @@ -0,0 +1,362 @@ +# Third-party +import pytest +import torch +from torch import nn + +# First-party +from neural_lam import config as nlconfig +from neural_lam import metrics +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + ProbabilisticARForecaster, + ProbabilisticForecasterModule, + StepPredictor, +) +from tests.conftest import init_datastore_example + + +class ZeroStepPredictor(StepPredictor): + """Deterministic predictor always predicting the zero state.""" + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = torch.zeros_like(prev_state) + pred_std = torch.zeros_like(prev_state) if self.output_std else None + return pred_state, pred_std + + +class NoisyStepPredictor(StepPredictor): + """Stochastic predictor sampling a fresh state at every call.""" + + def __init__(self, datastore, **kwargs): + super().__init__(datastore, **kwargs) + self.noise_scale = nn.Parameter(torch.tensor(1.0)) + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = self.noise_scale * torch.randn_like(prev_state) + return pred_state, None + + +class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster): + """ + Test-only concrete ``ProbabilisticARForecaster``. + + ``ProbabilisticARForecaster`` leaves ``compute_training_loss`` abstract + (no single default objective fits every stochastic model), so tests + that only need a working forecaster to instantiate use this example + ensemble-mean objective rather than the base class directly. + ``sample_ensemble`` always requires an explicit member count, so this + class takes its own ``train_num_members`` for the training objective. + """ + + def __init__(self, *args, train_num_members: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.train_num_members = train_num_members + + def compute_training_loss( + self, + init_states, + forcing_features, + target_states, + interior_mask_bool, + ): + ensemble, per_member_std = self.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=self.train_num_members, + ) + ensemble_mean = ensemble.mean(dim=1) + pred_std = ( + per_member_std.mean(dim=1) + if per_member_std is not None + else self.per_var_std + ) + batch_loss = torch.mean( + self.loss( + ensemble_mean, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} + + +def _example_batch(datastore, B=2, pred_steps=3): + """Create constant example input tensors matching the datastore dims.""" + num_grid_nodes = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + num_past_forcing_steps = 1 + num_future_forcing_steps = 1 + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + target_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) * 5.0 + return init_states, forcing_features, target_states + + +def test_ar_forecaster_training_loss_matches_direct_score(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore, loss="mse") + + init_states, forcing_features, target_states = _example_batch(datastore) + score_metric = metrics.get_metric("mse") + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + forecaster.per_var_std = torch.ones(d_state) + + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + prediction, _ = forecaster(init_states, forcing_features, target_states) + expected_loss = torch.mean( + score_metric( + prediction, + target_states, + forecaster.per_var_std, + mask=interior_mask_bool, + ) + ) + + assert batch_loss.shape == () + assert loss_components == {} + torch.testing.assert_close(batch_loss, expected_loss) + + +def test_sample_ensemble_shapes_and_member_variability(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) + + # Override masks to test boundary masking behaviour + forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) + forecaster.interior_mask[0, 0] = 1 # One node is interior + forecaster.boundary_mask = 1 - forecaster.interior_mask + + B, pred_steps, num_members = 2, 3, 4 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + num_grid_nodes = datastore.num_grid_points + d_state = target_states.shape[-1] + + torch.manual_seed(42) + ensemble, per_member_std = forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=num_members, + ) + + assert ensemble.shape == ( + B, + num_members, + pred_steps, + num_grid_nodes, + d_state, + ) + assert per_member_std is None + + # Members carry independent samples on the interior node + assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0]) + # Boundary nodes are overwritten with the true state in every member + assert torch.all(ensemble[:, :, :, 1:] == 5.0) + + +def test_probabilistic_training_loss_gradient_flow(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster( + predictor, datastore, loss="mse", train_num_members=2 + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + forecaster.per_var_std = torch.ones(d_state) + + torch.manual_seed(42) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + assert batch_loss.shape == () + assert loss_components == {} + assert torch.isfinite(batch_loss) + + batch_loss.backward() + assert predictor.noise_scale.grad is not None + assert predictor.noise_scale.grad != 0.0 + + +def test_sample_ensemble_rejects_empty_member_count(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) + init_states, forcing_features, target_states = _example_batch(datastore) + + with pytest.raises(ValueError, match="num_members"): + forecaster.sample_ensemble( + init_states, forcing_features, target_states, num_members=0 + ) + + +def test_probabilistic_ar_forecaster_is_abstract(): + """ProbabilisticARForecaster leaves compute_training_loss abstract, so + it cannot be instantiated directly; only a subclass that defines an + objective can.""" + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(TypeError, match="abstract"): + ProbabilisticARForecaster(predictor, datastore) + + +def test_module_training_step_delegates_to_forecaster(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + model = DeterministicForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + batch = (init_states, target_states, forcing_features, batch_times) + + batch_loss = model.training_step(batch) + + expected_loss, _ = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=model.interior_mask_bool, + ) + + torch.testing.assert_close(batch_loss, expected_loss) + + +class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster): + """ProbabilisticARForecaster recording the requested member count.""" + + def sample_ensemble(self, *args, **kwargs): + self.last_num_members = kwargs.get("num_members") + return super().sample_ensemble(*args, **kwargs) + + +def test_probabilistic_module_validation_scores_ensemble_mean(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, config=config + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=3, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(42) + model.validation_step(batch, 0) + + # Validation samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.val_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) + + +def test_probabilistic_module_rejects_empty_eval_ensemble(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = ConcreteProbabilisticARForecaster( + predictor, datastore, config=config + ) + + with pytest.raises(ValueError, match="eval_ensemble_size"): + ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=0, + ) + + +def test_probabilistic_module_test_step_scores_ensemble_mean(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, config=config + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=3, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(42) + model.test_step(batch, 0) + + # Test samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.test_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index eb4c2bc0..270c04cf 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -50,7 +50,8 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): def test_create_gif_forwarded_to_forecaster_module(): - """--create_gif must be forwarded to ForecasterModule.__init__.""" + """--create_gif must be forwarded to + DeterministicForecasterModule.__init__.""" mock_args = MagicMock() mock_args.eval = None mock_args.load = None @@ -84,7 +85,8 @@ def capture_init(_self, **kwargs): patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), patch("neural_lam.train_model.ARForecaster"), patch( - "neural_lam.models.module.ForecasterModule.__init__", + "neural_lam.models.modules.deterministic." + "DeterministicForecasterModule.__init__", capture_init, ), pytest.raises(SystemExit), @@ -93,7 +95,7 @@ def capture_init(_self, **kwargs): assert ( "create_gif" in captured_kwargs - ), "create_gif was not forwarded to ForecasterModule" + ), "create_gif was not forwarded to DeterministicForecasterModule" assert captured_kwargs["create_gif"] is True assert ( "train_steps_to_log" in captured_kwargs diff --git a/tests/test_training.py b/tests/test_training.py index bf1a5884..aca432c5 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -13,7 +13,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -105,7 +105,7 @@ def run_simple_training( ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS, ARForecaster @@ -123,13 +123,12 @@ def run_simple_training( output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, @@ -176,9 +175,9 @@ def all_gather(self, tensor_to_gather, sync_grads=False): return tensor_to_gather module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) # Simulate a 3D metric tensor: (N_eval, pred_steps, d_f) @@ -207,9 +206,9 @@ def all_gather(self, tensor, sync_grads=False): return torch.stack([tensor, tensor], dim=0) module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) tensor = torch.randn(4, 3, 5) # (N_eval, pred_steps, d_f)