diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 1eb1d526..c95b8bbf 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -387,6 +387,80 @@ def crps_gauss( ) +def crps_ens( + pred: torch.Tensor, + target: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ens_dim: int = 1, +) -> torch.Tensor: + """ + (Negative) Continuous Ranked Probability Score from ensemble samples. + + Unbiased ("fair") estimator of the CRPS computed directly from a set of + ensemble members, following Zamo and Naveau / WeatherBench 2. + + Parameters + ---------- + pred : torch.Tensor + Shape ``(..., M, ..., N, num_variables)``. Ensemble predictions, with + the ``M`` ensemble members laid out along ``ens_dim``. + target : torch.Tensor + Shape ``(..., N, num_variables)``. Ground-truth target (no ensemble + dimension). + pred_std : torch.Tensor or None, optional + Unused; accepted so the signature matches the other metrics. + mask : torch.Tensor or None, optional + Shape ``(N,)``. Boolean mask selecting grid nodes. ``None`` uses all. + average_grid : bool, optional + If True, average over the grid dimension ``N``. + sum_vars : bool, optional + If True, sum over the variable dimension ``num_variables``. + ens_dim : int, optional + Batch dimension along which the ensemble members are laid out and + reduced. Default ``1``. + + Returns + ------- + torch.Tensor + Reduced CRPS values; shape depends on ``average_grid`` and + ``sum_vars``. + """ + del pred_std # unused; present for a uniform metric signature + num_ens = pred.shape[ens_dim] + if num_ens == 1: + # With a single sample the CRPS reduces to the MAE + return mae( + pred.squeeze(ens_dim), + target, + None, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) + + mean_mae = torch.mean( + torch.abs(pred - target.unsqueeze(ens_dim)), dim=ens_dim + ) # (..., N, num_variables) + if num_ens == 2: + # Cheap exact estimator for a two-member ensemble + pair_diffs_term = -0.5 * torch.abs( + pred.select(ens_dim, 0) - pred.select(ens_dim, 1) + ) + else: + # Rank-based estimator, O(M log M) compute and O(M) memory. + # Ranks start at 1; two argsorts give the per-entry ranks. + ranks = pred.argsort(dim=ens_dim).argsort(ens_dim) + 1 + pair_diffs_term = (1 / (num_ens - 1)) * torch.mean( + (num_ens + 1 - 2 * ranks) * pred, dim=ens_dim + ) + crps_estimator = mean_mae + pair_diffs_term # (..., N, num_variables) + + return mask_and_reduce_metric(crps_estimator, mask, average_grid, sum_vars) + + DEFINED_METRICS = { "mse": mse, "mae": mae, diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..2c8920bb 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,14 +3,24 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.probabilistic import ProbabilisticARForecaster from .module import ForecasterModule +from .probabilistic_module import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel +from .step_predictors.graph.graph_efm import GraphEFM, GraphEFMMS 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 +from .step_predictors.probabilistic import LatentStepPredictor +# NOTE: GraphEFM/GraphEFMMS are intentionally NOT registered in MODELS yet. +# The shared construction call in train_model.py instantiates the chosen +# model with a fixed deterministic kwarg set -- datastore-first, no +# ``config``, and with ``mesh_aggr`` -- whereas the Graph-EFM models require +# ``config`` (for their per_var_std weighting) and take no ``mesh_aggr``. +# Registering them requires config-aware model assembly in train_model.py. MODELS = { "graph_lam": GraphLAM, "hi_lam": HiLAM, diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..120f67a9 --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,156 @@ +"""Autoregressive forecaster for Graph-EFM latent-variable predictors.""" + +# Standard library +from typing import Optional + +# Third-party +import torch + +# Local +from ..step_predictors.probabilistic import LatentStepPredictor +from .autoregressive import ARForecaster + + +class ProbabilisticARForecaster(ARForecaster): + """ + Autoregressive forecaster for latent-variable step predictors. + + The inherited :meth:`ARForecaster.forward` is the evaluation rollout: each + step calls the predictor (which samples its prior), conditions + autoregressively on the predicted state, and overwrites boundary nodes + with the true value. This subclass adds :meth:`training_rollout`, the + posterior-conditioned rollout used for training: at each step it asks the + predictor for the prior, the variational posterior (conditioned on the + true target) and the decoded reconstruction, autoregresses on the + predicted mean, and reduces the per-step KL between posterior and prior. + The KL is purely between the model's own distributions, so it is reduced + here; the likelihood and the ELBO assembly are left to the module. + """ + + # The wrapped predictor must expose per-step latent distributions + predictor: LatentStepPredictor + + def training_rollout( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + compute_kl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Posterior-conditioned rollout producing the per-step ELBO components. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``. + 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)``. The true next + states, used to condition the posterior and to overwrite the + boundary nodes. + compute_kl : bool, optional + If ``False`` (KL weight zero), skip the prior and return + ``kl_terms`` as ``None``. + + Returns + ------- + pred_means : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. Raw decoder + means (not boundary-overwritten), for the likelihood. + pred_stds : torch.Tensor or None + Shape ``(B, pred_steps, num_grid_nodes, d_state)`` when the + predictor emits its own std, otherwise ``None``. + kl_terms : torch.Tensor or None + Shape ``(B, pred_steps)``. Per-step KL of posterior from prior, + or ``None`` when ``compute_kl`` is ``False``. + """ + prev_prev_state = init_states[:, 0] + prev_state = init_states[:, 1] + mean_list = [] + std_list = [] + kl_list = [] + + for i in range(forcing_features.shape[1]): + target = target_states[:, i] + prior_dist, posterior_dist, pred_mean, pred_std = ( + self.predictor.step_distributions( + prev_state, + prev_prev_state, + forcing_features[:, i], + target=target, + compute_prior=compute_kl, + ) + ) + + mean_list.append(pred_mean) # raw, masked to interior later + if pred_std is not None: + std_list.append(pred_std) + if posterior_dist is not None and prior_dist is not None: + kl = torch.distributions.kl_divergence( + posterior_dist, prior_dist + ) + # Sum over all latent dimensions, leaving the batch dim + kl_list.append(kl.sum(dim=tuple(range(1, kl.ndim)))) + + # Autoregress on the predicted mean; boundary overwritten by truth + new_state = ( + self.boundary_mask * target + self.interior_mask * pred_mean + ) + prev_prev_state = prev_state + prev_state = new_state + + pred_means = torch.stack(mean_list, dim=1) + pred_stds = torch.stack(std_list, dim=1) if std_list else None + kl_terms = torch.stack(kl_list, dim=1) if kl_list else None + return pred_means, pred_stds, kl_terms + + def sample_trajectories( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_traj: int, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Sample an ensemble of prior trajectories. + + Runs the evaluation rollout (:meth:`ARForecaster.forward`, which + samples the prior each step) ``num_traj`` times and stacks the + members along a new ensemble dimension. Used for the CRPS term and + for ensemble evaluation. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. True states + used to overwrite boundary nodes at each step. + num_traj : int + Number of trajectories ``S`` to sample. + + Returns + ------- + traj_means : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, d_state)``. + traj_stds : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, d_state)`` when the + predictor emits its own std, otherwise ``None``. + """ + mean_list = [] + std_list = [] + for _ in range(num_traj): + prediction, pred_std = self( + init_states, forcing_features, boundary_states + ) + mean_list.append(prediction) + if pred_std is not None: + std_list.append(pred_std) + + traj_means = torch.stack(mean_list, dim=1) # (B, S, T, N, d_state) + traj_stds = torch.stack(std_list, dim=1) if std_list else None + return traj_means, traj_stds 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..bd19b111 --- /dev/null +++ b/neural_lam/models/latent/base_decoder.py @@ -0,0 +1,176 @@ +"""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 parameters of the next-state distribution 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 only) or + ``2 * num_state_vars`` outputs (mean, 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, last_state, graph_emb): + """ + Predict mean (and optionally std) of the next weather state. + + 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 mean is predicted + as a residual on top of ``last_state``. + + 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. + last_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. State at the + current time step, used as the base of the residual mean + prediction. + 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 + ------- + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted mean + of the next 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 + + pred_mean = last_state + mean_delta + + return pred_mean, 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..63e8344d --- /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 = graph_emb["m2m"][level] + if self.intra_up_gnns is not None: + new_mesh_rep, new_m2m_rep = self.intra_up_gnns[level]( + new_mesh_rep, new_m2m_rep + ) + + # 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/probabilistic_module.py b/neural_lam/models/probabilistic_module.py new file mode 100644 index 00000000..dcbedb01 --- /dev/null +++ b/neural_lam/models/probabilistic_module.py @@ -0,0 +1,100 @@ +"""Lightning module that trains latent-variable forecasters with the ELBO.""" + +# Local +from .forecasters.probabilistic import ProbabilisticARForecaster +from .module import ForecasterModule + + +class ProbabilisticForecasterModule(ForecasterModule): + """ + Lightning module for latent-variable forecasters. + + Extends :class:`ForecasterModule` by assembling the training loss as the + (beta-weighted) negative ELBO. The forecaster's ``training_rollout`` + returns the per-step decoder reconstruction and the per-step KL; this + module computes the likelihood from the reconstruction and the target + through ``self.loss`` (substituting the constant ``per_var_std`` when the + predictor emits no std), sums the likelihood and KL over the rollout, + weights the KL by ``kl_beta`` and averages over the batch. Evaluation + reuses the inherited single-member rollout. + """ + + # The wrapped forecaster must roll out a latent-variable predictor + forecaster: ProbabilisticARForecaster + + def __init__(self, *args, kl_beta: float = 1.0, **kwargs): + """ + Initialize the module and store the KL weight. + + Parameters + ---------- + *args + Positional arguments forwarded to + :meth:`ForecasterModule.__init__` (``forecaster``, ``config``, + ``datastore``, ...). + kl_beta : float, optional + Weight of the KL term in the ELBO. ``0`` trains a pure + autoencoder (the prior and KL are skipped). Default ``1.0``. + **kwargs + Keyword arguments forwarded to + :meth:`ForecasterModule.__init__` (``loss``, ``lr``, ...). + """ + super().__init__(*args, **kwargs) + self.kl_beta = kl_beta + + def training_step(self, batch): + """ + Assemble the beta-weighted negative ELBO over the rollout. + + Parameters + ---------- + batch : tuple + ``(init_states, target_states, forcing_features, batch_times)``. + + Returns + ------- + torch.Tensor + Scalar training loss. + """ + init_states, target_states, forcing_features, _ = batch + compute_kl = self.kl_beta > 0 + + pred_means, pred_stds, kl_terms = self.forecaster.training_rollout( + init_states, + forcing_features, + target_states, + compute_kl=compute_kl, + ) + pred_std = pred_stds if pred_stds is not None else self.per_var_std + + # Per-entry log-likelihood is the negative loss (exactly the + # log-likelihood for the nll loss); sum over interior grid + vars. + entry_log_lik = -self.loss( + pred_means, + target_states, + pred_std, + mask=self.interior_mask_bool, + average_grid=False, + sum_vars=False, + ) # (B, pred_steps, num_interior_nodes, d_state) + likelihood_terms = entry_log_lik.sum(dim=(2, 3)) # (B, pred_steps) + + # Sum the per-step terms over the rollout, mean over the batch + per_sample_likelihood = likelihood_terms.sum(dim=1) # (B,) + loss = -per_sample_likelihood.mean() + log_dict = {"elbo_likelihood": per_sample_likelihood.mean()} + if kl_terms is not None: + per_sample_kl = kl_terms.sum(dim=1) # (B,) + loss = loss + self.kl_beta * per_sample_kl.mean() + log_dict["elbo_kl"] = per_sample_kl.mean() + log_dict["train_loss"] = loss + + self.log_dict( + log_dict, + prog_bar=True, + on_step=True, + on_epoch=True, + sync_dist=True, + batch_size=init_states.shape[0], + ) + return loss diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 1a747971..9576b7a2 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -100,16 +100,9 @@ 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 - self.hierarchical, graph_ldict = utils.load_graph( - graph_dir_path=graph_dir_path + self.hierarchical = utils.load_and_register_graph( + self, datastore, graph_name ) - 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() @@ -119,14 +112,11 @@ 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.grid_input_dim( + datastore, + self.grid_static_features.shape[1], + 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..033f8b68 --- /dev/null +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -0,0 +1,960 @@ +"""Graph-based Ensemble Forecasting Model (Graph-EFM) single-step +predictors, for hierarchical (GraphEFM) and flat (GraphEFMMS) mesh +graphs.""" + +# Standard library +from typing import Callable, Dict, Optional + +# Third-party +import torch +from torch import nn + +# Local +from .... import utils +from ....config import NeuralLAMConfig +from ....datastore import BaseDatastore +from ....loss_weighting import get_state_feature_weighting +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. Besides ``forward`` (sampling a single step + from the prior) it exposes the per-step ELBO pieces + (``compute_step_loss`` -> ``(likelihood_term, kl_term, pred_mean, + pred_std)``). Rollout, ELBO assembly, ensemble logic and logging live + outside the predictor. + + This base class sets up everything that is independent of the mesh + graph type. Concrete subclasses are specific to a graph type (declared + by ``requires_hierarchical``): their constructors build the mesh + embedders and the prior/encoder/decoder latent modules, and they + implement :meth:`embedd_mesh`. See :class:`GraphEFM` (hierarchical + graph) and :class:`GraphEFMMS` (flat graph). + """ + + # Whether the concrete subclass requires a hierarchical mesh graph + requires_hierarchical: bool + + def __init__( + self, + config: NeuralLAMConfig, + datastore: BaseDatastore, + graph_name: str, + hidden_dim: int = 64, + hidden_layers: int = 1, + 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 constant per-variable std. Building the mesh + embedders and the prior/encoder/decoder latent modules is left to + the subclass constructor. + + Parameters + ---------- + config : NeuralLAMConfig + Full Neural-LAM configuration; used for the state feature + weighting that enters the constant per-variable std. + 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 + (``requires_hierarchical``). + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + 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, a constant per-variable std is used as + likelihood scale. + 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. + # NOTE: (IMPORTANT!) mesh nodes MUST have the first + # num_mesh_nodes indices. + self.hierarchical = utils.load_and_register_graph( + self, datastore, graph_name + ) + if self.hierarchical != self.requires_hierarchical: + required_type = ( + "hierarchical" if self.requires_hierarchical else "flat" + ) + loaded_type = "hierarchical" if self.hierarchical else "flat" + raise ValueError( + f"{type(self).__name__} requires a {required_type} mesh " + f"graph, but graph '{graph_name}' is {loaded_type}" + ) + + # 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.grid_input_dim( + datastore, + self.grid_static_features.shape[1], + 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) + + # Constant per-variable std used as the (homoscedastic) likelihood + # scale when the decoder does not output its own std. Mirrors + # ForecasterModule's per_var_std formula + # (state_diff_std_standardized / sqrt(state_feature_weights)); both + # copies are persistent=False so there is no checkpoint interaction. + if not self.output_std: + da_state_stats = datastore.get_standardization_dataarray( + category="state" + ) + state_diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + state_feature_weights = torch.tensor( + get_state_feature_weighting(config=config, datastore=datastore), + dtype=torch.float32, + ) + self.register_buffer( + "per_var_std", + state_diff_std / torch.sqrt(state_feature_weights), + persistent=False, + ) + else: + self.per_var_std = None + + # Compute indices and define clamping functions. GraphEFM's forward + # never clamps (the decoder outputs the full next state), so these are + # inert -- accepted for interface parity with other StepPredictors. + self.prepare_clamping_params(datastore) + + 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_and_graph(self, prev_state, prev_prev_state, forcing): + """ + Embed the grid (states up to t-1) and the full graph. + + 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 + ------- + grid_emb : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid embedding. + graph_emb : dict + Edge/mesh embeddings, each entry of shape ``(B, *, d_h)``. + """ + 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) + + grid_emb = self.grid_prev_embedder(grid_features) + # (B, num_grid_nodes, d_h) + + # Graph embedding. NOTE: this block depends only on static graph + # features, so it is constant across an autoregressive rollout. It is + # kept as a self-contained block so a future embedd_graph()/ + # embedd_grid() split (hoisting it out of the AR loop) is mechanical. + 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 grid_emb, graph_emb + + def estimate_likelihood( + self, + latent_dist, + current_state, + last_state, + grid_prev_emb, + graph_emb, + loss_fn: Callable, + interior_mask: torch.Tensor, + ): + """ + Estimate the (masked) likelihood using the given distribution over + latent variables. + + ``loss_fn`` and ``interior_mask`` are passed in (not stored on the + predictor): masks live on the forecaster/module, which supplies its + own loss function and boolean interior mask. + + Parameters + ---------- + latent_dist : torch.distributions.Distribution + Shape ``(B, num_mesh_nodes, d_latent)``. + current_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Target ``X_{t+1}``. + last_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + grid_prev_emb : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid embedding from + ``embedd_grid_and_graph``. + graph_emb : dict + Edge/mesh embeddings from ``embedd_grid_and_graph``. + loss_fn : Callable + Per-entry loss (e.g. ``metrics.nll``); likelihood is its negative. + interior_mask : torch.Tensor + Boolean ``(num_grid_nodes,)`` mask of interior nodes. + + Returns + ------- + likelihood_term : torch.Tensor + Shape ``(B,)``. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. + pred_std : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)`` (decoder) or ``(d_state,)`` + (constant ``per_var_std``). + """ + # Sample from variational distribution + latent_samples = latent_dist.rsample() # (B, num_mesh_nodes, d_latent) + + # Compute reconstruction (decoder) + pred_mean, model_pred_std = self.decoder( + grid_prev_emb, latent_samples, last_state, graph_emb + ) # both (B, num_grid_nodes, d_state) + + if self.output_std: + pred_std = model_pred_std # (B, num_grid_nodes, d_state) + else: + # Use constant set std.-devs. + pred_std = self.per_var_std # (d_f,) + + # Compute likelihood (negative loss, exactly likelihood for nll loss) + # Note: There are some round-off errors here due to float32 + # and large values + entry_likelihoods = -loss_fn( + pred_mean, + current_state, + pred_std, + mask=interior_mask, + average_grid=False, + sum_vars=False, + ) # (B, num_grid_nodes', d_state) + likelihood_term = torch.sum(entry_likelihoods, dim=(1, 2)) # (B,) + return likelihood_term, pred_mean, pred_std + + def compute_step_loss( + self, + prev_states, + current_state, + forcing_features, + loss_fn: Callable, + interior_mask: torch.Tensor, + compute_kl: bool = True, + ): + """ + Forward pass and per-step ELBO pieces for one time step. + + Parameters + ---------- + prev_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. ``X_{t-1}, X_t``. + current_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Target ``X_{t+1}``. + forcing_features : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + loss_fn : Callable + Per-entry loss used to compute the likelihood term. + interior_mask : torch.Tensor + Boolean ``(num_grid_nodes,)`` mask of interior nodes. + compute_kl : bool + When False, skip the prior and return ``kl_term = None`` (the + ``kl_beta == 0`` / pure-autoencoder case). The KL weight itself is + a training knob owned by the calling module. + + Returns + ------- + likelihood_term : torch.Tensor + Shape ``(B,)``. + kl_term : torch.Tensor or None + Shape ``(B,)``, or None when ``compute_kl`` is False. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. + pred_std : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)`` or ``(d_state,)``. + """ + # embed all features + grid_prev_emb, graph_emb = self.embedd_grid_and_graph( + prev_states[:, 1], + prev_states[:, 0], + forcing_features, + ) + # embed also including current grid state, for encoder + grid_current_emb = self.embedd_grid_with_target( + prev_states[:, 1], + prev_states[:, 0], + forcing_features, + current_state, + ) # (B, num_grid_nodes, d_h) + + # Compute variational approximation (encoder) + var_dist = self.encoder( + grid_current_emb, graph_emb=graph_emb + ) # Gaussian, (B, num_mesh_nodes, d_latent) + + # Compute likelihood + last_state = prev_states[:, -1] + likelihood_term, pred_mean, pred_std = self.estimate_likelihood( + var_dist, + current_state, + last_state, + grid_prev_emb, + graph_emb, + loss_fn, + interior_mask, + ) + if compute_kl: + # Compute prior + prior_dist = self.prior_model( + grid_prev_emb, graph_emb=graph_emb + ) # Gaussian, (B, num_mesh_nodes, d_latent) + + # Compute KL + kl_term = torch.sum( + torch.distributions.kl_divergence(var_dist, prior_dist), + dim=(1, 2), + ) # (B,) + else: + # If KL is off, do not need to even compute prior nor KL + kl_term = None # Set to None to crash if erroneously used + + return likelihood_term, kl_term, 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: embed features, sample the latent + from the prior, decode, and return 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. + """ + # embed all features + grid_prev_emb, graph_emb = self.embedd_grid_and_graph( + prev_state, prev_prev_state, forcing + ) + + # Compute prior + prior_dist = self.prior_model( + grid_prev_emb, graph_emb=graph_emb + ) # (B, num_mesh_nodes, d_latent) + + # Sample from prior + latent_samples = prior_dist.rsample() + # (B, num_mesh_nodes, d_latent) + + # Compute reconstruction (decoder). prev_state (X_t) is the state the + # decoder adds its predicted residual onto. + pred_mean, pred_std = self.decoder( + grid_prev_emb, latent_samples, prev_state, graph_emb + ) # (B, num_grid_nodes, d_state) + + 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``. + """ + + requires_hierarchical = True + + def __init__( + self, + config: NeuralLAMConfig, + 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 hierarchical latent modules. + + Parameters + ---------- + config : NeuralLAMConfig + Full Neural-LAM configuration; used for the state feature + weighting that enters the constant per-variable std. + 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. + prior_intra_level_layers : int + Number of intra-level GNN layers in the (learned) prior. + 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. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + 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, a constant per-variable std is used as + likelihood scale. + 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__( + config=config, + datastore=datastore, + graph_name=graph_name, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + 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 + ] + # The latent variable lives on the top mesh level + self.num_mesh_nodes = level_mesh_sizes[-1] + 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) + ] + ) + + latent_dim = latent_dim if latent_dim is not None else hidden_dim + + # Prior. When learn_prior, the prior is a graph encoder mapping the + # previous state to a latent distribution; otherwise it is a constant + # (input-independent) Normal. + if learn_prior: + self.prior_model = 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_intra_level_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + else: + self.prior_model = ConstantLatentEncoder( + latent_dim=latent_dim, + num_mesh_nodes=self.num_mesh_nodes, + output_dist=prior_dist, + ) + + # Encoder (variational posterior) + Decoder + self.encoder = 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=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=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 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: + # Need a placeholder otherwise, just use raw features + mesh_emb["m2m"] = list(self.m2m_features) + + return mesh_emb + + +class GraphEFMMS(BaseGraphEFM): + """ + Graph-based Ensemble Forecasting Model on a flat mesh graph + (Graph-EFM-MS, e.g. for multi-scale graphs). + + The latent variable lives on the mesh nodes. The prior and variational + encoder are ``GraphLatentEncoder``s and the decoder is a + ``GraphLatentDecoder``. + """ + + requires_hierarchical = False + + def __init__( + self, + config: NeuralLAMConfig, + 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 flat-graph latent modules. + + Parameters + ---------- + config : NeuralLAMConfig + Full Neural-LAM configuration; used for the state feature + weighting that enters the constant per-variable std. + 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. + prior_m2m_layers : int + Number of on-mesh (m2m) GNN layers in the (learned) prior. + 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. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + 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, a constant per-variable std is used as + likelihood scale. + 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__( + config=config, + datastore=datastore, + graph_name=graph_name, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + 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, + ) + + self.num_mesh_nodes = self.mesh_static_features.shape[0] + utils.log_on_rank_zero( + f"Loaded graph with " + f"{self.num_grid_nodes + self.num_mesh_nodes} nodes " + f"({self.num_grid_nodes} grid, {self.num_mesh_nodes} 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) + + latent_dim = latent_dim if latent_dim is not None else hidden_dim + + # Prior. When learn_prior, the prior is a graph encoder mapping the + # previous state to a latent distribution; otherwise it is a constant + # (input-independent) Normal. + if learn_prior: + self.prior_model = 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_m2m_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + else: + self.prior_model = ConstantLatentEncoder( + latent_dim=latent_dim, + num_mesh_nodes=self.num_mesh_nodes, + output_dist=prior_dist, + ) + + # Encoder (variational posterior) + Decoder + self.encoder = 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=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=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 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/models/step_predictors/probabilistic.py b/neural_lam/models/step_predictors/probabilistic.py new file mode 100644 index 00000000..80d9a007 --- /dev/null +++ b/neural_lam/models/step_predictors/probabilistic.py @@ -0,0 +1,114 @@ +"""Interface for latent-variable step predictors.""" + +# Standard library +from abc import abstractmethod +from typing import Optional, Tuple + +# Third-party +import torch + +# Local +from .base import StepPredictor + + +class LatentStepPredictor(StepPredictor): + """ + Abstract step predictor for latent-variable models. + + Extends :class:`StepPredictor` for predictors that advance the state by + sampling a latent variable and decoding it. Subclasses implement + :meth:`step_distributions`, which returns the latent prior, the + variational posterior and the decoded reconstruction for one time step. + The evaluation :meth:`forward` is provided here as the prior-sampling + special case. The predictor produces only distributions and a + reconstruction; likelihood, KL and the ELBO are assembled outside it. + """ + + @abstractmethod + def step_distributions( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + target: Optional[torch.Tensor] = None, + compute_prior: bool = True, + ) -> Tuple[ + Optional[torch.distributions.Distribution], + Optional[torch.distributions.Distribution], + torch.Tensor, + Optional[torch.Tensor], + ]: + """ + Latent distributions and decoded reconstruction for one time step. + + The latent sample is drawn from the variational posterior when a + ``target`` is given (training path) and from the prior otherwise + (evaluation path). This is the only entry point that takes the + target. The predictor computes no likelihood or KL. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings + for this step. + target : torch.Tensor or None, optional + Shape ``(B, num_grid_nodes, num_state_vars)``. The true next state + ``X_{t+1}``. When given, the latent is sampled from the + variational posterior conditioned on it; when ``None``, from the + prior. + compute_prior : bool, optional + If ``False``, skip the prior on the training path (``prior_dist`` + is returned as ``None``); used when the KL weight is zero. On the + evaluation path the prior is always computed. + + Returns + ------- + prior_dist : torch.distributions.Distribution or None + Prior over the latent variable, or ``None`` when skipped. + posterior_dist : torch.distributions.Distribution or None + Variational posterior over the latent variable, or ``None`` on the + evaluation path. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Decoded mean given + the latent sample. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` + is True, otherwise ``None``. + """ + + def forward( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Advance one step by sampling the latent prior and decoding. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings + for this step. + + Returns + ------- + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted next state + ``X_{t+1}``. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` + is True, otherwise ``None``. + """ + _, _, pred_mean, pred_std = self.step_distributions( + prev_state, prev_prev_state, forcing, target=None + ) + return pred_mean, pred_std diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 942eb206..211beb56 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -9,11 +9,12 @@ import warnings from functools import cache from pathlib import Path -from typing import Any, Iterator, Union, overload +from typing import TYPE_CHECKING, Any, Iterator, Union, overload # Third-party import pytorch_lightning as pl import torch +import torch_geometric as pyg from loguru import logger from pytorch_lightning.loggers import MLFlowLogger, WandbLogger from pytorch_lightning.utilities import rank_zero_only @@ -23,6 +24,11 @@ # Local from .custom_loggers import CustomMLFlowLogger +if TYPE_CHECKING: + # Imported only for type checking to avoid a runtime import cycle + # Local + from .datastore import BaseDatastore + class BufferList(nn.Module): """ @@ -436,6 +442,83 @@ def loads_file(fn: str) -> Any: } +def load_and_register_graph( + module: nn.Module, + datastore: "BaseDatastore", + graph_name: str, +) -> 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. + + 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) + 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", + grid_static_dim: int, + 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 and forcing variables. + grid_static_dim : int + Number of static features per grid node. + 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") + return ( + 2 * num_state_vars + + grid_static_dim + + num_forcing_vars + * (num_past_forcing_steps + num_future_forcing_steps + 1) + ) + + def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: """ Construct a multilayer perceptron from a blueprint of layer widths. @@ -471,6 +554,72 @@ def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: 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) + ], + ) + + @cache def has_working_latex() -> bool: """ diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py new file mode 100644 index 00000000..063fd81a --- /dev/null +++ b/tests/test_graph_efm_predictor.py @@ -0,0 +1,267 @@ +"""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 (GraphEFMMS) and +hierarchical (GraphEFM) variants on the real example datastore with a freshly +created graph, then exercise ``forward``, ``compute_step_loss`` and the +sampling helpers on synthetic tensors. +""" + +# Standard library +from pathlib import Path + +# Third-party +import pytest +import torch + +# First-party +from neural_lam import config as nlconfig +from neural_lam import metrics +from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.loss_weighting import get_state_feature_weighting +from neural_lam.models.step_predictors.graph.graph_efm import ( + GraphEFM, + GraphEFMMS, +) +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 + + +def _datastore_and_config_with_graph(graph_name): + """Create the example datastore and ensure ``graph_name`` exists.""" + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, + config_path=datastore.root_path, + ) + ) + + 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, config + + +def _build_predictor(graph_name, output_std=False): + datastore, config = _datastore_and_config_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 = GraphEFMMS + layer_kwargs = { + "prior_m2m_layers": 1, + "encoder_m2m_layers": 1, + "decoder_m2m_layers": 1, + } + predictor = predictor_class( + config=config, + 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, config + + +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_compute_step_loss_shapes_and_kl_toggle(graph_name): + """compute_step_loss returns (likelihood (B,), kl, pred_mean, pred_std); + kl is a (B,) tensor when compute_kl=True and None when disabled.""" + predictor, datastore, _ = _build_predictor(graph_name) + prev_state, prev_prev_state, forcing, d_state = _make_inputs( + predictor, datastore + ) + B = prev_state.shape[0] + prev_states = torch.stack([prev_prev_state, prev_state], dim=1) + current_state = torch.randn(B, predictor.num_grid_nodes, d_state) + interior_mask = torch.ones(predictor.num_grid_nodes, dtype=torch.bool) + + # KL on + likelihood, kl, pred_mean, pred_std = predictor.compute_step_loss( + prev_states, + current_state, + forcing, + loss_fn=metrics.nll, + interior_mask=interior_mask, + compute_kl=True, + ) + assert likelihood.shape == (B,) + assert kl is not None + assert kl.shape == (B,) + assert pred_mean.shape == (B, predictor.num_grid_nodes, d_state) + # output_std=False -> constant per-variable std (d_state,) + assert pred_std.shape == (d_state,) + + # KL off -> kl_term is None + likelihood_off, kl_off, _, _ = predictor.compute_step_loss( + prev_states, + current_state, + forcing, + loss_fn=metrics.nll, + interior_mask=interior_mask, + compute_kl=False, + ) + assert kl_off is None + assert likelihood_off.shape == (B,) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_compute_step_loss_is_differentiable(graph_name): + """The ELBO pieces are differentiable through the rsample paths, and the + gradient reaches encoder, decoder and prior parameters.""" + predictor, datastore, _ = _build_predictor(graph_name) + prev_state, prev_prev_state, forcing, d_state = _make_inputs( + predictor, datastore + ) + B = prev_state.shape[0] + prev_states = torch.stack([prev_prev_state, prev_state], dim=1) + current_state = torch.randn(B, predictor.num_grid_nodes, d_state) + interior_mask = torch.ones(predictor.num_grid_nodes, dtype=torch.bool) + + likelihood, kl, _, _ = predictor.compute_step_loss( + prev_states, + current_state, + forcing, + loss_fn=metrics.nll, + interior_mask=interior_mask, + compute_kl=True, + ) + elbo = (likelihood - kl).mean() + elbo.backward() + + for module in (predictor.encoder, predictor.decoder, predictor.prior_model): + assert any( + p.grad is not None and torch.any(p.grad != 0) + for p in module.parameters() + ), f"no gradient reached {module.__class__.__name__}" + + +@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) + + +def test_per_var_std_matches_module_formula(): + """per_var_std mirrors ForecasterModule's formula: + state_diff_std_standardized / sqrt(state_feature_weights).""" + predictor, datastore, config = _build_predictor("1level") + + da_state_stats = datastore.get_standardization_dataarray(category="state") + diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + feature_weights = torch.tensor( + get_state_feature_weighting(config=config, datastore=datastore), + dtype=torch.float32, + ) + expected = diff_std / torch.sqrt(feature_weights) + + assert predictor.per_var_std is not None + assert torch.allclose(predictor.per_var_std, expected) + + +def test_per_var_std_none_when_output_std(): + """When the decoder outputs its own std, the constant per_var_std is unused + and left as None (mirrors ForecasterModule).""" + predictor, _, _ = _build_predictor("1level", output_std=True) + assert predictor.per_var_std is None + + +@pytest.mark.parametrize( + "predictor_class, graph_name", + [(GraphEFM, "1level"), (GraphEFMMS, "hierarchical")], +) +def test_graph_type_mismatch_raises(predictor_class, graph_name): + """GraphEFM requires a hierarchical graph and GraphEFMMS a flat one; + constructing with the wrong graph type raises ValueError.""" + datastore, config = _datastore_and_config_with_graph(graph_name) + with pytest.raises(ValueError, match="mesh graph"): + predictor_class( + config=config, + 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..377182d1 --- /dev/null +++ b/tests/test_latent_modules.py @@ -0,0 +1,581 @@ +"""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"] + ) + last_state = torch.randn( + B, flat_dims["num_grid"], flat_dims["num_state_vars"] + ) + + pred_mean, pred_std = dec( + grid_rep, latent_samples, last_state, flat_graph_emb + ) + + expected_shape = (B, flat_dims["num_grid"], flat_dims["num_state_vars"]) + assert pred_mean.shape == expected_shape + assert pred_std is not None + assert pred_std.shape == expected_shape + assert (pred_std > 0).all() + + (pred_mean.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"] + ) + last_state = torch.randn( + B, flat_dims["num_grid"], flat_dims["num_state_vars"] + ) + + pred_mean, pred_std = dec( + grid_rep, latent_samples, last_state, flat_graph_emb + ) + assert pred_mean.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"] + ) + last_state = torch.randn( + B, flat_dims["num_grid"], flat_dims["num_state_vars"] + ) + pred_mean, _ = dec(grid_rep, latent_samples, last_state, flat_graph_emb) + assert pred_mean.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"]) + last_state = torch.randn(B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + + pred_mean, pred_std = dec( + grid_rep, latent_samples, last_state, hi_graph_emb + ) + expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + assert pred_mean.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) + last_state = torch.randn(B, num_grid, num_state_vars) + + pred_mean, pred_std = dec(grid_rep, latent_samples, last_state, graph_emb) + assert pred_mean.shape == (B, num_grid, num_state_vars) + assert pred_std.shape == (B, num_grid, num_state_vars) + + (pred_mean.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"]) + last_state = torch.randn(B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + + pred_mean, pred_std = dec( + grid_rep, latent_samples, last_state, hi_graph_emb + ) + expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) + assert pred_mean.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_probabilistic_interface.py b/tests/test_probabilistic_interface.py new file mode 100644 index 00000000..6366f900 --- /dev/null +++ b/tests/test_probabilistic_interface.py @@ -0,0 +1,256 @@ +"""Interface-level checks for the latent-variable forecasting stack. + +Drives a minimal, graph-free dummy ``LatentStepPredictor`` through +``ProbabilisticARForecaster`` and ``ProbabilisticForecasterModule`` to verify +the predictor -> forecaster -> module -> ELBO contract independently of any +concrete predictor implementation. The Graph-EFM wiring is tested separately. +""" + +# Standard library +import warnings + +# Third-party +import pytest +import pytorch_lightning as pl +import torch +import wandb +from torch import nn + +# First-party +from neural_lam import config as nlconfig +from neural_lam.models import ( + LatentStepPredictor, + ProbabilisticARForecaster, + ProbabilisticForecasterModule, +) +from neural_lam.weather_dataset import WeatherDataModule +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 +NUM_MESH_NODES = 3 +LATENT_DIM = 4 + + +class _DummyLatentPredictor(LatentStepPredictor): + """Minimal graph-free latent predictor exercising the interface.""" + + def __init__(self, datastore): + """ + Build a graph-free latent predictor over ``datastore``. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing grid metadata and state-variable counts. + """ + super().__init__(datastore=datastore, output_std=False) + d_state = datastore.get_num_data_vars(category="state") + self.encode_latent = nn.Linear(d_state, LATENT_DIM) + self.decode_latent = nn.Linear(LATENT_DIM, d_state) + + def _latent_dist(self, state): + """ + Build a diagonal-Normal latent distribution from a grid state. + + Parameters + ---------- + state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. State summarised + into latent parameters by mean-pooling over the grid. + + Returns + ------- + torch.distributions.Normal + Shape ``(B, NUM_MESH_NODES, LATENT_DIM)``. Unit-variance Normal. + """ + mean = self.encode_latent(state.mean(dim=1)) # (B, LATENT_DIM) + mean = mean.unsqueeze(1).expand(-1, NUM_MESH_NODES, -1) + return torch.distributions.Normal(mean, torch.ones_like(mean)) + + def step_distributions( + self, + prev_state, + prev_prev_state, + forcing, + target=None, + compute_prior=True, + ): + """ + Latent distributions and reconstruction for one step. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``; + unused by this dummy. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``; unused. + target : torch.Tensor or None, optional + Shape ``(B, num_grid_nodes, num_state_vars)``. True ``X_{t+1}``. + compute_prior : bool, optional + Whether to build the prior on the training path. + + Returns + ------- + prior_dist : torch.distributions.Normal or None + Prior over the latent, or ``None`` when skipped. + posterior_dist : torch.distributions.Normal or None + Posterior over the latent, or ``None`` on the eval path. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Decoded mean. + pred_std : None + This dummy never emits an std. + """ + if target is not None: + posterior_dist = self._latent_dist(target) + latent = posterior_dist.rsample() + prior_dist = ( + self._latent_dist(prev_state) if compute_prior else None + ) + else: + posterior_dist = None + prior_dist = self._latent_dist(prev_state) + latent = prior_dist.rsample() + + delta = self.decode_latent(latent.mean(dim=1)).unsqueeze(1) + pred_mean = prev_state + delta + return prior_dist, posterior_dist, pred_mean, None + + +def _build_module(kl_beta=1.0): + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = _DummyLatentPredictor(datastore) + forecaster = ProbabilisticARForecaster(predictor, datastore) + module = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="nll", + kl_beta=kl_beta, + ) + return datastore, module + + +def _make_batch(datastore, pred_steps=3, batch_size=2): + num_grid = datastore.num_grid_points + 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, d_state) + target_states = torch.randn(batch_size, pred_steps, num_grid, d_state) + forcing = torch.randn(batch_size, pred_steps, num_grid, d_forcing) + return init_states, target_states, forcing + + +@pytest.mark.parametrize("kl_beta", [1.0, 0.0]) +def test_training_rollout_terms(kl_beta): + """training_rollout returns per-step means/kl; kl is None when β == 0.""" + datastore, module = _build_module(kl_beta=kl_beta) + init_states, target_states, forcing = _make_batch(datastore) + + pred_means, pred_stds, kl_terms = module.forecaster.training_rollout( + init_states, + forcing, + target_states, + compute_kl=kl_beta > 0, + ) + b, t = init_states.shape[0], forcing.shape[1] + assert pred_means.shape == ( + b, + t, + datastore.num_grid_points, + datastore.get_num_data_vars(category="state"), + ) + assert pred_stds is None # output_std=False + if kl_beta > 0: + assert kl_terms.shape == (b, t) + else: + assert kl_terms is None + + +def test_sample_trajectories_shape(): + """sample_trajectories stacks an ensemble dimension of prior rollouts.""" + datastore, module = _build_module() + init_states, target_states, forcing = _make_batch(datastore) + + num_traj = 4 + traj_means, traj_stds = module.forecaster.sample_trajectories( + init_states, forcing, target_states, num_traj + ) + b, t = init_states.shape[0], forcing.shape[1] + assert traj_means.shape == ( + b, + num_traj, + t, + datastore.num_grid_points, + datastore.get_num_data_vars(category="state"), + ) + assert traj_stds is None # output_std=False + + +def test_elbo_loss_is_finite_and_differentiable(): + """The assembled ELBO loss is finite and gradients reach the predictor.""" + datastore, module = _build_module() + init_states, target_states, forcing = _make_batch(datastore) + + pred_means, _, kl_terms = module.forecaster.training_rollout( + init_states, + forcing, + target_states, + compute_kl=True, + ) + entry_log_lik = -module.loss( + pred_means, + target_states, + module.per_var_std, + mask=module.interior_mask_bool, + average_grid=False, + sum_vars=False, + ) + likelihood = entry_log_lik.sum(dim=(2, 3)).sum(dim=1) # (B,) + loss = -likelihood.mean() + module.kl_beta * kl_terms.sum(dim=1).mean() + + assert torch.isfinite(loss) + loss.backward() + assert any( + p.grad is not None and torch.any(p.grad != 0) + for p in module.forecaster.predictor.parameters() + ) + + +def test_end_to_end_probabilistic_training_runs(): + """A full Trainer epoch (train + val) runs through the whole pipeline.""" + datastore, module = _build_module() + + data_module = WeatherDataModule( + datastore=datastore, + ar_steps_train=3, + ar_steps_eval=5, + batch_size=2, + num_workers=1, + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + ) + + trainer = pl.Trainer( + max_epochs=1, + accelerator="cpu", + devices=1, + log_every_n_steps=1, + enable_checkpointing=False, + logger=False, + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + wandb.init(mode="disabled") + trainer.fit(module, data_module)