From e8dcc66f88412bca62d6d9435f53388cea0bab20 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 28 May 2026 05:27:21 +0530 Subject: [PATCH 01/51] feat: add latent encoder/decoder infrastructure for Graph-EFM port Adds neural_lam/models/latent/ with the encoder and decoder submodules needed by the probabilistic GraphEFM model (issue #62). Ported from the prob_model_lam branch with adaptations for the current main architecture: - constants.GRID_STATE_DIM replaced by a num_state_vars constructor arg - interaction_net imports updated to neural_lam.gnn_layers - GraphLatentDecoder.processor unified with the other four GNN-seq constructions to use utils.make_gnn_seq (handles processor_layers=0) - HiGraph{Encoder,Decoder} guard against single-level meshes where the latent variable would be silently ignored - ConstantLatentEncoder docstring documents the N(1,1) vs N(0,1) discrepancy with the prob_model_lam CLI help (open question upstream) Also adds to neural_lam/utils.py: - IdentityModule: pass-through nn.Module for multi-arg sequential GNNs - make_gnn_seq: builds a pyg.nn.Sequential of InteractionNets, or an IdentityModule when num_gnn_layers=0; lazy-imports gnn_layers to avoid the existing gnn_layers -> utils circular dependency 17 tests in tests/test_latent_modules.py cover output shapes, distribution properties, backprop to every parameter, 2- and 3-level hierarchical graphs, intra_level_layers=0, and the single-level guard. --- neural_lam/models/latent/__init__.py | 18 + neural_lam/models/latent/base_decoder.py | 96 ++++ neural_lam/models/latent/base_encoder.py | 65 +++ neural_lam/models/latent/constant_encoder.py | 35 ++ neural_lam/models/latent/graph_decoder.py | 70 +++ neural_lam/models/latent/graph_encoder.py | 61 +++ neural_lam/models/latent/hi_graph_decoder.py | 187 +++++++ neural_lam/models/latent/hi_graph_encoder.py | 120 +++++ neural_lam/utils.py | 35 ++ tests/test_latent_modules.py | 539 +++++++++++++++++++ 10 files changed, 1226 insertions(+) create mode 100644 neural_lam/models/latent/__init__.py create mode 100644 neural_lam/models/latent/base_decoder.py create mode 100644 neural_lam/models/latent/base_encoder.py create mode 100644 neural_lam/models/latent/constant_encoder.py create mode 100644 neural_lam/models/latent/graph_decoder.py create mode 100644 neural_lam/models/latent/graph_encoder.py create mode 100644 neural_lam/models/latent/hi_graph_decoder.py create mode 100644 neural_lam/models/latent/hi_graph_encoder.py create mode 100644 tests/test_latent_modules.py diff --git a/neural_lam/models/latent/__init__.py b/neural_lam/models/latent/__init__.py new file mode 100644 index 00000000..f50d2ac6 --- /dev/null +++ b/neural_lam/models/latent/__init__.py @@ -0,0 +1,18 @@ +# 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..5f72beca --- /dev/null +++ b/neural_lam/models/latent/base_decoder.py @@ -0,0 +1,96 @@ +# 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 + softplus std) depending on + ``output_std``. + """ + + def __init__( + self, + hidden_dim, + latent_dim, + num_state_vars, + hidden_layers=1, + output_std=True, + ): + 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. + + original_grid_rep: (B, num_grid_nodes, d_h) + latent_rep: (B, num_mesh_nodes, d_h) + residual_grid_rep: (B, num_grid_nodes, d_h) + graph_emb: dict of graph edge / node embeddings + + Returns: + combined_grid_rep: (B, num_grid_nodes, d_h) + """ + 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. + + grid_rep: (B, num_grid_nodes, d_h) + latent_samples: (B, num_mesh_nodes, latent_dim) + last_state: (B, num_grid_nodes, num_state_vars) + graph_emb: dict with at least ``g2m``, ``m2m``, ``m2g`` entries + + Returns: + pred_mean: (B, num_grid_nodes, num_state_vars) + pred_std: (B, num_grid_nodes, num_state_vars) or ``None`` + """ + latent_emb = self.latent_embedder(latent_samples) + + 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: + 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..889014b0 --- /dev/null +++ b/neural_lam/models/latent/base_encoder.py @@ -0,0 +1,65 @@ +# 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"): + 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. + + grid_rep: (B, num_grid_nodes, d_h) + + Returns: + parameters: (B, num_mesh_nodes, output_dim) + """ + raise NotImplementedError("compute_dist_params not implemented") + + def forward(self, grid_rep, **kwargs): + """ + Compute the Gaussian distribution over the latent variable. + + grid_rep: (B, num_grid_nodes, d_h) + + Returns: + distribution: ``torch.distributions.Normal`` of 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..403b7cb6 --- /dev/null +++ b/neural_lam/models/latent/constant_encoder.py @@ -0,0 +1,35 @@ +# Third-party +import torch + +# Local +from .base_encoder import BaseLatentEncoder + + +class ConstantLatentEncoder(BaseLatentEncoder): + """ + Latent encoder that returns a constant (input-independent) distribution. + + Used as a non-learned prior in ``GraphEFM`` when ``learn_prior`` is + disabled. ``compute_dist_params`` returns a tensor of ones, so the + resulting Normal is ``Normal(mean=1, std=1)`` for ``output_dist= + "isotropic"`` and ``Normal(mean=1, std=softplus(1)+eps)`` for + ``output_dist="diagonal"``. (Note: the ``train_model.py`` CLI help on + ``prob_model_lam`` describes this prior as "mean 0"; the code itself + has always produced mean 1. Preserved as-is during the port for + behavioral parity — open question for upstream.) + """ + + def __init__(self, latent_dim, num_mesh_nodes, output_dist="isotropic"): + super().__init__(latent_dim, output_dist) + self.num_mesh_nodes = num_mesh_nodes + + def compute_dist_params(self, grid_rep, **kwargs): + """ + Return constant parameters of shape (B, num_mesh_nodes, output_dim). + """ + return torch.ones( + 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..7db8fe96 --- /dev/null +++ b/neural_lam/models/latent/graph_decoder.py @@ -0,0 +1,70 @@ +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import InteractionNet, PropagationNet + +# Local +from .base_decoder import BaseGraphLatentDecoder + + +class GraphLatentDecoder(BaseGraphLatentDecoder): + """ + Latent decoder for a flat (non-hierarchical) graph. Encodes grid into + mesh with an InteractionNet, processes on mesh, and reads back out to + grid with a PropagationNet. 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, + processor_layers, + hidden_layers=1, + output_std=True, + ): + super().__init__( + hidden_dim, latent_dim, num_state_vars, hidden_layers, output_std + ) + + self.g2m_gnn = InteractionNet( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + self.processor = utils.make_gnn_seq( + m2m_edge_index, processor_layers, hidden_layers, hidden_dim + ) + + self.m2g_gnn = PropagationNet( + 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 -> processor -> m2g. + + original_grid_rep: (B, num_grid_nodes, d_h) + latent_rep: (B, num_mesh_nodes, d_h) + residual_grid_rep: (B, num_grid_nodes, d_h) + + Returns: + grid_rep: (B, num_grid_nodes, d_h) + """ + mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) + + mesh_rep, _ = self.processor(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..8c14054c --- /dev/null +++ b/neural_lam/models/latent/graph_encoder.py @@ -0,0 +1,61 @@ +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import PropagationNet + +# 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 PropagationNet followed by a stack of + on-mesh InteractionNet processor layers. + """ + + def __init__( + self, + latent_dim, + g2m_edge_index, + m2m_edge_index, + hidden_dim, + processor_layers, + hidden_layers=1, + output_dist="isotropic", + ): + super().__init__(latent_dim, output_dist) + + self.g2m_gnn = PropagationNet( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + self.processor = utils.make_gnn_seq( + m2m_edge_index, processor_layers, hidden_layers, hidden_dim + ) + + 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. + + grid_rep: (B, num_grid_nodes, d_h) + graph_emb: dict with at least + - ``mesh``: (B, num_mesh_nodes, d_h) + - ``g2m``: (B, M_g2m, d_h) + - ``m2m``: (B, M_m2m, d_h) + + Returns: + parameters: (B, num_mesh_nodes, output_dim) + """ + mesh_rep = self.g2m_gnn(grid_rep, graph_emb["mesh"], graph_emb["g2m"]) + mesh_rep, _ = self.processor(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..434faed0 --- /dev/null +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -0,0 +1,187 @@ +# Third-party +from torch import nn + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import InteractionNet, PropagationNet + +# 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 a PropagationNet. + """ + + 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, + output_std=True, + ): + 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 = InteractionNet( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + self.m2g_gnn = PropagationNet( + m2g_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + 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 + ] + ) + 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 + ] + ) + + # Identity mappings if intra_level_layers == 0 + 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 + ] + ) + 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 + ] + ) + + 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. + + original_grid_rep: (B, num_grid_nodes, d_h) + latent_rep: (B, num_mesh_nodes[L], d_h) + residual_grid_rep: (B, num_grid_nodes, d_h) + graph_emb: dict with at least + - ``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) + - ``m2g``: (B, M_m2g, d_h) + + Returns: + grid_rep: (B, num_grid_nodes, d_h) + """ + 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 ( + up_gnn, + intra_gnn_seq, + mesh_up_level_rep, + m2m_level_rep, + mesh_level_rep, + ) in zip( + self.mesh_up_gnns, + self.intra_up_gnns[:-1], + graph_emb["mesh_up"], + graph_emb["m2m"][:-1], + graph_emb["mesh"][1:-1] + [latent_rep], + ): + new_mesh_rep, new_m2m_rep = intra_gnn_seq( + current_mesh_rep, m2m_level_rep + ) + + 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 + 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 ( + down_gnn, + intra_gnn_seq, + mesh_down_level_rep, + m2m_level_rep, + mesh_level_rep, + ) in zip( + reversed(self.mesh_down_gnns), + reversed(self.intra_down_gnns), + reversed(graph_emb["mesh_down"]), + reversed(m2m_level_reps), + reversed(mesh_level_reps), + ): + new_mesh_rep = down_gnn( + current_mesh_rep, mesh_level_rep, mesh_down_level_rep + ) + current_mesh_rep, _ = intra_gnn_seq(new_mesh_rep, m2m_level_rep) + + 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..ee57b8ee --- /dev/null +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -0,0 +1,120 @@ +# Third-party +from torch import nn + +# First-party +from neural_lam import utils +from neural_lam.gnn_layers import PropagationNet + +# Local +from .base_encoder import BaseLatentEncoder + + +class HiGraphLatentEncoder(BaseLatentEncoder): + """ + Latent encoder for a hierarchical mesh: grid -> bottom mesh level via a + PropagationNet, then propagates upward through mesh levels using + 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, + output_dist="isotropic", + ): + 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 = PropagationNet( + g2m_edge_index, + hidden_dim, + hidden_layers=hidden_layers, + update_edges=False, + ) + + 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 + ] + ) + + # Identity mappings if intra_level_layers == 0 + 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 + ] + ) + + 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. + + grid_rep: (B, num_grid_nodes, d_h) + graph_emb: dict with at least + - ``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) + + Returns: + parameters: (B, num_mesh_nodes[L], output_dim) + """ + current_mesh_rep = self.g2m_gnn( + grid_rep, graph_emb["mesh"][0], graph_emb["g2m"] + ) + + # Same-level processing on level 0 + current_mesh_rep, _ = self.intra_level_gnns[0]( + current_mesh_rep, graph_emb["m2m"][0] + ) + + # Walk up levels 1..L + for ( + up_gnn, + intra_gnn_seq, + mesh_up_level_rep, + m2m_level_rep, + mesh_level_rep, + ) in zip( + self.mesh_up_gnns, + self.intra_level_gnns[1:], + graph_emb["mesh_up"], + graph_emb["m2m"][1:], + graph_emb["mesh"][1:], + ): + new_node_rep = up_gnn( + current_mesh_rep, mesh_level_rep, mesh_up_level_rep + ) + current_mesh_rep, _ = intra_gnn_seq(new_node_rep, m2m_level_rep) + + return self.latent_param_map(current_mesh_rep) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 942eb206..48c280ee 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -14,6 +14,7 @@ # 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 @@ -471,6 +472,40 @@ def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: return nn.Sequential(*layers) +class IdentityModule(nn.Module): + """Identity operator that accepts and returns multiple positional inputs.""" + + def forward(self, *args): + return args + + +def make_gnn_seq(edge_index, num_gnn_layers, hidden_layers, hidden_dim): + """ + Build a sequential stack of InteractionNet layers that propagates both + node and edge representations. Returns an IdentityModule if + num_gnn_layers is 0. + """ + # First-party + from neural_lam.gnn_layers import InteractionNet + + if num_gnn_layers == 0: + return IdentityModule() + return pyg.nn.Sequential( + "mesh_rep, edge_rep", + [ + ( + InteractionNet( + 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_latent_modules.py b/tests/test_latent_modules.py new file mode 100644 index 00000000..3b845ed4 --- /dev/null +++ b/tests/test_latent_modules.py @@ -0,0 +1,539 @@ +"""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 IdentityModule, 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, + "processor_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_identity_module_passes_args_through(): + module = IdentityModule() + a, b, c = torch.randn(3), torch.randn(2), torch.randn(1) + out = module(a, b, c) + assert out == (a, b, c) + + +def test_make_gnn_seq_zero_layers_returns_identity(): + edge_index = _fully_connected_edge_index(3, 3) + seq = make_gnn_seq( + edge_index, num_gnn_layers=0, hidden_layers=1, hidden_dim=8 + ) + assert isinstance(seq, IdentityModule) + + 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 torch.equal(out_mesh, mesh_rep) + assert torch.equal(out_edge, edge_rep) + + +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) + + +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"], + processor_layers=flat_dims["processor_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"], + processor_layers=flat_dims["processor_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"], + processor_layers=flat_dims["processor_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 + + +# --- 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 routes intra-processing through IdentityModule + (the make_gnn_seq branch). Exercise that path 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"], + ) From 6c23ff371186eeed7b0c2f5f2541754bc7392824 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 5 Jun 2026 06:31:22 +0530 Subject: [PATCH 02/51] Update neural_lam/models/latent/constant_encoder.py Co-authored-by: Joel Oskarsson --- neural_lam/models/latent/constant_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 403b7cb6..27aef1da 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -27,7 +27,7 @@ def compute_dist_params(self, grid_rep, **kwargs): """ Return constant parameters of shape (B, num_mesh_nodes, output_dim). """ - return torch.ones( + return torch.zeros( grid_rep.shape[0], self.num_mesh_nodes, self.output_dim, From 3cebe5abfa8dcdd1c6d8ea4e48975191f2c27475 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 5 Jun 2026 17:02:45 +0530 Subject: [PATCH 03/51] feat: address review feedback on latent encoder/decoder infra Make GNN types configurable and tidy up the latent modules per PR review: - make_gnn_seq: accept a gnn_type arg (resolved via get_gnn_class) so it is not limited to InteractionNet, and make it strict (raise on num_gnn_layers < 1) instead of silently returning an IdentityModule; callers now own the no-op (identity) case explicitly. - graph/hi encoders and decoders: expose g2m/m2g/mesh_up/mesh_down gnn_type parameters wired to get_gnn_class, with defaults matching prob_model_lam. - graph encoder/decoder: rename processor_layers -> m2m_layers (and the self.processor attribute -> self.m2m_gnns); "processor" was misleading in an encoder/decoder context. - ConstantLatentEncoder: return zeros instead of ones so the static prior is mean 0 (fixes the prob_model_lam mean-1 bug; matches its own CLI help). - tests: update for the renamed arg and strict make_gnn_seq, add coverage for the flat zero-m2m identity path, and assert the constant prior is N(0, 1). --- neural_lam/models/latent/constant_encoder.py | 15 ++-- neural_lam/models/latent/graph_decoder.py | 29 ++++--- neural_lam/models/latent/graph_encoder.py | 21 +++-- neural_lam/models/latent/hi_graph_decoder.py | 24 ++++-- neural_lam/models/latent/hi_graph_encoder.py | 18 +++-- neural_lam/utils.py | 31 ++++++-- tests/test_latent_modules.py | 80 ++++++++++++++++---- 7 files changed, 157 insertions(+), 61 deletions(-) diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 27aef1da..600f0a09 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -10,13 +10,14 @@ class ConstantLatentEncoder(BaseLatentEncoder): Latent encoder that returns a constant (input-independent) distribution. Used as a non-learned prior in ``GraphEFM`` when ``learn_prior`` is - disabled. ``compute_dist_params`` returns a tensor of ones, so the - resulting Normal is ``Normal(mean=1, std=1)`` for ``output_dist= - "isotropic"`` and ``Normal(mean=1, std=softplus(1)+eps)`` for - ``output_dist="diagonal"``. (Note: the ``train_model.py`` CLI help on - ``prob_model_lam`` describes this prior as "mean 0"; the code itself - has always produced mean 1. Preserved as-is during the port for - behavioral parity — open question for upstream.) + disabled. ``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"``. (Note: ``prob_model_lam`` returned a tensor + of ones here, giving mean 1, while its ``train_model.py`` CLI help + described the prior as "mean 0". The mean 1 was a bug -- it is only a + constant offset, but a mean-0 prior is what is intended, so the port + uses zeros.) """ def __init__(self, latent_dim, num_mesh_nodes, output_dist="isotropic"): diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index 7db8fe96..b613878b 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -1,6 +1,6 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import InteractionNet, PropagationNet +from neural_lam.gnn_layers import get_gnn_class # Local from .base_decoder import BaseGraphLatentDecoder @@ -9,9 +9,10 @@ class GraphLatentDecoder(BaseGraphLatentDecoder): """ Latent decoder for a flat (non-hierarchical) graph. Encodes grid into - mesh with an InteractionNet, processes on mesh, and reads back out to - grid with a PropagationNet. The grid representation also goes through a - residual MLP that is added back to the mesh-to-grid output. + 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__( @@ -22,26 +23,32 @@ def __init__( hidden_dim, latent_dim, num_state_vars, - processor_layers, + m2m_layers, hidden_layers=1, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="PropagationNet", output_std=True, ): super().__init__( hidden_dim, latent_dim, num_state_vars, hidden_layers, output_std ) - self.g2m_gnn = InteractionNet( + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( g2m_edge_index, hidden_dim, hidden_layers=hidden_layers, update_edges=False, ) - self.processor = utils.make_gnn_seq( - m2m_edge_index, processor_layers, hidden_layers, hidden_dim + self.m2m_gnns = ( + utils.make_gnn_seq( + m2m_edge_index, m2m_layers, hidden_layers, hidden_dim + ) + if m2m_layers > 0 + else utils.IdentityModule() ) - self.m2g_gnn = PropagationNet( + self.m2g_gnn = get_gnn_class(m2g_gnn_type)( m2g_edge_index, hidden_dim, hidden_layers=hidden_layers, @@ -52,7 +59,7 @@ def combine_with_latent( self, original_grid_rep, latent_rep, residual_grid_rep, graph_emb ): """ - Fuse grid and latent reps via g2m -> processor -> m2g. + Fuse grid and latent reps via g2m -> m2m -> m2g. original_grid_rep: (B, num_grid_nodes, d_h) latent_rep: (B, num_mesh_nodes, d_h) @@ -63,7 +70,7 @@ def combine_with_latent( """ mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) - mesh_rep, _ = self.processor(mesh_rep, graph_emb["m2m"]) + mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) grid_rep = self.m2g_gnn(mesh_rep, residual_grid_rep, graph_emb["m2g"]) diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py index 8c14054c..1dd824a3 100644 --- a/neural_lam/models/latent/graph_encoder.py +++ b/neural_lam/models/latent/graph_encoder.py @@ -1,6 +1,6 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import PropagationNet +from neural_lam.gnn_layers import get_gnn_class # Local from .base_encoder import BaseLatentEncoder @@ -10,8 +10,8 @@ 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 PropagationNet followed by a stack of - on-mesh InteractionNet processor layers. + (non-hierarchical) graph: one g2m GNN (type set by ``g2m_gnn_type``) + followed by a stack of on-mesh (m2m) InteractionNet layers. """ def __init__( @@ -20,21 +20,26 @@ def __init__( g2m_edge_index, m2m_edge_index, hidden_dim, - processor_layers, + m2m_layers, hidden_layers=1, + g2m_gnn_type="PropagationNet", output_dist="isotropic", ): super().__init__(latent_dim, output_dist) - self.g2m_gnn = PropagationNet( + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( g2m_edge_index, hidden_dim, hidden_layers=hidden_layers, update_edges=False, ) - self.processor = utils.make_gnn_seq( - m2m_edge_index, processor_layers, hidden_layers, hidden_dim + self.m2m_gnns = ( + utils.make_gnn_seq( + m2m_edge_index, m2m_layers, hidden_layers, hidden_dim + ) + if m2m_layers > 0 + else utils.IdentityModule() ) self.latent_param_map = utils.make_mlp( @@ -57,5 +62,5 @@ def compute_dist_params(self, grid_rep, graph_emb, **kwargs): parameters: (B, num_mesh_nodes, output_dim) """ mesh_rep = self.g2m_gnn(grid_rep, graph_emb["mesh"], graph_emb["g2m"]) - mesh_rep, _ = self.processor(mesh_rep, graph_emb["m2m"]) + 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 index 434faed0..906ce7d6 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -3,7 +3,7 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import InteractionNet, PropagationNet +from neural_lam.gnn_layers import get_gnn_class # Local from .base_decoder import BaseGraphLatentDecoder @@ -16,7 +16,9 @@ class HiGraphLatentDecoder(BaseGraphLatentDecoder): *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 a PropagationNet. + via an m2g GNN (type set by ``m2g_gnn_type``). The g2m, mesh-up and + mesh-down GNN types are set by ``g2m_gnn_type``, ``mesh_up_gnn_type`` + and ``mesh_down_gnn_type`` respectively. """ def __init__( @@ -31,6 +33,10 @@ def __init__( num_state_vars, intra_level_layers, hidden_layers=1, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="PropagationNet", + mesh_up_gnn_type="InteractionNet", + mesh_down_gnn_type="PropagationNet", output_std=True, ): super().__init__( @@ -47,22 +53,23 @@ def __init__( "flat graphs." ) - self.g2m_gnn = InteractionNet( + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( g2m_edge_index, hidden_dim, hidden_layers=hidden_layers, update_edges=False, ) - self.m2g_gnn = PropagationNet( + self.m2g_gnn = get_gnn_class(m2g_gnn_type)( m2g_edge_index, hidden_dim, hidden_layers=hidden_layers, update_edges=False, ) + mesh_up_class = get_gnn_class(mesh_up_gnn_type) self.mesh_up_gnns = nn.ModuleList( [ - InteractionNet( + mesh_up_class( edge_index, hidden_dim, hidden_layers=hidden_layers, @@ -71,9 +78,10 @@ def __init__( for edge_index in mesh_up_edge_index ] ) + mesh_down_class = get_gnn_class(mesh_down_gnn_type) self.mesh_down_gnns = nn.ModuleList( [ - PropagationNet( + mesh_down_class( edge_index, hidden_dim, hidden_layers=hidden_layers, @@ -89,6 +97,8 @@ def __init__( utils.make_gnn_seq( edge_index, intra_level_layers, hidden_layers, hidden_dim ) + if intra_level_layers > 0 + else utils.IdentityModule() for edge_index in m2m_edge_index ] ) @@ -97,6 +107,8 @@ def __init__( utils.make_gnn_seq( edge_index, intra_level_layers, hidden_layers, hidden_dim ) + if intra_level_layers > 0 + else utils.IdentityModule() for edge_index in list(m2m_edge_index)[:-1] # Top level (L) does not need a down intra-level GNN ] diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index ee57b8ee..73677634 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -3,7 +3,7 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import PropagationNet +from neural_lam.gnn_layers import get_gnn_class # Local from .base_encoder import BaseLatentEncoder @@ -12,9 +12,10 @@ class HiGraphLatentEncoder(BaseLatentEncoder): """ Latent encoder for a hierarchical mesh: grid -> bottom mesh level via a - PropagationNet, then propagates upward through mesh levels using - PropagationNets, with optional intra-level processing at each level. - The latent distribution is read out from the top mesh level. + g2m GNN (type set by ``g2m_gnn_type``), then propagates upward through + mesh levels using mesh-up GNNs (type set by ``mesh_up_gnn_type``), with + optional intra-level processing at each level. The latent distribution + is read out from the top mesh level. """ def __init__( @@ -26,6 +27,8 @@ def __init__( hidden_dim, intra_level_layers, hidden_layers=1, + g2m_gnn_type="PropagationNet", + mesh_up_gnn_type="PropagationNet", output_dist="isotropic", ): super().__init__(latent_dim, output_dist) @@ -40,16 +43,17 @@ def __init__( "flat graphs." ) - self.g2m_gnn = PropagationNet( + self.g2m_gnn = get_gnn_class(g2m_gnn_type)( g2m_edge_index, hidden_dim, hidden_layers=hidden_layers, update_edges=False, ) + mesh_up_class = get_gnn_class(mesh_up_gnn_type) self.mesh_up_gnns = nn.ModuleList( [ - PropagationNet( + mesh_up_class( edge_index, hidden_dim, hidden_layers=hidden_layers, @@ -65,6 +69,8 @@ def __init__( utils.make_gnn_seq( edge_index, intra_level_layers, hidden_layers, hidden_dim ) + if intra_level_layers > 0 + else utils.IdentityModule() for edge_index in m2m_edge_index ] ) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 48c280ee..18284dd7 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -479,22 +479,37 @@ def forward(self, *args): return args -def make_gnn_seq(edge_index, num_gnn_layers, hidden_layers, hidden_dim): +def make_gnn_seq( + edge_index, + num_gnn_layers, + hidden_layers, + hidden_dim, + gnn_type="InteractionNet", +): """ - Build a sequential stack of InteractionNet layers that propagates both - node and edge representations. Returns an IdentityModule if - num_gnn_layers is 0. + Build a sequential stack of GNN layers that propagates both node and + edge representations. The layer type is set by ``gnn_type`` (any key in + ``gnn_layers.GNN_TYPES``, default ``InteractionNet``); all such layers + share the ``(send, rec, edge) -> (rec, edge)`` interface. + + ``num_gnn_layers`` must be at least 1. Callers that want a no-op stage + (e.g. zero intra-level layers) should substitute an ``IdentityModule`` + themselves rather than calling this with 0. """ # First-party - from neural_lam.gnn_layers import InteractionNet + from neural_lam.gnn_layers import get_gnn_class - if num_gnn_layers == 0: - return IdentityModule() + if num_gnn_layers < 1: + raise ValueError( + "make_gnn_seq requires num_gnn_layers >= 1 " + f"(got {num_gnn_layers}); use an IdentityModule for a no-op stage." + ) + gnn_class = get_gnn_class(gnn_type) return pyg.nn.Sequential( "mesh_rep, edge_rep", [ ( - InteractionNet( + gnn_class( edge_index, hidden_dim, hidden_layers=hidden_layers, diff --git a/tests/test_latent_modules.py b/tests/test_latent_modules.py index 3b845ed4..6194e052 100644 --- a/tests/test_latent_modules.py +++ b/tests/test_latent_modules.py @@ -53,7 +53,7 @@ def flat_dims(): "latent_dim": 4, "num_state_vars": 2, "hidden_layers": 1, - "processor_layers": 2, + "m2m_layers": 2, } @@ -87,18 +87,14 @@ def test_identity_module_passes_args_through(): assert out == (a, b, c) -def test_make_gnn_seq_zero_layers_returns_identity(): +def test_make_gnn_seq_zero_layers_raises(): + """make_gnn_seq must build a real sequence; the no-op (identity) case is + the caller's responsibility, exercised via the zero-intra-layer tests.""" edge_index = _fully_connected_edge_index(3, 3) - seq = make_gnn_seq( - edge_index, num_gnn_layers=0, hidden_layers=1, hidden_dim=8 - ) - assert isinstance(seq, IdentityModule) - - 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 torch.equal(out_mesh, mesh_rep) - assert torch.equal(out_edge, edge_rep) + 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(): @@ -164,6 +160,10 @@ def test_constant_encoder_is_input_independent(): assert torch.equal(a.mean, b.mean) assert torch.equal(a.stddev, b.stddev) assert a.mean.shape == (2, 3, 4) + # Prior is a mean-0 standard normal (isotropic): fixes the prob_model_lam + # mean-1 bug, see ConstantLatentEncoder docstring. + 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( @@ -174,7 +174,7 @@ def test_graph_encoder_shapes_and_backprop( g2m_edge_index=flat_edges["g2m"], m2m_edge_index=flat_edges["m2m"], hidden_dim=flat_dims["hidden_dim"], - processor_layers=flat_dims["processor_layers"], + m2m_layers=flat_dims["m2m_layers"], hidden_layers=flat_dims["hidden_layers"], output_dist="diagonal", ) @@ -204,7 +204,7 @@ def test_graph_decoder_shapes_with_output_std( hidden_dim=flat_dims["hidden_dim"], latent_dim=flat_dims["latent_dim"], num_state_vars=flat_dims["num_state_vars"], - processor_layers=flat_dims["processor_layers"], + m2m_layers=flat_dims["m2m_layers"], hidden_layers=flat_dims["hidden_layers"], output_std=True, ) @@ -241,7 +241,7 @@ def test_graph_decoder_no_output_std_returns_none( hidden_dim=flat_dims["hidden_dim"], latent_dim=flat_dims["latent_dim"], num_state_vars=flat_dims["num_state_vars"], - processor_layers=flat_dims["processor_layers"], + m2m_layers=flat_dims["m2m_layers"], hidden_layers=flat_dims["hidden_layers"], output_std=False, ) @@ -265,6 +265,56 @@ def test_graph_decoder_no_output_std_returns_none( assert pred_std is None +def test_flat_modules_zero_m2m_layers_use_identity( + flat_dims, flat_edges, flat_graph_emb +): + """m2m_layers=0 routes on-mesh processing through IdentityModule at the + call site (make_gnn_seq itself rejects 0). 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 isinstance(enc.m2m_gnns, IdentityModule) + + 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 isinstance(dec.m2m_gnns, IdentityModule) + + 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 ---------------------------------------- From df632b79602190b70ea9daf7b2607e668438274b Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 5 Jun 2026 17:10:11 +0530 Subject: [PATCH 04/51] style: apply black formatting to hierarchical latent modules --- neural_lam/models/latent/hi_graph_decoder.py | 26 ++++++++++++++------ neural_lam/models/latent/hi_graph_encoder.py | 13 +++++++--- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index 906ce7d6..2e11c432 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -94,21 +94,31 @@ def __init__( # Identity mappings if intra_level_layers == 0 self.intra_up_gnns = nn.ModuleList( [ - utils.make_gnn_seq( - edge_index, intra_level_layers, hidden_layers, hidden_dim + ( + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + if intra_level_layers > 0 + else utils.IdentityModule() ) - if intra_level_layers > 0 - else utils.IdentityModule() for edge_index in m2m_edge_index ] ) self.intra_down_gnns = nn.ModuleList( [ - utils.make_gnn_seq( - edge_index, intra_level_layers, hidden_layers, hidden_dim + ( + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + if intra_level_layers > 0 + else utils.IdentityModule() ) - if intra_level_layers > 0 - else utils.IdentityModule() for edge_index in list(m2m_edge_index)[:-1] # Top level (L) does not need a down intra-level GNN ] diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index 73677634..de8d87a2 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -66,11 +66,16 @@ def __init__( # Identity mappings if intra_level_layers == 0 self.intra_level_gnns = nn.ModuleList( [ - utils.make_gnn_seq( - edge_index, intra_level_layers, hidden_layers, hidden_dim + ( + utils.make_gnn_seq( + edge_index, + intra_level_layers, + hidden_layers, + hidden_dim, + ) + if intra_level_layers > 0 + else utils.IdentityModule() ) - if intra_level_layers > 0 - else utils.IdentityModule() for edge_index in m2m_edge_index ] ) From 1274b029220d6cd429f36db886d805e5617ae749 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 11 Jun 2026 14:11:15 +0530 Subject: [PATCH 05/51] feat: add GraphEFM single-step probabilistic predictor (PR-6) Port prob_model_lam's GraphEFM single-step half onto the StepPredictor interface, reusing the latent encoder/decoder infra. The predictor owns its conditional prior, variational encoder, and latent decoder, plus the per-step ELBO pieces (compute_step_loss) and sampling helpers; rollout, ELBO assembly, ensemble logic, and logging stay outside it. - forward is source's predict_step (prior rsample -> decode -> sampled next state); no rescaling/clamping - loss_fn and interior_mask are threaded parameters, not predictor state; compute_step_loss takes compute_kl (kl_term=None when off) - per_var_std mirrors ForecasterModule's formula, hence the config arg - one class for flat + hierarchical meshes, resolved from self.hierarchical - not registered in MODELS yet (needs config / no mesh_aggr); config-aware assembly deferred to the ensemble-forecaster PR Adds tests/test_graph_efm_predictor.py covering forward shapes, output_std, compute_step_loss + KL toggle, differentiability, member stochasticity, sample_obs_noise, and the per_var_std formula (flat + hierarchical). --- neural_lam/models/__init__.py | 8 + .../models/step_predictors/graph/graph_efm.py | 680 ++++++++++++++++++ tests/test_graph_efm_predictor.py | 255 +++++++ 3 files changed, 943 insertions(+) create mode 100644 neural_lam/models/step_predictors/graph/graph_efm.py create mode 100644 tests/test_graph_efm_predictor.py diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..aa589250 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -6,11 +6,19 @@ from .module import ForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel +from .step_predictors.graph.graph_efm import GraphEFM 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 +# NOTE: GraphEFM is intentionally NOT registered in MODELS yet. The shared +# construction call in train_model.py (e.g. line 34) instantiates the chosen +# model with a fixed deterministic kwarg set -- datastore-first, no ``config``, +# and with ``mesh_aggr`` -- whereas GraphEFM requires ``config`` (for its +# per_var_std weighting) and takes no ``mesh_aggr``. Registering it now would +# break that call. Wiring up config-aware assembly is deferred to the +# ensemble-forecaster PR (see open question Q3). MODELS = { "graph_lam": GraphLAM, "hi_lam": HiLAM, 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..6eb440bd --- /dev/null +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -0,0 +1,680 @@ +# 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 GraphEFM(StepPredictor): + """ + Graph-based Ensemble Forecasting Model -- single-step predictor. + + Port of ``prob_model_lam``'s ``GraphEFM`` (``forward`` is the source's + ``predict_step``) onto the ``StepPredictor`` interface. The predictor owns + its own conditional-prior / variational-encoder / latent-decoder, each of + which carries its own g2m/processor/m2g GNNs, so the encode-process-decode + backbone of ``BaseGraphModel`` does not apply -- this extends + ``StepPredictor`` directly. It is self-contained: besides ``forward`` it + exposes the per-step ELBO pieces (``compute_step_loss`` -> + ``(likelihood_term, kl_term, pred_mean, pred_std)``) and the sampling + helpers used by a future rollout/ensemble module. Rollout, ELBO assembly, + ensemble logic and logging live outside the predictor. + + One class handles both flat and hierarchical meshes, resolved at + construction from ``self.hierarchical`` (set by ``utils.load_graph``). + """ + + 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_processor_layers: int = 2, + encoder_processor_layers: int = 2, + processor_layers: int = 4, + learn_prior: bool = True, + prior_dist: str = "isotropic", + num_past_forcing_steps: int = 1, + num_future_forcing_steps: int = 1, + output_std: bool = False, + sample_obs_noise: bool = False, + output_clamping_lower: Optional[Dict[str, float]] = None, + output_clamping_upper: Optional[Dict[str, float]] = None, + ): + super().__init__( + datastore=datastore, + output_std=output_std, + output_clamping_lower=output_clamping_lower, + output_clamping_upper=output_clamping_upper, + ) + + # Whether to sample observation noise during rollout. When False, + # sample_next_state returns the predicted mean. + self.sample_obs_noise = bool(sample_obs_noise) + + # Load graph with static features (same pattern as BaseGraphModel). + # 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 + ) + 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 (datastore-driven; replaces source's + # constants.GRID_STATE_DIM / GRID_FORCING_DIM). + 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] + # grid_dim: total grid input dim, same formula as BaseGraphModel. The + # cat ORDER in embedd_all/embedd_current follows source + # (prev_prev, prev, forcing, static[, current]); the size is unchanged. + self.grid_dim = ( + 2 * num_state_vars + + grid_static_dim + + num_forcing_vars + * (num_past_forcing_steps + num_future_forcing_steps + 1) + ) + 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) + + if self.hierarchical: + level_mesh_sizes = [ + mesh_feat.shape[0] for mesh_feat in self.mesh_static_features + ] + 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 processor layers, no need to embed m2m + self.embedd_m2m = ( + max( + prior_processor_layers, + encoder_processor_layers, + processor_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) + ] + ) + else: + 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)" + ) + 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: + if self.hierarchical: + 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_processor_layers, + hidden_layers=hidden_layers, + output_dist=prior_dist, + ) + else: + 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_processor_layers, + hidden_layers=hidden_layers, + 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. The latent modules take + # num_state_vars (datastore-driven) where source used GRID_STATE_DIM. + if self.hierarchical: + 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_processor_layers, + hidden_layers=hidden_layers, + 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=num_state_vars, + intra_level_layers=processor_layers, + hidden_layers=hidden_layers, + output_std=bool(output_std), + ) + else: + 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_processor_layers, + hidden_layers=hidden_layers, + 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=num_state_vars, + m2m_layers=processor_layers, + hidden_layers=hidden_layers, + output_std=bool(output_std), + ) + + # 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 sample_next_state(self, pred_mean, pred_std): + """ + Sample state at next time step given a Gaussian observation model. + If ``self.sample_obs_noise`` is False, only return the mean. + + Parameters + ---------- + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Predicted mean. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, d_state)``, or None when the decoder + does not output a std (``output_std=False``). + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Next state. + """ + if not self.output_std: + pred_std = self.per_var_std # (d_f,) + + if self.sample_obs_noise: + return torch.distributions.Normal(pred_mean, pred_std).rsample() + # (B, num_grid_nodes, d_state) + + return pred_mean # (B, num_grid_nodes, d_state) + + def embedd_current( + 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_all(self, prev_state, prev_prev_state, forcing): + """ + Embed all node and edge representations. + + 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) + } + + if self.hierarchical: + graph_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) + + if self.embedd_m2m: + graph_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 + graph_emb["m2m"] = list(self.m2m_features) + + graph_emb["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 + ) + ] + graph_emb["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 + ) + ] + else: + graph_emb["mesh"] = self.expand_to_batch( + self.mesh_embedder(self.mesh_static_features), batch_size + ) # (B, num_mesh_nodes, d_h) + graph_emb["m2m"] = self.expand_to_batch( + self.m2m_embedder(self.m2m_features), batch_size + ) # (B, M_m2m, d_h) + + 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_all``. + graph_emb : dict + Edge/mesh embeddings from ``embedd_all``. + 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_all( + prev_states[:, 1], + prev_states[:, 0], + forcing_features, + ) + # embed also including current grid state, for encoder + grid_current_emb = self.embedd_current( + 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 (source's ``predict_step``): + embed features, sample the latent from the prior, decode, and return + the sampled next state. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + + Returns + ------- + new_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Sampled ``X_{t+1}``. + 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_all( + 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) + last_state = prev_state + pred_mean, pred_std = self.decoder( + grid_prev_emb, latent_samples, last_state, graph_emb + ) # (B, num_grid_nodes, d_state) + + return self.sample_next_state(pred_mean, pred_std), pred_std diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py new file mode 100644 index 00000000..c2aab4e3 --- /dev/null +++ b/tests/test_graph_efm_predictor.py @@ -0,0 +1,255 @@ +"""Unit tests for the GraphEFM single-step probabilistic predictor. + +These mirror the smoke-test pattern used for the deterministic predictors +(see ``tests/test_gnn_layers.py``): build flat and hierarchical 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 +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, sample_obs_noise=False): + datastore, config = _datastore_and_config_with_graph(graph_name) + predictor = GraphEFM( + config=config, + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + prior_processor_layers=1, + encoder_processor_layers=1, + processor_layers=1, + 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, + sample_obs_noise=sample_obs_noise, + ) + 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_sample_next_state_respects_sample_obs_noise(): + """sample_next_state returns the mean when sample_obs_noise is False and a + stochastic draw (different from the mean) when True.""" + deterministic, datastore, _ = _build_predictor( + "1level", sample_obs_noise=False + ) + d_state = datastore.get_num_data_vars(category="state") + # Last dim must match per_var_std (d_state,) for the obs-noise broadcast. + pred_mean = torch.randn(2, 5, d_state) + + out_mean = deterministic.sample_next_state(pred_mean, pred_std=None) + assert torch.equal(out_mean, pred_mean) + + stochastic, _, _ = _build_predictor("1level", sample_obs_noise=True) + # per_var_std is registered (output_std=False); the draw should differ + # from the mean. + out_sampled = stochastic.sample_next_state(pred_mean, pred_std=None) + assert out_sampled.shape == pred_mean.shape + assert not torch.allclose(out_sampled, pred_mean) + + +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 From ea9ab9c06f9c51dc1506c7230c6a0625c559ab03 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 12 Jun 2026 23:00:49 +0530 Subject: [PATCH 06/51] refactor: hard-code constrained latent GNN types, reuse existing flags for the rest Per review discussion: the architecturally constrained edge sets in the hierarchical latent modules get fixed GNN types instead of parameters: - HiGraphLatentEncoder mesh-up: PropagationNet (must push grid info up into the latent readout) - HiGraphLatentDecoder mesh-up: InteractionNet (PropagationNet residual would bypass Z at the top level, leaving it unused at initialization) - HiGraphLatentDecoder mesh-down: PropagationNet (must push Z down the hierarchy to reach the grid output) All remaining choices (g2m/m2g) stay configurable and default to InteractionNet for consistency with the rest of the codebase. GraphEFM now accepts g2m_gnn_type/m2g_gnn_type and passes them through to the prior, encoder and decoder, ready for wiring to the existing argparse flags. --- neural_lam/models/latent/graph_decoder.py | 2 +- neural_lam/models/latent/graph_encoder.py | 2 +- neural_lam/models/latent/hi_graph_decoder.py | 28 +++++++++++-------- neural_lam/models/latent/hi_graph_encoder.py | 18 ++++++------ .../models/step_predictors/graph/graph_efm.py | 10 +++++++ 5 files changed, 39 insertions(+), 21 deletions(-) diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index b613878b..b57ee580 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -26,7 +26,7 @@ def __init__( m2m_layers, hidden_layers=1, g2m_gnn_type="InteractionNet", - m2g_gnn_type="PropagationNet", + m2g_gnn_type="InteractionNet", output_std=True, ): super().__init__( diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py index 1dd824a3..8b4a28ae 100644 --- a/neural_lam/models/latent/graph_encoder.py +++ b/neural_lam/models/latent/graph_encoder.py @@ -22,7 +22,7 @@ def __init__( hidden_dim, m2m_layers, hidden_layers=1, - g2m_gnn_type="PropagationNet", + g2m_gnn_type="InteractionNet", output_dist="isotropic", ): super().__init__(latent_dim, output_dist) diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index 2e11c432..54cbba30 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -3,7 +3,11 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import get_gnn_class +from neural_lam.gnn_layers import ( + InteractionNet, + PropagationNet, + get_gnn_class, +) # Local from .base_decoder import BaseGraphLatentDecoder @@ -16,9 +20,9 @@ class HiGraphLatentDecoder(BaseGraphLatentDecoder): *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, mesh-up and - mesh-down GNN types are set by ``g2m_gnn_type``, ``mesh_up_gnn_type`` - and ``mesh_down_gnn_type`` respectively. + 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__( @@ -34,9 +38,7 @@ def __init__( intra_level_layers, hidden_layers=1, g2m_gnn_type="InteractionNet", - m2g_gnn_type="PropagationNet", - mesh_up_gnn_type="InteractionNet", - mesh_down_gnn_type="PropagationNet", + m2g_gnn_type="InteractionNet", output_std=True, ): super().__init__( @@ -66,10 +68,12 @@ def __init__( update_edges=False, ) - mesh_up_class = get_gnn_class(mesh_up_gnn_type) + # 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( [ - mesh_up_class( + InteractionNet( edge_index, hidden_dim, hidden_layers=hidden_layers, @@ -78,10 +82,12 @@ def __init__( for edge_index in mesh_up_edge_index ] ) - mesh_down_class = get_gnn_class(mesh_down_gnn_type) + # 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( [ - mesh_down_class( + PropagationNet( edge_index, hidden_dim, hidden_layers=hidden_layers, diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index de8d87a2..bc4f8545 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -3,7 +3,7 @@ # First-party from neural_lam import utils -from neural_lam.gnn_layers import get_gnn_class +from neural_lam.gnn_layers import PropagationNet, get_gnn_class # Local from .base_encoder import BaseLatentEncoder @@ -13,9 +13,9 @@ 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 GNNs (type set by ``mesh_up_gnn_type``), with - optional intra-level processing at each level. The latent distribution - is read out from the top mesh level. + 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__( @@ -27,8 +27,7 @@ def __init__( hidden_dim, intra_level_layers, hidden_layers=1, - g2m_gnn_type="PropagationNet", - mesh_up_gnn_type="PropagationNet", + g2m_gnn_type="InteractionNet", output_dist="isotropic", ): super().__init__(latent_dim, output_dist) @@ -50,10 +49,13 @@ def __init__( update_edges=False, ) - mesh_up_class = get_gnn_class(mesh_up_gnn_type) + # 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( [ - mesh_up_class( + PropagationNet( edge_index, hidden_dim, hidden_layers=hidden_layers, diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 6eb440bd..58245854 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -54,6 +54,8 @@ def __init__( 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, sample_obs_noise: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, @@ -210,6 +212,7 @@ def __init__( hidden_dim=hidden_dim, intra_level_layers=prior_processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, output_dist=prior_dist, ) else: @@ -220,6 +223,7 @@ def __init__( hidden_dim=hidden_dim, m2m_layers=prior_processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, output_dist=prior_dist, ) else: @@ -240,6 +244,7 @@ def __init__( hidden_dim=hidden_dim, intra_level_layers=encoder_processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, output_dist="diagonal", ) self.decoder = HiGraphLatentDecoder( @@ -253,6 +258,8 @@ def __init__( num_state_vars=num_state_vars, intra_level_layers=processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + m2g_gnn_type=m2g_gnn_type, output_std=bool(output_std), ) else: @@ -263,6 +270,7 @@ def __init__( hidden_dim=hidden_dim, m2m_layers=encoder_processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, output_dist="diagonal", ) self.decoder = GraphLatentDecoder( @@ -274,6 +282,8 @@ def __init__( num_state_vars=num_state_vars, m2m_layers=processor_layers, hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + m2g_gnn_type=m2g_gnn_type, output_std=bool(output_std), ) From cd3d4f0dc6f432809aca2d2dfb2dc66c84786179 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 12 Jun 2026 23:17:57 +0530 Subject: [PATCH 07/51] fix: satisfy interrogate docstring coverage from upstream main Upstream main added an interrogate pre-commit hook requiring 100% docstring coverage, which failed on this branch's CI after merging. - Remove the branch's pre-reorganization duplicates (forecaster.py, ar_forecaster.py, step_predictor.py, forecaster_module.py); main carries the same code under models/forecasters/, models/ step_predictors/ and models/module.py, and all imports already go through the new layout. - Add the missing module and __init__ docstrings (numpy style) in the latent modules, GraphEFM and utils.IdentityModule.forward. --- neural_lam/models/latent/__init__.py | 4 ++ neural_lam/models/latent/base_decoder.py | 19 ++++++ neural_lam/models/latent/base_encoder.py | 13 ++++ neural_lam/models/latent/constant_encoder.py | 15 +++++ neural_lam/models/latent/graph_decoder.py | 34 +++++++++++ neural_lam/models/latent/graph_encoder.py | 26 ++++++++ neural_lam/models/latent/hi_graph_decoder.py | 43 +++++++++++++ neural_lam/models/latent/hi_graph_encoder.py | 31 ++++++++++ .../models/step_predictors/graph/graph_efm.py | 60 +++++++++++++++++++ neural_lam/utils.py | 13 ++++ 10 files changed, 258 insertions(+) diff --git a/neural_lam/models/latent/__init__.py b/neural_lam/models/latent/__init__.py index f50d2ac6..fba7ed42 100644 --- a/neural_lam/models/latent/__init__.py +++ b/neural_lam/models/latent/__init__.py @@ -1,3 +1,7 @@ +"""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 diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index 5f72beca..022da927 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -1,3 +1,5 @@ +"""Abstract base class for graph-based latent decoders.""" + # Third-party from torch import nn @@ -25,6 +27,23 @@ def __init__( 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_dim : int + Dimensionality of the latent variable at each mesh node. + num_state_vars : int + Number of state variables predicted at each grid node. + hidden_layers : int + Number of hidden layers in the internal MLPs. + output_std : bool + If True, the decoder outputs both mean and std of the next-state + distribution; if False, only the mean. + """ super().__init__() self.grid_update_mlp = utils.make_mlp( diff --git a/neural_lam/models/latent/base_encoder.py b/neural_lam/models/latent/base_encoder.py index 889014b0..6cf99b4c 100644 --- a/neural_lam/models/latent/base_encoder.py +++ b/neural_lam/models/latent/base_encoder.py @@ -1,3 +1,5 @@ +"""Abstract base class for latent encoders.""" + # Third-party import torch from torch import distributions as tdists @@ -16,6 +18,17 @@ class BaseLatentEncoder(nn.Module): """ 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 diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 600f0a09..8abc4cdd 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -1,3 +1,6 @@ +"""Constant (input-independent) latent encoder, used as a non-learned +prior.""" + # Third-party import torch @@ -21,6 +24,18 @@ class ConstantLatentEncoder(BaseLatentEncoder): """ 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 diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index b57ee580..ae116b95 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -1,3 +1,5 @@ +"""Latent decoder for flat (non-hierarchical) graphs.""" + # First-party from neural_lam import utils from neural_lam.gnn_layers import get_gnn_class @@ -29,6 +31,38 @@ def __init__( 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 ) diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py index 8b4a28ae..021ce7d0 100644 --- a/neural_lam/models/latent/graph_encoder.py +++ b/neural_lam/models/latent/graph_encoder.py @@ -1,3 +1,5 @@ +"""Latent encoder for flat (non-hierarchical) graphs.""" + # First-party from neural_lam import utils from neural_lam.gnn_layers import get_gnn_class @@ -25,6 +27,30 @@ def __init__( 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)( diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index 54cbba30..cef56b19 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -1,3 +1,5 @@ +"""Latent decoder for hierarchical graphs.""" + # Third-party from torch import nn @@ -41,6 +43,47 @@ def __init__( 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 ) diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index bc4f8545..142161c0 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -1,3 +1,5 @@ +"""Latent encoder for hierarchical graphs.""" + # Third-party from torch import nn @@ -30,6 +32,35 @@ def __init__( 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 diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 58245854..470bda0f 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -1,3 +1,6 @@ +"""Graph-based Ensemble Forecasting Model (GraphEFM) single-step +predictor.""" + # Standard library from typing import Callable, Dict, Optional @@ -61,6 +64,63 @@ def __init__( output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None, ): + """ + Build the prior, variational encoder and latent decoder sub-models. + + 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. + Both flat and hierarchical graphs are supported; which latent + modules are built is resolved from the loaded 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_processor_layers : int + Number of processor GNN layers in the (learned) prior. + encoder_processor_layers : int + Number of processor GNN layers in the variational encoder. + processor_layers : int + Number of processor 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. + sample_obs_noise : bool + If True, sample observation noise when rolling out; if False, + ``sample_next_state`` returns the predicted mean. + 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, diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 18284dd7..ea384f5b 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -476,6 +476,19 @@ class IdentityModule(nn.Module): """Identity operator that accepts and returns multiple positional inputs.""" def forward(self, *args): + """ + Return all positional inputs unchanged. + + Parameters + ---------- + *args : tuple + Any positional arguments. + + Returns + ------- + tuple + The inputs, unchanged. + """ return args From 8f41e5ef265fcc84e74f4877a61a5b05368c78d5 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 12 Jun 2026 23:29:55 +0530 Subject: [PATCH 08/51] docs: describe current code only in docstrings and comments Remove references to the original prob_model_lam implementation and other work meta-information from docstrings and comments, per review. Docstrings now describe what each class/function does; usage context is left to call sites. --- neural_lam/models/__init__.py | 11 +++--- neural_lam/models/latent/constant_encoder.py | 15 +++----- .../models/step_predictors/graph/graph_efm.py | 36 +++++++++---------- tests/test_latent_modules.py | 3 +- 4 files changed, 26 insertions(+), 39 deletions(-) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index aa589250..66ae3ec3 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -13,12 +13,11 @@ from .step_predictors.graph.hierarchical import BaseHiGraphModel # NOTE: GraphEFM is intentionally NOT registered in MODELS yet. The shared -# construction call in train_model.py (e.g. line 34) instantiates the chosen -# model with a fixed deterministic kwarg set -- datastore-first, no ``config``, -# and with ``mesh_aggr`` -- whereas GraphEFM requires ``config`` (for its -# per_var_std weighting) and takes no ``mesh_aggr``. Registering it now would -# break that call. Wiring up config-aware assembly is deferred to the -# ensemble-forecaster PR (see open question Q3). +# 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 GraphEFM requires ``config`` (for its per_var_std +# weighting) and takes no ``mesh_aggr``. Registering it requires config-aware +# model assembly in train_model.py. MODELS = { "graph_lam": GraphLAM, "hi_lam": HiLAM, diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 8abc4cdd..60cfe109 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -1,5 +1,4 @@ -"""Constant (input-independent) latent encoder, used as a non-learned -prior.""" +"""Constant (input-independent) latent encoder.""" # Third-party import torch @@ -12,15 +11,9 @@ class ConstantLatentEncoder(BaseLatentEncoder): """ Latent encoder that returns a constant (input-independent) distribution. - Used as a non-learned prior in ``GraphEFM`` when ``learn_prior`` is - disabled. ``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"``. (Note: ``prob_model_lam`` returned a tensor - of ones here, giving mean 1, while its ``train_model.py`` CLI help - described the prior as "mean 0". The mean 1 was a bug -- it is only a - constant offset, but a mean-0 prior is what is intended, so the port - uses zeros.) + ``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"): diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 470bda0f..0ac7de5a 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -27,16 +27,15 @@ class GraphEFM(StepPredictor): """ Graph-based Ensemble Forecasting Model -- single-step predictor. - Port of ``prob_model_lam``'s ``GraphEFM`` (``forward`` is the source's - ``predict_step``) onto the ``StepPredictor`` interface. The predictor owns - its own conditional-prior / variational-encoder / latent-decoder, each of - which carries its own g2m/processor/m2g GNNs, so the encode-process-decode - backbone of ``BaseGraphModel`` does not apply -- this extends - ``StepPredictor`` directly. It is self-contained: besides ``forward`` it - exposes the per-step ELBO pieces (``compute_step_loss`` -> - ``(likelihood_term, kl_term, pred_mean, pred_std)``) and the sampling - helpers used by a future rollout/ensemble module. Rollout, ELBO assembly, - ensemble logic and logging live outside the predictor. + A latent-variable step predictor consisting of a conditional prior, a + variational encoder and a latent decoder, each of which carries its own + g2m/processor/m2g 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)``) and sampling helpers. Rollout, ELBO assembly, ensemble + logic and logging live outside the predictor. One class handles both flat and hierarchical meshes, resolved at construction from ``self.hierarchical`` (set by ``utils.load_graph``). @@ -146,14 +145,13 @@ def __init__( else: setattr(self, name, attr_value) - # Specify dimensions of data (datastore-driven; replaces source's - # constants.GRID_STATE_DIM / GRID_FORCING_DIM). + # Specify dimensions of data 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] - # grid_dim: total grid input dim, same formula as BaseGraphModel. The - # cat ORDER in embedd_all/embedd_current follows source - # (prev_prev, prev, forcing, static[, current]); the size is unchanged. + # grid_dim: total grid input dim, same formula as BaseGraphModel, + # matching the cat order in embedd_all/embedd_current + # (prev_prev, prev, forcing, static[, current]). self.grid_dim = ( 2 * num_state_vars + grid_static_dim @@ -293,8 +291,7 @@ def __init__( output_dist=prior_dist, ) - # Encoder (variational posterior) + Decoder. The latent modules take - # num_state_vars (datastore-driven) where source used GRID_STATE_DIM. + # Encoder (variational posterior) + Decoder if self.hierarchical: self.encoder = HiGraphLatentEncoder( latent_dim=latent_dim, @@ -706,9 +703,8 @@ def forward( forcing: torch.Tensor, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """ - Sample one time step prediction (source's ``predict_step``): - embed features, sample the latent from the prior, decode, and return - the sampled next state. + Sample one time step prediction: embed features, sample the latent + from the prior, decode, and return the sampled next state. Parameters ---------- diff --git a/tests/test_latent_modules.py b/tests/test_latent_modules.py index 6194e052..e4f06803 100644 --- a/tests/test_latent_modules.py +++ b/tests/test_latent_modules.py @@ -160,8 +160,7 @@ def test_constant_encoder_is_input_independent(): assert torch.equal(a.mean, b.mean) assert torch.equal(a.stddev, b.stddev) assert a.mean.shape == (2, 3, 4) - # Prior is a mean-0 standard normal (isotropic): fixes the prob_model_lam - # mean-1 bug, see ConstantLatentEncoder docstring. + # 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)) From 0e50d9ce5096c11b1b83b48dfdb1cc919d974250 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 12 Jun 2026 23:36:07 +0530 Subject: [PATCH 09/51] docs: convert latent module and make_gnn_seq docstrings to numpy style Add proper Parameters/Returns sections following the numpydoc convention, per review. --- neural_lam/models/latent/base_decoder.py | 56 ++++++++++++++------ neural_lam/models/latent/base_encoder.py | 30 ++++++++--- neural_lam/models/latent/constant_encoder.py | 15 +++++- neural_lam/models/latent/graph_decoder.py | 22 ++++++-- neural_lam/models/latent/graph_encoder.py | 22 +++++--- neural_lam/models/latent/hi_graph_decoder.py | 35 +++++++----- neural_lam/models/latent/hi_graph_encoder.py | 27 ++++++---- neural_lam/utils.py | 37 ++++++++++--- 8 files changed, 182 insertions(+), 62 deletions(-) diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index 022da927..94deecce 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -70,13 +70,23 @@ def combine_with_latent( """ Fuse grid and latent representations and return a grid-shaped output. - original_grid_rep: (B, num_grid_nodes, d_h) - latent_rep: (B, num_mesh_nodes, d_h) - residual_grid_rep: (B, num_grid_nodes, d_h) - graph_emb: dict of graph edge / node embeddings - - Returns: - combined_grid_rep: (B, num_grid_nodes, d_h) + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes, d_h)``. Embedded latent sample. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation to use + for residual connections. + graph_emb : dict + Embedded graph node and edge features. + + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation. """ raise NotImplementedError("combine_with_latent not implemented") @@ -84,14 +94,30 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): """ Predict mean (and optionally std) of the next weather state. - grid_rep: (B, num_grid_nodes, d_h) - latent_samples: (B, num_mesh_nodes, latent_dim) - last_state: (B, num_grid_nodes, num_state_vars) - graph_emb: dict with at least ``g2m``, ``m2m``, ``m2g`` entries - - Returns: - pred_mean: (B, num_grid_nodes, num_state_vars) - pred_std: (B, num_grid_nodes, num_state_vars) or ``None`` + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation. + latent_samples : torch.Tensor + Shape ``(B, num_mesh_nodes, latent_dim)``. Sample of the + latent variable. + 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 + prediction. + graph_emb : dict + Embedded graph node and edge features, with at least ``g2m``, + ``m2m`` and ``m2g`` entries. + + 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. """ latent_emb = self.latent_embedder(latent_samples) diff --git a/neural_lam/models/latent/base_encoder.py b/neural_lam/models/latent/base_encoder.py index 6cf99b4c..cad46d4f 100644 --- a/neural_lam/models/latent/base_encoder.py +++ b/neural_lam/models/latent/base_encoder.py @@ -47,10 +47,19 @@ def compute_dist_params(self, grid_rep, **kwargs): """ Compute raw distribution parameters from the grid representation. - grid_rep: (B, num_grid_nodes, d_h) + 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: - parameters: (B, num_mesh_nodes, output_dim) + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of + the latent distribution. """ raise NotImplementedError("compute_dist_params not implemented") @@ -58,11 +67,18 @@ def forward(self, grid_rep, **kwargs): """ Compute the Gaussian distribution over the latent variable. - grid_rep: (B, num_grid_nodes, d_h) + 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: - distribution: ``torch.distributions.Normal`` of shape - (B, num_mesh_nodes, latent_dim) + 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) diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 60cfe109..35f31fda 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -34,7 +34,20 @@ def __init__(self, latent_dim, num_mesh_nodes, output_dist="isotropic"): def compute_dist_params(self, grid_rep, **kwargs): """ - Return constant parameters of shape (B, num_mesh_nodes, output_dim). + Return constant (zero) distribution parameters. + + Parameters + ---------- + grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Used only to determine + batch size and device; the values do not affect the output. + **kwargs + Ignored. + + Returns + ------- + torch.Tensor + Shape ``(B, num_mesh_nodes, output_dim)``. All zeros. """ return torch.zeros( grid_rep.shape[0], diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index ae116b95..cf9df1d5 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -95,12 +95,24 @@ def combine_with_latent( """ Fuse grid and latent reps via g2m -> m2m -> m2g. - original_grid_rep: (B, num_grid_nodes, d_h) - latent_rep: (B, num_mesh_nodes, d_h) - residual_grid_rep: (B, num_grid_nodes, d_h) + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes, d_h)``. Embedded latent sample. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation used + as receiver in the mesh-to-grid step. + graph_emb : dict + Embedded graph node and edge features, with at least ``g2m``, + ``m2m`` and ``m2g`` entries. - Returns: - grid_rep: (B, num_grid_nodes, d_h) + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation. """ mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py index 021ce7d0..e596ac61 100644 --- a/neural_lam/models/latent/graph_encoder.py +++ b/neural_lam/models/latent/graph_encoder.py @@ -78,14 +78,22 @@ def compute_dist_params(self, grid_rep, graph_emb, **kwargs): """ Compute distribution parameters on mesh from grid features. - grid_rep: (B, num_grid_nodes, d_h) - graph_emb: dict with at least - - ``mesh``: (B, num_mesh_nodes, d_h) - - ``g2m``: (B, M_g2m, d_h) - - ``m2m``: (B, M_m2m, d_h) + 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: - parameters: (B, num_mesh_nodes, output_dim) + 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"]) mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index cef56b19..a2640c44 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -179,19 +179,30 @@ def combine_with_latent( """ Hierarchical up-then-down fusion of grid and latent reps. - original_grid_rep: (B, num_grid_nodes, d_h) - latent_rep: (B, num_mesh_nodes[L], d_h) - residual_grid_rep: (B, num_grid_nodes, d_h) - graph_emb: dict with at least - - ``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) - - ``m2g``: (B, M_m2g, d_h) + Parameters + ---------- + original_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + latent_rep : torch.Tensor + Shape ``(B, num_mesh_nodes[L], d_h)``. Embedded latent sample + on the top mesh level ``L``. + residual_grid_rep : torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Grid representation used + as receiver in the mesh-to-grid step. + 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)``, + ``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)``. - Returns: - grid_rep: (B, num_grid_nodes, d_h) + Returns + ------- + torch.Tensor + Shape ``(B, num_grid_nodes, d_h)``. Combined grid + representation. """ current_mesh_rep = self.g2m_gnn( original_grid_rep, graph_emb["mesh"][0], graph_emb["g2m"] diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index 142161c0..2b1230ca 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -123,15 +123,24 @@ def compute_dist_params(self, grid_rep, graph_emb, **kwargs): """ Compute distribution parameters on the top mesh level. - grid_rep: (B, num_grid_nodes, d_h) - graph_emb: dict with at least - - ``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) - - Returns: - parameters: (B, num_mesh_nodes[L], output_dim) + 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"] diff --git a/neural_lam/utils.py b/neural_lam/utils.py index ea384f5b..52ebec8a 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -501,13 +501,38 @@ def make_gnn_seq( ): """ Build a sequential stack of GNN layers that propagates both node and - edge representations. The layer type is set by ``gnn_type`` (any key in - ``gnn_layers.GNN_TYPES``, default ``InteractionNet``); all such layers - share the ``(send, rec, edge) -> (rec, edge)`` interface. + edge representations. - ``num_gnn_layers`` must be at least 1. Callers that want a no-op stage - (e.g. zero intra-level layers) should substitute an ``IdentityModule`` - themselves rather than calling this with 0. + 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 + substitute an ``IdentityModule`` themselves 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 From 706fbe126497fcbbf4a52357103db8e76d8b4fba Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 12 Jun 2026 23:57:53 +0530 Subject: [PATCH 10/51] refactor: skip on-mesh/intra-level processing explicitly instead of IdentityModule When m2m_layers / intra_level_layers is 0, the latent modules now set the corresponding GNN attribute to None and skip the update in the forward pass, instead of routing representations through a no-op IdentityModule. This makes it clear from the forward code that no processing happens in that case. IdentityModule is removed from utils. The hierarchical up/down loops index levels explicitly to accommodate the conditional; outputs are unchanged (verified bit-identical against the previous implementation). --- neural_lam/models/latent/graph_decoder.py | 7 +- neural_lam/models/latent/graph_encoder.py | 7 +- neural_lam/models/latent/hi_graph_decoder.py | 97 +++++++++----------- neural_lam/models/latent/hi_graph_encoder.py | 52 +++++------ neural_lam/utils.py | 27 +----- tests/test_latent_modules.py | 27 ++---- 6 files changed, 94 insertions(+), 123 deletions(-) diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index cf9df1d5..0ac4b180 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -74,12 +74,14 @@ def __init__( 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 utils.IdentityModule() + else None ) self.m2g_gnn = get_gnn_class(m2g_gnn_type)( @@ -116,7 +118,8 @@ def combine_with_latent( """ mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) - mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) + 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"]) diff --git a/neural_lam/models/latent/graph_encoder.py b/neural_lam/models/latent/graph_encoder.py index e596ac61..9082f999 100644 --- a/neural_lam/models/latent/graph_encoder.py +++ b/neural_lam/models/latent/graph_encoder.py @@ -60,12 +60,14 @@ def __init__( 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 utils.IdentityModule() + else None ) self.latent_param_map = utils.make_mlp( @@ -96,5 +98,6 @@ def compute_dist_params(self, grid_rep, graph_emb, **kwargs): the latent distribution. """ mesh_rep = self.g2m_gnn(grid_rep, graph_emb["mesh"], graph_emb["g2m"]) - mesh_rep, _ = self.m2m_gnns(mesh_rep, graph_emb["m2m"]) + 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 index a2640c44..eb972705 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -140,37 +140,38 @@ def __init__( ] ) - # Identity mappings if intra_level_layers == 0 - self.intra_up_gnns = nn.ModuleList( - [ - ( + # 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, ) - if intra_level_layers > 0 - else utils.IdentityModule() - ) - for edge_index in m2m_edge_index - ] + for edge_index in m2m_edge_index + ] + ) + if intra_level_layers > 0 + else None ) - self.intra_down_gnns = nn.ModuleList( - [ - ( + self.intra_down_gnns = ( + nn.ModuleList( + [ utils.make_gnn_seq( edge_index, intra_level_layers, hidden_layers, hidden_dim, ) - if intra_level_layers > 0 - else utils.IdentityModule() - ) - for edge_index in list(m2m_edge_index)[:-1] - # Top level (L) does not need a down intra-level GNN - ] + 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( @@ -213,23 +214,21 @@ def combine_with_latent( # so the latent is fused in at the top of the hierarchy. mesh_level_reps = [] m2m_level_reps = [] - for ( - up_gnn, - intra_gnn_seq, - mesh_up_level_rep, - m2m_level_rep, - mesh_level_rep, - ) in zip( - self.mesh_up_gnns, - self.intra_up_gnns[:-1], - graph_emb["mesh_up"], - graph_emb["m2m"][:-1], - graph_emb["mesh"][1:-1] + [latent_rep], - ): - new_mesh_rep, new_m2m_rep = intra_gnn_seq( - current_mesh_rep, m2m_level_rep + 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) @@ -238,29 +237,23 @@ def combine_with_latent( ) # Top level processing - current_mesh_rep, _ = self.intra_up_gnns[-1]( - current_mesh_rep, graph_emb["m2m"][-1] - ) + 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 ( - down_gnn, - intra_gnn_seq, - mesh_down_level_rep, - m2m_level_rep, - mesh_level_rep, - ) in zip( - reversed(self.mesh_down_gnns), - reversed(self.intra_down_gnns), - reversed(graph_emb["mesh_down"]), - reversed(m2m_level_reps), - reversed(mesh_level_reps), - ): - new_mesh_rep = down_gnn( - current_mesh_rep, mesh_level_rep, mesh_down_level_rep + 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], ) - current_mesh_rep, _ = intra_gnn_seq(new_mesh_rep, m2m_level_rep) + 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"] diff --git a/neural_lam/models/latent/hi_graph_encoder.py b/neural_lam/models/latent/hi_graph_encoder.py index 2b1230ca..e0148b38 100644 --- a/neural_lam/models/latent/hi_graph_encoder.py +++ b/neural_lam/models/latent/hi_graph_encoder.py @@ -96,21 +96,22 @@ def __init__( ] ) - # Identity mappings if intra_level_layers == 0 - self.intra_level_gnns = nn.ModuleList( - [ - ( + # 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, ) - if intra_level_layers > 0 - else utils.IdentityModule() - ) - for edge_index in m2m_edge_index - ] + for edge_index in m2m_edge_index + ] + ) + if intra_level_layers > 0 + else None ) self.latent_param_map = utils.make_mlp( @@ -147,27 +148,26 @@ def compute_dist_params(self, grid_rep, graph_emb, **kwargs): ) # Same-level processing on level 0 - current_mesh_rep, _ = self.intra_level_gnns[0]( - current_mesh_rep, graph_emb["m2m"][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 ( - up_gnn, - intra_gnn_seq, - mesh_up_level_rep, - m2m_level_rep, - mesh_level_rep, - ) in zip( - self.mesh_up_gnns, - self.intra_level_gnns[1:], - graph_emb["mesh_up"], - graph_emb["m2m"][1:], - graph_emb["mesh"][1:], + 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, ): - new_node_rep = up_gnn( + current_mesh_rep = up_gnn( current_mesh_rep, mesh_level_rep, mesh_up_level_rep ) - current_mesh_rep, _ = intra_gnn_seq(new_node_rep, m2m_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/utils.py b/neural_lam/utils.py index 52ebec8a..eba72d1a 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -472,26 +472,6 @@ def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: return nn.Sequential(*layers) -class IdentityModule(nn.Module): - """Identity operator that accepts and returns multiple positional inputs.""" - - def forward(self, *args): - """ - Return all positional inputs unchanged. - - Parameters - ---------- - *args : tuple - Any positional arguments. - - Returns - ------- - tuple - The inputs, unchanged. - """ - return args - - def make_gnn_seq( edge_index, num_gnn_layers, @@ -513,9 +493,8 @@ def make_gnn_seq( 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 - substitute an ``IdentityModule`` themselves rather than calling - this with 0. + 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 @@ -540,7 +519,7 @@ def make_gnn_seq( if num_gnn_layers < 1: raise ValueError( "make_gnn_seq requires num_gnn_layers >= 1 " - f"(got {num_gnn_layers}); use an IdentityModule for a no-op stage." + f"(got {num_gnn_layers}); skip the stage for a no-op." ) gnn_class = get_gnn_class(gnn_type) return pyg.nn.Sequential( diff --git a/tests/test_latent_modules.py b/tests/test_latent_modules.py index e4f06803..377182d1 100644 --- a/tests/test_latent_modules.py +++ b/tests/test_latent_modules.py @@ -19,7 +19,7 @@ HiGraphLatentDecoder, HiGraphLatentEncoder, ) -from neural_lam.utils import IdentityModule, make_gnn_seq +from neural_lam.utils import make_gnn_seq def _fully_connected_edge_index(n_send, n_rec): @@ -80,16 +80,9 @@ def flat_graph_emb(flat_dims, flat_edges): } -def test_identity_module_passes_args_through(): - module = IdentityModule() - a, b, c = torch.randn(3), torch.randn(2), torch.randn(1) - out = module(a, b, c) - assert out == (a, b, c) - - def test_make_gnn_seq_zero_layers_raises(): - """make_gnn_seq must build a real sequence; the no-op (identity) case is - the caller's responsibility, exercised via the zero-intra-layer tests.""" + """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( @@ -264,11 +257,11 @@ def test_graph_decoder_no_output_std_returns_none( assert pred_std is None -def test_flat_modules_zero_m2m_layers_use_identity( +def test_flat_modules_zero_m2m_layers_skip_processing( flat_dims, flat_edges, flat_graph_emb ): - """m2m_layers=0 routes on-mesh processing through IdentityModule at the - call site (make_gnn_seq itself rejects 0). Exercise both flat modules.""" + """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"], @@ -277,7 +270,7 @@ def test_flat_modules_zero_m2m_layers_use_identity( m2m_layers=0, hidden_layers=flat_dims["hidden_layers"], ) - assert isinstance(enc.m2m_gnns, IdentityModule) + assert enc.m2m_gnns is None dec = GraphLatentDecoder( g2m_edge_index=flat_edges["g2m"], @@ -289,7 +282,7 @@ def test_flat_modules_zero_m2m_layers_use_identity( m2m_layers=0, hidden_layers=flat_dims["hidden_layers"], ) - assert isinstance(dec.m2m_gnns, IdentityModule) + assert dec.m2m_gnns is None B = flat_dims["batch_size"] grid_rep = torch.randn(B, flat_dims["num_grid"], flat_dims["hidden_dim"]) @@ -537,8 +530,8 @@ def test_hi_graph_modules_reject_single_level(): def test_hi_graph_decoder_zero_intra_layers(hi_dims, hi_edges, hi_graph_emb): - """intra_level_layers=0 routes intra-processing through IdentityModule - (the make_gnn_seq branch). Exercise that path end-to-end.""" + """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"], From d2d249b6b309e75b02d78be84376eae794fbf1a5 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 00:00:37 +0530 Subject: [PATCH 11/51] docs: align ConstantLatentEncoder.compute_dist_params docstring with base class Use the base class summary and expand on it with the constant-specific behavior, per review. --- neural_lam/models/latent/constant_encoder.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/neural_lam/models/latent/constant_encoder.py b/neural_lam/models/latent/constant_encoder.py index 35f31fda..f1bf3be5 100644 --- a/neural_lam/models/latent/constant_encoder.py +++ b/neural_lam/models/latent/constant_encoder.py @@ -34,20 +34,25 @@ def __init__(self, latent_dim, num_mesh_nodes, output_dist="isotropic"): def compute_dist_params(self, grid_rep, **kwargs): """ - Return constant (zero) distribution parameters. + 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)``. Used only to determine - batch size and device; the values do not affect the output. + Shape ``(B, num_grid_nodes, d_h)``. Grid input representation, + used only to determine batch size and device. **kwargs - Ignored. + Ignored; accepted for compatibility with the base class + interface. Returns ------- torch.Tensor - Shape ``(B, num_mesh_nodes, output_dim)``. All zeros. + Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of + the latent distribution, all zeros. """ return torch.zeros( grid_rep.shape[0], From 527227c41ea35127fd76ae77869ab458ac7d0e8a Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 00:01:51 +0530 Subject: [PATCH 12/51] Update neural_lam/models/latent/base_decoder.py Co-authored-by: Joel Oskarsson --- neural_lam/models/latent/base_decoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index 94deecce..e42c9e60 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -15,7 +15,7 @@ class BaseGraphLatentDecoder(nn.Module): 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 + softplus std) depending on + ``2 * num_state_vars`` outputs (mean, std) depending on ``output_std``. """ From b9d22ced8b5639b3a789cc8414d8f9e1b5b884ec Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 00:05:52 +0530 Subject: [PATCH 13/51] docs: expand parameter descriptions in latent decoder docstrings Describe the role of each representation in the message passing (sender/receiver, where the latent enters, purpose of the residual grid rep) in BaseGraphLatentDecoder and the inheriting decoders, per review. --- neural_lam/models/latent/base_decoder.py | 60 ++++++++++++++------ neural_lam/models/latent/graph_decoder.py | 22 ++++--- neural_lam/models/latent/hi_graph_decoder.py | 24 +++++--- 3 files changed, 74 insertions(+), 32 deletions(-) diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index e42c9e60..671b0709 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -34,15 +34,23 @@ def __init__( ---------- 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. + 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. + 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. + 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; if False, only the mean. + 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__() @@ -73,20 +81,29 @@ def combine_with_latent( Parameters ---------- original_grid_rep : torch.Tensor - Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + 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)``. Embedded latent sample. + 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)``. Grid representation to use - for residual connections. + 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 graph node and edge features. + 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. + representation, incorporating both the grid input and the + latent sample. """ raise NotImplementedError("combine_with_latent not implemented") @@ -94,20 +111,28 @@ 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)``. Grid input representation. + 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. + 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 + current time step, used as the base of the residual mean prediction. graph_emb : dict - Embedded graph node and edge features, with at least ``g2m``, - ``m2m`` and ``m2g`` entries. + Embedded static graph node and edge features, forwarded to + ``combine_with_latent``; includes at least the ``g2m``, + ``m2m`` and ``m2g`` edge embeddings. Returns ------- @@ -117,7 +142,8 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): 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. + next state, obtained from the output parameter map through a + softplus. """ latent_emb = self.latent_embedder(latent_samples) diff --git a/neural_lam/models/latent/graph_decoder.py b/neural_lam/models/latent/graph_decoder.py index 0ac4b180..d1ab64f5 100644 --- a/neural_lam/models/latent/graph_decoder.py +++ b/neural_lam/models/latent/graph_decoder.py @@ -100,21 +100,29 @@ def combine_with_latent( Parameters ---------- original_grid_rep : torch.Tensor - Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + 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)``. Embedded latent sample. + 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)``. Grid representation used - as receiver in the mesh-to-grid step. + 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 graph node and edge features, with at least ``g2m``, - ``m2m`` and ``m2g`` entries. + 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. + representation, incorporating both the grid input and the + latent sample. """ mesh_rep = self.g2m_gnn(original_grid_rep, latent_rep, graph_emb["g2m"]) diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index eb972705..63e8344d 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -183,27 +183,35 @@ def combine_with_latent( Parameters ---------- original_grid_rep : torch.Tensor - Shape ``(B, num_grid_nodes, d_h)``. Grid representation. + 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)``. Embedded latent sample - on the top mesh level ``L``. + 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)``. Grid representation used - as receiver in the mesh-to-grid step. + 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 graph node and edge features, with at least entries + 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)``. + ``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. + 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"] From f38036da4d5b07f368652b334ef335efaba17fcf Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 00:08:31 +0530 Subject: [PATCH 14/51] docs: clarify role of grid update MLP in latent decoder forward --- neural_lam/models/latent/base_decoder.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index 671b0709..936320bb 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -147,6 +147,10 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): """ 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( From c45d8ac5d911247513d06be850e53d6256499ec7 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 00:09:50 +0530 Subject: [PATCH 15/51] docs: explain mean/std chunking of decoder output parameters --- neural_lam/models/latent/base_decoder.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index 936320bb..bd19b111 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -160,6 +160,11 @@ def forward(self, grid_rep, latent_samples, last_state, 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: From cdaa6f9c1bc527290259e998badb85bf9bc645e8 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 01:03:58 +0530 Subject: [PATCH 16/51] refactor: split Graph-EFM into hierarchical and flat subclasses Replace the single GraphEFM class that resolved flat vs hierarchical at construction with an explicit class per graph type, per review: - BaseGraphEFM: graph-type independent setup (graph loading, grid and grid-mesh edge embedders, per-variable std) and all shared behavior (forward, compute_step_loss, estimate_likelihood, sampling helpers). Validates the loaded graph against the subclass's requires_hierarchical and exposes an embedd_mesh hook used by embedd_all. - GraphEFM: hierarchical mesh graphs; builds per-level mesh embedders and HiGraphLatentEncoder/HiGraphLatentDecoder modules. - GraphEFMMS: flat (e.g. multi-scale) mesh graphs; builds the flat mesh embedders and GraphLatentEncoder/GraphLatentDecoder modules. learn_prior remains a constructor flag on both subclasses. Tests select the class per graph type and cover the graph-type mismatch error. --- neural_lam/models/__init__.py | 14 +- .../models/step_predictors/graph/graph_efm.py | 721 ++++++++++++------ tests/test_graph_efm_predictor.py | 35 +- 3 files changed, 507 insertions(+), 263 deletions(-) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 66ae3ec3..68e2a4e1 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -6,18 +6,18 @@ from .module import ForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel -from .step_predictors.graph.graph_efm import GraphEFM +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 -# NOTE: GraphEFM is 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 GraphEFM requires ``config`` (for its per_var_std -# weighting) and takes no ``mesh_aggr``. Registering it requires config-aware -# model assembly in train_model.py. +# 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/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 0ac7de5a..e70a4546 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -1,5 +1,6 @@ -"""Graph-based Ensemble Forecasting Model (GraphEFM) single-step -predictor.""" +"""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 @@ -23,9 +24,10 @@ from ..base import StepPredictor -class GraphEFM(StepPredictor): +class BaseGraphEFM(StepPredictor): """ - Graph-based Ensemble Forecasting Model -- single-step predictor. + 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 @@ -37,34 +39,38 @@ class GraphEFM(StepPredictor): pred_std)``) and sampling helpers. Rollout, ELBO assembly, ensemble logic and logging live outside the predictor. - One class handles both flat and hierarchical meshes, resolved at - construction from ``self.hierarchical`` (set by ``utils.load_graph``). + 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 = "hierarchical", + graph_name: str, hidden_dim: int = 64, hidden_layers: int = 1, - latent_dim: Optional[int] = None, - prior_processor_layers: int = 2, - encoder_processor_layers: int = 2, - processor_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, sample_obs_noise: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None, ): """ - Build the prior, variational encoder and latent decoder sub-models. + 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 ---------- @@ -76,38 +82,16 @@ def __init__( and variable counts. graph_name : str Name of the graph directory (under ``/graph``) to load. - Both flat and hierarchical graphs are supported; which latent - modules are built is resolved from the loaded graph. + 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. - latent_dim : int, optional - Dimensionality of the latent variable at each mesh node; - defaults to ``hidden_dim`` when None. - prior_processor_layers : int - Number of processor GNN layers in the (learned) prior. - encoder_processor_layers : int - Number of processor GNN layers in the variational encoder. - processor_layers : int - Number of processor 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 @@ -138,6 +122,15 @@ def __init__( self.hierarchical, graph_ldict = utils.load_graph( graph_dir_path=graph_dir_path ) + 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}" + ) for name, attr_value in graph_ldict.items(): # Make BufferLists module members and register tensors as buffers if isinstance(attr_value, torch.Tensor): @@ -146,7 +139,8 @@ def __init__( setattr(self, name, attr_value) # Specify dimensions of data - num_state_vars = datastore.get_num_data_vars(category="state") + self.num_state_vars = datastore.get_num_data_vars(category="state") + num_state_vars = self.num_state_vars num_forcing_vars = datastore.get_num_data_vars(category="forcing") grid_static_dim = self.grid_static_features.shape[1] # grid_dim: total grid input dim, same formula as BaseGraphModel, @@ -175,175 +169,6 @@ def __init__( self.g2m_embedder = utils.make_mlp([g2m_dim] + self.mlp_blueprint_end) self.m2g_embedder = utils.make_mlp([m2g_dim] + self.mlp_blueprint_end) - if self.hierarchical: - level_mesh_sizes = [ - mesh_feat.shape[0] for mesh_feat in self.mesh_static_features - ] - 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 processor layers, no need to embed m2m - self.embedd_m2m = ( - max( - prior_processor_layers, - encoder_processor_layers, - processor_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) - ] - ) - else: - 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)" - ) - 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: - if self.hierarchical: - 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_processor_layers, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - output_dist=prior_dist, - ) - else: - 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_processor_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 - if self.hierarchical: - 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_processor_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=num_state_vars, - intra_level_layers=processor_layers, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - m2g_gnn_type=m2g_gnn_type, - output_std=bool(output_std), - ) - else: - 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_processor_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=num_state_vars, - m2m_layers=processor_layers, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - m2g_gnn_type=m2g_gnn_type, - output_std=bool(output_std), - ) - # 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 @@ -445,6 +270,25 @@ def embedd_current( 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_all(self, prev_state, prev_prev_state, forcing): """ Embed all node and edge representations. @@ -492,46 +336,7 @@ def embedd_all(self, prev_state, prev_prev_state, forcing): self.m2g_embedder(self.m2g_features), batch_size ), # (B, M_m2g, d_h) } - - if self.hierarchical: - graph_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) - - if self.embedd_m2m: - graph_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 - graph_emb["m2m"] = list(self.m2m_features) - - graph_emb["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 - ) - ] - graph_emb["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 - ) - ] - else: - graph_emb["mesh"] = self.expand_to_batch( - self.mesh_embedder(self.mesh_static_features), batch_size - ) # (B, num_mesh_nodes, d_h) - graph_emb["m2m"] = self.expand_to_batch( - self.m2m_embedder(self.m2m_features), batch_size - ) # (B, M_m2m, d_h) + graph_emb.update(self.embedd_mesh(batch_size)) return grid_emb, graph_emb @@ -744,3 +549,419 @@ def forward( ) # (B, num_grid_nodes, d_state) return self.sample_next_state(pred_mean, pred_std), 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_processor_layers: int = 2, + encoder_processor_layers: int = 2, + processor_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, + sample_obs_noise: 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. + + See :meth:`BaseGraphEFM.__init__` for the shared parameters + (``config``, ``datastore``, ``graph_name``, ``hidden_dim``, + ``hidden_layers``, ``num_past_forcing_steps``, + ``num_future_forcing_steps``, ``output_std``, ``sample_obs_noise`` + and the clamping limits). + + Parameters + ---------- + latent_dim : int, optional + Dimensionality of the latent variable at each top-level mesh + node; defaults to ``hidden_dim`` when None. + prior_processor_layers : int + Number of intra-level GNN layers in the (learned) prior. + encoder_processor_layers : int + Number of intra-level GNN layers in the variational encoder. + processor_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"``. + 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``). + """ + 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, + sample_obs_noise=sample_obs_noise, + 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 processor layers, no need to embed m2m + self.embedd_m2m = ( + max( + prior_processor_layers, + encoder_processor_layers, + processor_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_processor_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_processor_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=processor_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_processor_layers: int = 2, + encoder_processor_layers: int = 2, + processor_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, + sample_obs_noise: 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. + + See :meth:`BaseGraphEFM.__init__` for the shared parameters + (``config``, ``datastore``, ``graph_name``, ``hidden_dim``, + ``hidden_layers``, ``num_past_forcing_steps``, + ``num_future_forcing_steps``, ``output_std``, ``sample_obs_noise`` + and the clamping limits). + + Parameters + ---------- + latent_dim : int, optional + Dimensionality of the latent variable at each mesh node; + defaults to ``hidden_dim`` when None. + prior_processor_layers : int + Number of on-mesh (m2m) GNN layers in the (learned) prior. + encoder_processor_layers : int + Number of on-mesh (m2m) GNN layers in the variational encoder. + processor_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"``. + 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``). + """ + 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, + sample_obs_noise=sample_obs_noise, + 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_processor_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_processor_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=processor_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/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py index c2aab4e3..922a4634 100644 --- a/tests/test_graph_efm_predictor.py +++ b/tests/test_graph_efm_predictor.py @@ -1,9 +1,10 @@ -"""Unit tests for the GraphEFM single-step probabilistic predictor. +"""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 flat and hierarchical variants on the -real example datastore with a freshly created graph, then exercise ``forward``, -``compute_step_loss`` and the sampling helpers on synthetic tensors. +(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 @@ -18,7 +19,10 @@ 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 +from neural_lam.models.step_predictors.graph.graph_efm import ( + GraphEFM, + GraphEFMMS, +) from tests.conftest import init_datastore_example NUM_PAST_FORCING_STEPS = 1 @@ -55,7 +59,8 @@ def _datastore_and_config_with_graph(graph_name): def _build_predictor(graph_name, output_std=False, sample_obs_noise=False): datastore, config = _datastore_and_config_with_graph(graph_name) - predictor = GraphEFM( + predictor_class = GraphEFM if graph_name == "hierarchical" else GraphEFMMS + predictor = predictor_class( config=config, datastore=datastore, graph_name=graph_name, @@ -253,3 +258,21 @@ def test_per_var_std_none_when_output_std(): 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, + ) From 0663a875797fb1d1f36c33544cf52e0b406ff1a1 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 01:10:12 +0530 Subject: [PATCH 17/51] refactor: avoid processor terminology in Graph-EFM parameters Rename the layer-count parameters to say what the layers are, matching the latent module parameter names: prior/encoder/decoder_intra_level_ layers on GraphEFM (hierarchical) and prior/encoder/decoder_m2m_layers on GraphEFMMS (flat), per review. --- .../models/step_predictors/graph/graph_efm.py | 47 ++++++++++--------- tests/test_graph_efm_predictor.py | 19 ++++++-- 2 files changed, 39 insertions(+), 27 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index e70a4546..ade70418 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -31,7 +31,8 @@ class BaseGraphEFM(StepPredictor): A latent-variable step predictor consisting of a conditional prior, a variational encoder and a latent decoder, each of which carries its own - g2m/processor/m2g GNNs. The encode-process-decode backbone of + 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 @@ -570,9 +571,9 @@ def __init__( hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: Optional[int] = None, - prior_processor_layers: int = 2, - encoder_processor_layers: int = 2, - processor_layers: int = 4, + 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, @@ -598,11 +599,11 @@ def __init__( latent_dim : int, optional Dimensionality of the latent variable at each top-level mesh node; defaults to ``hidden_dim`` when None. - prior_processor_layers : int + prior_intra_level_layers : int Number of intra-level GNN layers in the (learned) prior. - encoder_processor_layers : int + encoder_intra_level_layers : int Number of intra-level GNN layers in the variational encoder. - processor_layers : int + 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 @@ -678,12 +679,12 @@ def __init__( for _ in range(num_levels - 1) ] ) - # If not using any processor layers, no need to embed m2m + # If not using any intra-level layers, no need to embed m2m self.embedd_m2m = ( max( - prior_processor_layers, - encoder_processor_layers, - processor_layers, + prior_intra_level_layers, + encoder_intra_level_layers, + decoder_intra_level_layers, ) > 0 ) @@ -707,7 +708,7 @@ def __init__( m2m_edge_index=self.m2m_edge_index, mesh_up_edge_index=self.mesh_up_edge_index, hidden_dim=hidden_dim, - intra_level_layers=prior_processor_layers, + intra_level_layers=prior_intra_level_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, output_dist=prior_dist, @@ -726,7 +727,7 @@ def __init__( m2m_edge_index=self.m2m_edge_index, mesh_up_edge_index=self.mesh_up_edge_index, hidden_dim=hidden_dim, - intra_level_layers=encoder_processor_layers, + intra_level_layers=encoder_intra_level_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, output_dist="diagonal", @@ -740,7 +741,7 @@ def __init__( hidden_dim=hidden_dim, latent_dim=latent_dim, num_state_vars=self.num_state_vars, - intra_level_layers=processor_layers, + intra_level_layers=decoder_intra_level_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, m2g_gnn_type=m2g_gnn_type, @@ -817,9 +818,9 @@ def __init__( hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: Optional[int] = None, - prior_processor_layers: int = 2, - encoder_processor_layers: int = 2, - processor_layers: int = 4, + 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, @@ -845,11 +846,11 @@ def __init__( latent_dim : int, optional Dimensionality of the latent variable at each mesh node; defaults to ``hidden_dim`` when None. - prior_processor_layers : int + prior_m2m_layers : int Number of on-mesh (m2m) GNN layers in the (learned) prior. - encoder_processor_layers : int + encoder_m2m_layers : int Number of on-mesh (m2m) GNN layers in the variational encoder. - processor_layers : int + 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 @@ -905,7 +906,7 @@ def __init__( g2m_edge_index=self.g2m_edge_index, m2m_edge_index=self.m2m_edge_index, hidden_dim=hidden_dim, - m2m_layers=prior_processor_layers, + m2m_layers=prior_m2m_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, output_dist=prior_dist, @@ -923,7 +924,7 @@ def __init__( g2m_edge_index=self.g2m_edge_index, m2m_edge_index=self.m2m_edge_index, hidden_dim=hidden_dim, - m2m_layers=encoder_processor_layers, + m2m_layers=encoder_m2m_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, output_dist="diagonal", @@ -935,7 +936,7 @@ def __init__( hidden_dim=hidden_dim, latent_dim=latent_dim, num_state_vars=self.num_state_vars, - m2m_layers=processor_layers, + m2m_layers=decoder_m2m_layers, hidden_layers=hidden_layers, g2m_gnn_type=g2m_gnn_type, m2g_gnn_type=m2g_gnn_type, diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py index 922a4634..51ea4c5f 100644 --- a/tests/test_graph_efm_predictor.py +++ b/tests/test_graph_efm_predictor.py @@ -59,7 +59,20 @@ def _datastore_and_config_with_graph(graph_name): def _build_predictor(graph_name, output_std=False, sample_obs_noise=False): datastore, config = _datastore_and_config_with_graph(graph_name) - predictor_class = GraphEFM if graph_name == "hierarchical" else GraphEFMMS + 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, @@ -67,15 +80,13 @@ def _build_predictor(graph_name, output_std=False, sample_obs_noise=False): hidden_dim=4, hidden_layers=1, latent_dim=4, - prior_processor_layers=1, - encoder_processor_layers=1, - processor_layers=1, 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, sample_obs_noise=sample_obs_noise, + **layer_kwargs, ) return predictor, datastore, config From 5d1e75bb1aca9b746b11e04320fc1805034f4470 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 01:31:23 +0530 Subject: [PATCH 18/51] docs: document all constructor parameters in Graph-EFM subclasses Make the GraphEFM and GraphEFMMS __init__ docstrings self-contained with the full parameter list, instead of pointing at the base class for the shared ones, per review. --- .../models/step_predictors/graph/graph_efm.py | 68 +++++++++++++++---- 1 file changed, 56 insertions(+), 12 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index ade70418..39325755 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -588,14 +588,21 @@ def __init__( """ Build the mesh embedders and hierarchical latent modules. - See :meth:`BaseGraphEFM.__init__` for the shared parameters - (``config``, ``datastore``, ``graph_name``, ``hidden_dim``, - ``hidden_layers``, ``num_past_forcing_steps``, - ``num_future_forcing_steps``, ``output_std``, ``sample_obs_noise`` - and the clamping limits). - 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. @@ -612,12 +619,27 @@ def __init__( 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. + sample_obs_noise : bool + If True, sample observation noise when rolling out; if False, + ``sample_next_state`` returns the predicted mean. + 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, @@ -835,14 +857,21 @@ def __init__( """ Build the mesh embedders and flat-graph latent modules. - See :meth:`BaseGraphEFM.__init__` for the shared parameters - (``config``, ``datastore``, ``graph_name``, ``hidden_dim``, - ``hidden_layers``, ``num_past_forcing_steps``, - ``num_future_forcing_steps``, ``output_std``, ``sample_obs_noise`` - and the clamping limits). - 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. @@ -859,12 +888,27 @@ def __init__( 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. + sample_obs_noise : bool + If True, sample observation noise when rolling out; if False, + ``sample_next_state`` returns the predicted mean. + 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, From e29e0c7660aea97a7d686bf8a8a86464073227da Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 01:43:05 +0530 Subject: [PATCH 19/51] refactor: remove sample_obs_noise option from Graph-EFM Sampling uncorrelated Gaussian observation noise per grid node is not useful in practice, so the option is removed entirely rather than left to tempt users, per review. forward now returns the decoder mean directly (the prediction is stochastic only through the latent sample) and the trivial sample_next_state helper is dropped. --- .../models/step_predictors/graph/graph_efm.py | 58 +++---------------- tests/test_graph_efm_predictor.py | 24 +------- 2 files changed, 9 insertions(+), 73 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 39325755..d147ec04 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -37,8 +37,8 @@ class BaseGraphEFM(StepPredictor): ``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)``) and sampling helpers. Rollout, ELBO assembly, ensemble - logic and logging live outside the predictor. + 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 @@ -61,7 +61,6 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - sample_obs_noise: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None, ): @@ -97,9 +96,6 @@ def __init__( If True, the decoder outputs a per-variable std alongside the mean; if False, a constant per-variable std is used as likelihood scale. - sample_obs_noise : bool - If True, sample observation noise when rolling out; if False, - ``sample_next_state`` returns the predicted mean. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional @@ -112,10 +108,6 @@ def __init__( output_clamping_upper=output_clamping_upper, ) - # Whether to sample observation noise during rollout. When False, - # sample_next_state returns the predicted mean. - self.sample_obs_noise = bool(sample_obs_noise) - # Load graph with static features (same pattern as BaseGraphModel). # NOTE: (IMPORTANT!) mesh nodes MUST have the first # num_mesh_nodes indices. @@ -200,33 +192,6 @@ def __init__( # inert -- accepted for interface parity with other StepPredictors. self.prepare_clamping_params(datastore) - def sample_next_state(self, pred_mean, pred_std): - """ - Sample state at next time step given a Gaussian observation model. - If ``self.sample_obs_noise`` is False, only return the mean. - - Parameters - ---------- - pred_mean : torch.Tensor - Shape ``(B, num_grid_nodes, d_state)``. Predicted mean. - pred_std : torch.Tensor or None - Shape ``(B, num_grid_nodes, d_state)``, or None when the decoder - does not output a std (``output_std=False``). - - Returns - ------- - torch.Tensor - Shape ``(B, num_grid_nodes, d_state)``. Next state. - """ - if not self.output_std: - pred_std = self.per_var_std # (d_f,) - - if self.sample_obs_noise: - return torch.distributions.Normal(pred_mean, pred_std).rsample() - # (B, num_grid_nodes, d_state) - - return pred_mean # (B, num_grid_nodes, d_state) - def embedd_current( self, prev_state, @@ -510,7 +475,9 @@ def forward( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """ Sample one time step prediction: embed features, sample the latent - from the prior, decode, and return the sampled next state. + 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 ---------- @@ -524,7 +491,8 @@ def forward( Returns ------- new_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_state)``. Sampled ``X_{t+1}``. + 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. @@ -549,7 +517,7 @@ def forward( grid_prev_emb, latent_samples, last_state, graph_emb ) # (B, num_grid_nodes, d_state) - return self.sample_next_state(pred_mean, pred_std), pred_std + return pred_mean, pred_std class GraphEFM(BaseGraphEFM): @@ -581,7 +549,6 @@ def __init__( g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", output_std: bool = False, - sample_obs_noise: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None, ): @@ -633,9 +600,6 @@ def __init__( If True, the decoder outputs a per-variable std alongside the mean; if False, a constant per-variable std is used as likelihood scale. - sample_obs_noise : bool - If True, sample observation noise when rolling out; if False, - ``sample_next_state`` returns the predicted mean. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional @@ -650,7 +614,6 @@ def __init__( num_past_forcing_steps=num_past_forcing_steps, num_future_forcing_steps=num_future_forcing_steps, output_std=output_std, - sample_obs_noise=sample_obs_noise, output_clamping_lower=output_clamping_lower, output_clamping_upper=output_clamping_upper, ) @@ -850,7 +813,6 @@ def __init__( g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", output_std: bool = False, - sample_obs_noise: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None, ): @@ -902,9 +864,6 @@ def __init__( If True, the decoder outputs a per-variable std alongside the mean; if False, a constant per-variable std is used as likelihood scale. - sample_obs_noise : bool - If True, sample observation noise when rolling out; if False, - ``sample_next_state`` returns the predicted mean. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional @@ -919,7 +878,6 @@ def __init__( num_past_forcing_steps=num_past_forcing_steps, num_future_forcing_steps=num_future_forcing_steps, output_std=output_std, - sample_obs_noise=sample_obs_noise, output_clamping_lower=output_clamping_lower, output_clamping_upper=output_clamping_upper, ) diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py index 51ea4c5f..063fd81a 100644 --- a/tests/test_graph_efm_predictor.py +++ b/tests/test_graph_efm_predictor.py @@ -57,7 +57,7 @@ def _datastore_and_config_with_graph(graph_name): return datastore, config -def _build_predictor(graph_name, output_std=False, sample_obs_noise=False): +def _build_predictor(graph_name, output_std=False): datastore, config = _datastore_and_config_with_graph(graph_name) if graph_name == "hierarchical": predictor_class = GraphEFM @@ -85,7 +85,6 @@ def _build_predictor(graph_name, output_std=False, sample_obs_noise=False): num_past_forcing_steps=NUM_PAST_FORCING_STEPS, num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, output_std=output_std, - sample_obs_noise=sample_obs_noise, **layer_kwargs, ) return predictor, datastore, config @@ -223,27 +222,6 @@ def test_forward_member_stochasticity(graph_name): assert not torch.allclose(out_a, out_b) -def test_sample_next_state_respects_sample_obs_noise(): - """sample_next_state returns the mean when sample_obs_noise is False and a - stochastic draw (different from the mean) when True.""" - deterministic, datastore, _ = _build_predictor( - "1level", sample_obs_noise=False - ) - d_state = datastore.get_num_data_vars(category="state") - # Last dim must match per_var_std (d_state,) for the obs-noise broadcast. - pred_mean = torch.randn(2, 5, d_state) - - out_mean = deterministic.sample_next_state(pred_mean, pred_std=None) - assert torch.equal(out_mean, pred_mean) - - stochastic, _, _ = _build_predictor("1level", sample_obs_noise=True) - # per_var_std is registered (output_std=False); the draw should differ - # from the mean. - out_sampled = stochastic.sample_next_state(pred_mean, pred_std=None) - assert out_sampled.shape == pred_mean.shape - assert not torch.allclose(out_sampled, pred_mean) - - def test_per_var_std_matches_module_formula(): """per_var_std mirrors ForecasterModule's formula: state_diff_std_standardized / sqrt(state_feature_weights).""" From 441be4b84c0285f5d4ba8cb7236f2cc283f5e3a1 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 11:02:07 +0530 Subject: [PATCH 20/51] refactor: extract shared graph-setup helpers to remove duplication BaseGraphModel and BaseGraphEFM duplicated their graph loading, buffer registration and grid-input-dim computation. Factor these into two utils helpers used by both: - utils.load_and_register_graph(module, datastore, graph_name): loads the graph and registers its tensors/BufferLists on the module, returning whether it is hierarchical. - utils.grid_input_dim(datastore, grid_static_dim, num_past_forcing_ steps, num_future_forcing_steps): the total grid input dimensionality. This keeps the two model families' grid-feature setup in one place (e.g. for a future boundary-forcing input) without coupling their differing forward passes or submodule sets via inheritance. --- .../models/step_predictors/graph/base.py | 24 ++---- .../models/step_predictors/graph/graph_efm.py | 30 +++---- neural_lam/utils.py | 84 ++++++++++++++++++- 3 files changed, 100 insertions(+), 38 deletions(-) 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 index d147ec04..4da9d398 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -108,12 +108,11 @@ def __init__( output_clamping_upper=output_clamping_upper, ) - # Load graph with static features (same pattern as BaseGraphModel). + # 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 ) if self.hierarchical != self.requires_hierarchical: required_type = ( @@ -124,26 +123,17 @@ def __init__( f"{type(self).__name__} requires a {required_type} mesh " f"graph, but graph '{graph_name}' is {loaded_type}" ) - 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_state_vars = datastore.get_num_data_vars(category="state") num_state_vars = self.num_state_vars - num_forcing_vars = datastore.get_num_data_vars(category="forcing") - grid_static_dim = self.grid_static_features.shape[1] - # grid_dim: total grid input dim, same formula as BaseGraphModel, - # matching the cat order in embedd_all/embedd_current - # (prev_prev, prev, forcing, static[, current]). - self.grid_dim = ( - 2 * num_state_vars - + grid_static_dim - + num_forcing_vars - * (num_past_forcing_steps + num_future_forcing_steps + 1) + # 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] diff --git a/neural_lam/utils.py b/neural_lam/utils.py index eba72d1a..2cb02e05 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -9,7 +9,7 @@ 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 @@ -24,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): """ @@ -437,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 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. From e12f13955ced3d38e894f22d7b37906643448d74 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 11:24:33 +0530 Subject: [PATCH 21/51] refactor: give descriptive names to Graph-EFM grid embedding methods Rename embedd_all -> embedd_grid_and_graph (embeds the grid for states up to t-1 plus the full graph) and embedd_current -> embedd_grid_with_target (embeds the grid including the target state, for the encoder), per review. --- .../models/step_predictors/graph/graph_efm.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 4da9d398..44549049 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -182,7 +182,7 @@ def __init__( # inert -- accepted for interface parity with other StepPredictors. self.prepare_clamping_params(datastore) - def embedd_current( + def embedd_grid_with_target( self, prev_state, prev_prev_state, @@ -245,9 +245,9 @@ def embedd_mesh(self, batch_size): """ raise NotImplementedError("embedd_mesh not implemented") - def embedd_all(self, prev_state, prev_prev_state, forcing): + def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): """ - Embed all node and edge representations. + Embed the grid (states up to t-1) and the full graph. Parameters ---------- @@ -324,9 +324,9 @@ def estimate_likelihood( 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_all``. + ``embedd_grid_and_graph``. graph_emb : dict - Edge/mesh embeddings from ``embedd_all``. + 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 @@ -411,13 +411,13 @@ def compute_step_loss( Shape ``(B, num_grid_nodes, d_state)`` or ``(d_state,)``. """ # embed all features - grid_prev_emb, graph_emb = self.embedd_all( + 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_current( + grid_current_emb = self.embedd_grid_with_target( prev_states[:, 1], prev_states[:, 0], forcing_features, @@ -488,7 +488,7 @@ def forward( otherwise None. """ # embed all features - grid_prev_emb, graph_emb = self.embedd_all( + grid_prev_emb, graph_emb = self.embedd_grid_and_graph( prev_state, prev_prev_state, forcing ) From 4a11d2025a886f8225fd253c77f034ebdb34558d Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 13 Jun 2026 11:29:20 +0530 Subject: [PATCH 22/51] refactor: drop redundant last_state alias in Graph-EFM forward In forward, prev_state is already X_t, so pass it directly to the decoder instead of aliasing it to last_state, per review. --- neural_lam/models/step_predictors/graph/graph_efm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 44549049..033f8b68 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -501,10 +501,10 @@ def forward( latent_samples = prior_dist.rsample() # (B, num_mesh_nodes, d_latent) - # Compute reconstruction (decoder) - last_state = prev_state + # 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, last_state, graph_emb + grid_prev_emb, latent_samples, prev_state, graph_emb ) # (B, num_grid_nodes, d_state) return pred_mean, pred_std From 99d86c8f053ea2858b6624d4f3796e3ec8aca2f2 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 18 Jun 2026 14:52:14 +0530 Subject: [PATCH 23/51] Update neural_lam/utils.py Co-authored-by: Joel Oskarsson --- neural_lam/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 2cb02e05..211beb56 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -480,7 +480,7 @@ def load_and_register_graph( return hierarchical -def grid_input_dim( +def compute_grid_input_dim( datastore: "BaseDatastore", grid_static_dim: int, num_past_forcing_steps: int, From 76c0dabd91a131d2c4c7c8a315b7e408dae7d79d Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 03:15:10 +0530 Subject: [PATCH 24/51] refactor: move graph-type compatibility check into GraphEFM subclasses --- .../models/step_predictors/graph/base.py | 2 +- .../models/step_predictors/graph/graph_efm.py | 39 ++++++++----------- 2 files changed, 18 insertions(+), 23 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 9576b7a2..17ccc61b 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -112,7 +112,7 @@ def __init__( ) # Compute grid_input_dim: total input dimensionality on the grid - self.grid_input_dim = utils.grid_input_dim( + self.grid_input_dim = utils.compute_grid_input_dim( datastore, self.grid_static_features.shape[1], num_past_forcing_steps, diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 033f8b68..81b0129e 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -41,16 +41,13 @@ class BaseGraphEFM(StepPredictor): 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 + graph type. Concrete subclasses are specific to a graph type: their + constructors verify the loaded graph is of the expected type, 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, @@ -82,8 +79,7 @@ def __init__( 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``). + Must be of the graph type required by the concrete subclass. hidden_dim : int Dimensionality of internal node and edge representations. hidden_layers : int @@ -114,22 +110,13 @@ def __init__( 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( + self.grid_dim = utils.compute_grid_input_dim( datastore, self.grid_static_features.shape[1], num_past_forcing_steps, @@ -519,8 +506,6 @@ class GraphEFM(BaseGraphEFM): decoder is a ``HiGraphLatentDecoder``. """ - requires_hierarchical = True - def __init__( self, config: NeuralLAMConfig, @@ -608,6 +593,12 @@ def __init__( output_clamping_upper=output_clamping_upper, ) + if not self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a hierarchical mesh graph, " + f"but graph '{graph_name}' is flat" + ) + level_mesh_sizes = [ mesh_feat.shape[0] for mesh_feat in self.mesh_static_features ] @@ -783,8 +774,6 @@ class GraphEFMMS(BaseGraphEFM): ``GraphLatentDecoder``. """ - requires_hierarchical = False - def __init__( self, config: NeuralLAMConfig, @@ -872,6 +861,12 @@ def __init__( output_clamping_upper=output_clamping_upper, ) + if self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a flat mesh graph, " + f"but graph '{graph_name}' is hierarchical" + ) + self.num_mesh_nodes = self.mesh_static_features.shape[0] utils.log_on_rank_zero( f"Loaded graph with " From 3bec96cc819fd3aa0ed7073f8bafbdfe9ae462cc Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 03:30:12 +0530 Subject: [PATCH 25/51] refactor: remove stale mesh-node-index NOTE from BaseGraphEFM --- neural_lam/models/step_predictors/graph/graph_efm.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 81b0129e..a5e86b87 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -105,8 +105,6 @@ def __init__( ) # 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 ) From ab005c34ba5d7b8e01e0d4d89c4265832c74c911 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 03:55:30 +0530 Subject: [PATCH 26/51] refactor: hoist shared prior construction into BaseGraphEFM Move the constant-prior branch (identical across subclasses) into BaseGraphEFM.build_prior and delegate the graph-specific learnable prior to build_learnable_prior, which GraphEFM and GraphEFMMS implement. --- .../models/step_predictors/graph/graph_efm.py | 245 +++++++++++++++--- 1 file changed, 204 insertions(+), 41 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index a5e86b87..f061c465 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -167,6 +167,101 @@ def __init__( # inert -- accepted for interface parity with other StepPredictors. self.prepare_clamping_params(datastore) + def build_prior( + self, + learn_prior, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the prior over the latent variable. + + When ``learn_prior`` is True the (graph-type specific) learnable prior + is delegated to :meth:`build_learnable_prior`; otherwise the constant + ``Normal(0, 1)`` prior, which is identical for every graph type, is + built here. Must be called after ``self.num_mesh_nodes`` is set. + + Parameters + ---------- + learn_prior : bool + If True, build a learnable prior conditioned on the previous + state; if False, build a constant (input-independent) prior. + latent_dim : int + Dimensionality of the latent variable at each mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the learnable prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh GNN layers in the learnable prior. + + Returns + ------- + torch.nn.Module + The prior latent encoder. + """ + if learn_prior: + return self.build_learnable_prior( + latent_dim=latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_layers, + ) + return ConstantLatentEncoder( + latent_dim=latent_dim, + num_mesh_nodes=self.num_mesh_nodes, + output_dist=prior_dist, + ) + + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the graph-type specific learnable prior encoder. + + Implemented by the concrete subclass, which knows the mesh graph type + and therefore the appropriate latent encoder class. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh GNN layers in the prior. + + Returns + ------- + torch.nn.Module + The learnable prior latent encoder. + """ + raise NotImplementedError("build_learnable_prior not implemented") + def embedd_grid_with_target( self, prev_state, @@ -662,27 +757,16 @@ def __init__( 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, - ) + # Prior (constant prior shared via the base class) + self.prior_model = self.build_prior( + learn_prior=learn_prior, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_intra_level_layers, + ) # Encoder (variational posterior) + Decoder self.encoder = HiGraphLatentEncoder( @@ -712,6 +796,51 @@ def __init__( output_std=bool(output_std), ) + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the hierarchical learnable prior encoder. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each top-level mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of intra-level GNN layers in the prior. + + Returns + ------- + HiGraphLatentEncoder + The learnable prior latent encoder. + """ + return HiGraphLatentEncoder( + latent_dim=latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + mesh_up_edge_index=self.mesh_up_edge_index, + hidden_dim=hidden_dim, + intra_level_layers=prior_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + def embedd_mesh(self, batch_size): """ Embed static mesh node and intra-mesh edge features per level. @@ -882,26 +1011,16 @@ def __init__( 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, - ) + # Prior (constant prior shared via the base class) + self.prior_model = self.build_prior( + learn_prior=learn_prior, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_m2m_layers, + ) # Encoder (variational posterior) + Decoder self.encoder = GraphLatentEncoder( @@ -928,6 +1047,50 @@ def __init__( output_std=bool(output_std), ) + def build_learnable_prior( + self, + latent_dim, + hidden_dim, + hidden_layers, + g2m_gnn_type, + prior_dist, + prior_layers, + ): + """ + Build the flat-graph learnable prior encoder. + + Parameters + ---------- + latent_dim : int + Dimensionality of the latent variable at each mesh node. + hidden_dim : int + Dimensionality of internal node and edge representations. + hidden_layers : int + Number of hidden layers in internal MLPs. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh (m2m) GNN layers in the prior. + + Returns + ------- + GraphLatentEncoder + The learnable prior latent encoder. + """ + return GraphLatentEncoder( + latent_dim=latent_dim, + g2m_edge_index=self.g2m_edge_index, + m2m_edge_index=self.m2m_edge_index, + hidden_dim=hidden_dim, + m2m_layers=prior_layers, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + output_dist=prior_dist, + ) + def embedd_mesh(self, batch_size): """ Embed static mesh node and intra-mesh edge features. From 00ec230188b7d354b437bb0f7653770dd95872b5 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 04:00:15 +0530 Subject: [PATCH 27/51] refactor: rename GraphEFMMS to GraphEFMMultiScale Spell out the multi-scale variant's name instead of the EFMMS abbreviation, keeping the GraphEFM prefix shared with the hierarchical variant. --- neural_lam/models/__init__.py | 5 +++-- neural_lam/models/step_predictors/graph/graph_efm.py | 8 ++++---- tests/test_graph_efm_predictor.py | 10 +++++----- 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 68e2a4e1..e986a872 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -6,13 +6,14 @@ from .module import ForecasterModule 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_efm import GraphEFM, GraphEFMMultiScale from .step_predictors.graph.graph_lam import GraphLAM from .step_predictors.graph.hi_lam import HiLAM from .step_predictors.graph.hi_lam_parallel import HiLAMParallel from .step_predictors.graph.hierarchical import BaseHiGraphModel -# NOTE: GraphEFM/GraphEFMMS are intentionally NOT registered in MODELS yet. +# NOTE: GraphEFM/GraphEFMMultiScale 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 diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index f061c465..8220999a 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -1,5 +1,5 @@ """Graph-based Ensemble Forecasting Model (Graph-EFM) single-step -predictors, for hierarchical (GraphEFM) and flat (GraphEFMMS) mesh +predictors, for hierarchical (GraphEFM) and flat (GraphEFMMultiScale) mesh graphs.""" # Standard library @@ -45,7 +45,7 @@ class BaseGraphEFM(StepPredictor): constructors verify the loaded graph is of the expected type, 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). + graph) and :class:`GraphEFMMultiScale` (flat graph). """ def __init__( @@ -891,10 +891,10 @@ def embedd_mesh(self, batch_size): return mesh_emb -class GraphEFMMS(BaseGraphEFM): +class GraphEFMMultiScale(BaseGraphEFM): """ Graph-based Ensemble Forecasting Model on a flat mesh graph - (Graph-EFM-MS, e.g. for multi-scale graphs). + (e.g. a multi-scale graph). The latent variable lives on the mesh nodes. The prior and variational encoder are ``GraphLatentEncoder``s and the decoder is a diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py index 063fd81a..3633c3fb 100644 --- a/tests/test_graph_efm_predictor.py +++ b/tests/test_graph_efm_predictor.py @@ -1,7 +1,7 @@ """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 +(see ``tests/test_gnn_layers.py``): build the flat (GraphEFMMultiScale) and hierarchical (GraphEFM) variants on the real example datastore with a freshly created graph, then exercise ``forward``, ``compute_step_loss`` and the sampling helpers on synthetic tensors. @@ -21,7 +21,7 @@ from neural_lam.loss_weighting import get_state_feature_weighting from neural_lam.models.step_predictors.graph.graph_efm import ( GraphEFM, - GraphEFMMS, + GraphEFMMultiScale, ) from tests.conftest import init_datastore_example @@ -67,7 +67,7 @@ def _build_predictor(graph_name, output_std=False): "decoder_intra_level_layers": 1, } else: - predictor_class = GraphEFMMS + predictor_class = GraphEFMMultiScale layer_kwargs = { "prior_m2m_layers": 1, "encoder_m2m_layers": 1, @@ -251,10 +251,10 @@ def test_per_var_std_none_when_output_std(): @pytest.mark.parametrize( "predictor_class, graph_name", - [(GraphEFM, "1level"), (GraphEFMMS, "hierarchical")], + [(GraphEFM, "1level"), (GraphEFMMultiScale, "hierarchical")], ) def test_graph_type_mismatch_raises(predictor_class, graph_name): - """GraphEFM requires a hierarchical graph and GraphEFMMS a flat one; + """GraphEFM requires a hierarchical graph and GraphEFMMultiScale 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"): From dc6600ee8be2d2dcf6d9327555166b7900bfdb69 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 04:13:47 +0530 Subject: [PATCH 28/51] refactor: derive grid static dim from datastore in compute_grid_input_dim The static feature count is already available from the datastore via get_num_data_vars(category="static"), so drop the redundant grid_static_dim argument and query it inside the function. --- neural_lam/models/step_predictors/graph/base.py | 1 - neural_lam/models/step_predictors/graph/graph_efm.py | 1 - neural_lam/utils.py | 8 +++----- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 17ccc61b..f23c5b22 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -114,7 +114,6 @@ def __init__( # Compute grid_input_dim: total input dimensionality on the grid self.grid_input_dim = utils.compute_grid_input_dim( datastore, - self.grid_static_features.shape[1], num_past_forcing_steps, num_future_forcing_steps, ) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 8220999a..4f54c94e 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -116,7 +116,6 @@ def __init__( # includes the target state, for the encoder input. self.grid_dim = utils.compute_grid_input_dim( datastore, - self.grid_static_features.shape[1], num_past_forcing_steps, num_future_forcing_steps, ) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index 211beb56..fc3c7d19 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -482,7 +482,6 @@ def load_and_register_graph( def compute_grid_input_dim( datastore: "BaseDatastore", - grid_static_dim: int, num_past_forcing_steps: int, num_future_forcing_steps: int, ) -> int: @@ -496,9 +495,7 @@ def compute_grid_input_dim( Parameters ---------- datastore : BaseDatastore - Datastore providing the number of state and forcing variables. - grid_static_dim : int - Number of static features per grid node. + Datastore providing the number of state, static and forcing variables. num_past_forcing_steps : int Number of past forcing steps included in the input window. num_future_forcing_steps : int @@ -510,10 +507,11 @@ def compute_grid_input_dim( Total grid input dimensionality. """ num_state_vars = datastore.get_num_data_vars(category="state") + num_static_vars = datastore.get_num_data_vars(category="static") num_forcing_vars = datastore.get_num_data_vars(category="forcing") return ( 2 * num_state_vars - + grid_static_dim + + num_static_vars + num_forcing_vars * (num_past_forcing_steps + num_future_forcing_steps + 1) ) From a4767114abd5ad41ed6c2fa04b270fd1532a9a5d Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Thu, 25 Jun 2026 04:25:26 +0530 Subject: [PATCH 29/51] fix: handle missing static dataarray in compute_grid_input_dim get_num_data_vars("static") can report a nonzero count even when the datastore provides no static dataarray, in which case the grid static buffer is empty. Mirror the buffer construction by treating a None static dataarray as zero static features, fixing the no-static-features case. --- neural_lam/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index fc3c7d19..02a5fe09 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -507,8 +507,16 @@ def compute_grid_input_dim( Total grid input dimensionality. """ num_state_vars = datastore.get_num_data_vars(category="state") - num_static_vars = datastore.get_num_data_vars(category="static") num_forcing_vars = datastore.get_num_data_vars(category="forcing") + # The static category is optional: when the datastore provides no static + # data array the grid carries no static features, mirroring the empty + # (N, 0) static buffer the step predictor builds in that case. + da_static = datastore.get_dataarray(category="static", split=None) + num_static_vars = ( + 0 + if da_static is None + else datastore.get_num_data_vars(category="static") + ) return ( 2 * num_state_vars + num_static_vars From 1713911ea38c1cc121deed3948fa16798b1c8915 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sun, 5 Jul 2026 09:31:42 +0530 Subject: [PATCH 30/51] feat: move training loss onto Forecaster, add probabilistic interface Add abstract Forecaster.compute_training_loss returning a finished (loss, loss_components) pair, so each forecaster owns its complete training objective. ForecasterModule.training_step now only injects the configured scoring rule, interior mask and per_var_std, and logs the result. The deterministic ARForecaster loss is unchanged in value. Add the abstract ProbabilisticForecaster (sample_ensemble capability), ProbabilisticARForecaster (sequential sampled rollouts, trains on the configured score of the ensemble mean) and a minimal ProbabilisticForecasterModule whose validation samples an ensemble and logs the RMSE of the ensemble mean. Interface design from #685. --- CHANGELOG.md | 17 + neural_lam/models/__init__.py | 5 + neural_lam/models/forecasters/__init__.py | 1 + .../models/forecasters/autoregressive.py | 73 +++++ neural_lam/models/forecasters/base.py | 64 ++++ .../models/forecasters/probabilistic.py | 257 +++++++++++++++ neural_lam/models/module.py | 30 +- neural_lam/models/probabilistic_module.py | 140 ++++++++ tests/test_probabilistic_forecaster.py | 304 ++++++++++++++++++ 9 files changed, 878 insertions(+), 13 deletions(-) create mode 100644 neural_lam/models/forecasters/probabilistic.py create mode 100644 neural_lam/models/probabilistic_module.py create mode 100644 tests/test_probabilistic_forecaster.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fd123dee..d5ab6ee8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add a general probabilistic forecasting interface: an abstract + `ProbabilisticForecaster` capable of sampling ensemble forecasts + (`sample_ensemble`, members stacked along a new dimension after batch), + its auto-regressive implementation `ProbabilisticARForecaster` (samples + independent trajectories through a stochastic step predictor and by + default trains on the configured scoring rule applied to the ensemble + mean) and a `ProbabilisticForecasterModule` whose validation samples an + ensemble and logs the RMSE of the ensemble mean. Move ownership of the + training objective from `ForecasterModule` onto the `Forecaster`: the + new abstract `Forecaster.compute_training_loss` returns a finished + `(loss, loss_components)` pair and `ForecasterModule.training_step` only + injects the configured scoring rule and interior mask and logs the + result. The deterministic `ARForecaster` training loss is unchanged in + value, only computed by the forecaster itself. + [\#685](https://github.com/mllam/neural-lam/issues/685) + @Sir-Sloth-The-Lazy + - Add `PropagationNet` GNN layer that incentivises directional message propagation from sender to receiver nodes, and expose it alongside `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`, diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..1bfe9eb5 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,7 +3,12 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.probabilistic import ( + ProbabilisticARForecaster, + ProbabilisticForecaster, +) from .module import ForecasterModule +from .probabilistic_module import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_lam import GraphLAM diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 7ea9f6fd..254c4ba0 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -5,3 +5,4 @@ # Local from .autoregressive import ARForecaster from .base import Forecaster +from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index a8135f62..92a4e938 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,6 +1,7 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" # Standard library +from typing import Callable # Third-party import torch @@ -144,3 +145,75 @@ def forward( pred_std = None return prediction, pred_std + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score the deterministic rollout with the injected scoring rule. + + Unrolls a single forecast over the full rollout, scores it against + the target states on interior nodes and averages over batch and + time. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during the rollout. + Dims: same as the prediction. + score_fn : Callable + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the wrapped predictor does not + output an std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the rollout, averaged over + batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; the deterministic objective has no separate components. + """ + prediction, pred_std = self( + init_states, forcing_features, target_states + ) + if pred_std is None: + pred_std = per_var_std + + batch_loss = torch.mean( + score_fn( + prediction, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4d957916..4b50c030 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,6 +2,7 @@ # Standard library from abc import ABC, abstractmethod +from typing import Callable # Third-party import torch @@ -79,3 +80,66 @@ def forward( per-variable std is substituted upstream by ``ForecasterModule``. Dims: same as ``prediction``. """ + + @abstractmethod + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the training objective for one batch. + + The forecaster owns its complete training objective: which forecasts + to produce from the batch, which loss terms to compute from them and + how to combine those terms into a single scalar. The wrapping + ``ForecasterModule`` only injects the configured scoring rule and + mask, logs the returned components and optimizes the returned loss. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during forecasting. + Dims: same as the prediction. + score_fn : Callable + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the forecaster does not predict its + own std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The full training loss for the batch, to take gradients + of. + loss_components : dict of {str: torch.Tensor} + Scalar loss-related quantities to log alongside the loss, keyed + by component name. The wrapping module prefixes the names with + the training phase. Empty when the objective has no separate + components worth logging. + """ diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..c6609b7d --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,257 @@ +"""Forecasters producing probabilistic (ensemble) forecasts.""" + +# Standard library +from abc import abstractmethod +from typing import Callable + +# Third-party +import torch + +# Local +from ...datastore import BaseDatastore +from ..step_predictors.base import StepPredictor +from .autoregressive import ARForecaster +from .base import Forecaster + + +class ProbabilisticForecaster(Forecaster): + """ + Forecaster whose forecasts are samples from a predictive distribution. + + Adds the capability that probabilistic evaluation and ensemble-based + objectives build on: sampling an ensemble of forecasts. How the + members are produced (auto-regressive sampling, diffusion, ...) is + left to subclasses; consumers only rely on the shape of the returned + ensemble. + """ + + @abstractmethod + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each + predicted step, identically in every member. Dims: same as one + member. + num_members : int or None + Number of ensemble members ``S`` to sample. When ``None``, the + forecaster's configured ensemble size is used. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + ensemble_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` + when the forecaster predicts an std, otherwise ``None``. Dims: + same as ``ensemble``. + """ + + +class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): + """ + Auto-regressive forecaster for step predictors that sample their output. + + Each call to the wrapped predictor draws a fresh sample of the next + state, so the inherited ``ARForecaster.forward`` unrolls one sampled + trajectory. This class adds ensemble forecasting on top: unrolling + several trajectories and stacking them along an ensemble dimension. + The default training objective scores the ensemble mean with the + injected scoring rule; forecasters with model-specific objectives + (ensemble scoring rules, variational objectives) override + ``compute_training_loss``. + """ + + def __init__( + self, + predictor: StepPredictor, + datastore: BaseDatastore, + ensemble_size: int, + ) -> None: + """ + Initialize the ProbabilisticARForecaster. + + Parameters + ---------- + predictor : StepPredictor + The predictor to use for each step. Each call should draw a + fresh sample of the next state. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + ensemble_size : int + Number of ensemble members to sample when no explicit member + count is given, in particular for the training objective. + """ + super().__init__(predictor, datastore) + if ensemble_size < 1: + raise ValueError( + f"ensemble_size must be at least 1, got {ensemble_size}" + ) + self.ensemble_size = ensemble_size + + def sample_ensemble( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_members: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Sample an ensemble of forecasts. + + Unrolls ``num_members`` independent forecasts, each sampling fresh + randomness at every step, and stacks them along a new ensemble + dimension after the batch dimension. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used only to overwrite boundary nodes at each AR + step, identically in every member. Dims: same as one member. + num_members : int or None + Number of ensemble members ``S`` to sample. When ``None``, + ``self.ensemble_size`` is used. + + Returns + ------- + ensemble : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + The sampled forecasts, stacked along the ensemble dimension + ``S``. + ensemble_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` + when the wrapped predictor outputs an std, otherwise ``None``. + Dims: same as ``ensemble``. + """ + if num_members is None: + num_members = self.ensemble_size + + member_list = [] + member_std_list = [] + for _ in range(num_members): + prediction, pred_std = self( + init_states, forcing_features, boundary_states + ) + member_list.append(prediction) + if pred_std is not None: + member_std_list.append(pred_std) + + ensemble = torch.stack(member_list, dim=1) + ensemble_std = ( + torch.stack(member_std_list, dim=1) if member_std_list else None + ) + return ensemble, ensemble_std + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + score_fn: Callable[..., torch.Tensor], + interior_mask_bool: torch.Tensor, + per_var_std: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Score the ensemble mean with the injected scoring rule. + + Samples an ensemble of ``self.ensemble_size`` forecasts, averages + the members into an ensemble mean forecast, scores it against the + target states on interior nodes and averages over batch and time. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used both as the prediction + targets and to overwrite boundary nodes during the rollouts. + Dims: same as one ensemble member. + score_fn : Callable + The configured scoring rule from ``neural_lam.metrics``, called + as ``score_fn(prediction, target, pred_std, mask=...)``. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes are scored. + per_var_std : torch.Tensor or None + Shape ``(num_state_vars,)``. Constant per-variable standard + deviation to score with when the wrapped predictor does not + output an std, otherwise ``None``. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The scoring rule applied to the ensemble mean, averaged + over batch and time. + loss_components : dict of {str: torch.Tensor} + Empty; this objective has no separate components. + """ + ensemble, ensemble_std = self.sample_ensemble( + init_states, forcing_features, target_states + ) + ensemble_mean = ensemble.mean(dim=1) + if ensemble_std is not None: + pred_std = ensemble_std.mean(dim=1) + else: + pred_std = per_var_std + + batch_loss = torch.mean( + score_fn( + ensemble_mean, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 71ce7951..62791733 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -338,7 +338,7 @@ def on_after_batch_transfer(self, batch, dataloader_idx): def common_step(self, batch): """ - Perform a common prediction step for training, validation, and testing. + Perform a common prediction step for validation and testing. Parameters ---------- @@ -362,6 +362,10 @@ def training_step(self, batch): """ Perform a single training step. + The training objective is fully assembled by the wrapped forecaster; + this method injects the configured scoring rule and interior mask, + then logs the loss and any loss components the forecaster returns. + Parameters ---------- batch : tuple @@ -372,20 +376,20 @@ def training_step(self, batch): torch.Tensor The computed loss for the training step. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.per_var_std - - batch_loss = torch.mean( - self.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ) + init_states, target_states, forcing_features, _ = batch + batch_loss, loss_components = self.forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=self.loss, + interior_mask_bool=self.interior_mask_bool, + per_var_std=self.per_var_std, ) - log_dict = {"train_loss": batch_loss} + log_dict = { + f"train_{name}": value for name, value in loss_components.items() + } + log_dict["train_loss"] = batch_loss self.log_dict( log_dict, prog_bar=True, diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py new file mode 100644 index 00000000..24f83049 --- /dev/null +++ b/neural_lam/models/probabilistic_module.py @@ -0,0 +1,140 @@ +"""Lightning module evaluating probabilistic forecasters as ensembles.""" + +# Third-party +import torch + +# Local +from .. import metrics +from .forecasters.probabilistic import ProbabilisticForecaster +from .module import ForecasterModule + + +class ProbabilisticForecasterModule(ForecasterModule): + """ + Lightning module for forecasters that sample ensemble forecasts. + + Training is inherited unchanged from ``ForecasterModule``: the wrapped + forecaster assembles its own training loss. Validation is ensemble + based instead of deterministic: an ensemble is sampled from the + forecaster and scored through its ensemble mean (root-mean-squared + error of the ensemble mean). The module only assumes that the + forecaster can sample ensemble forecasts of the correct shape; it makes + no assumption on how the members are produced. + """ + + # The wrapped forecaster must be able to sample ensemble forecasts + forecaster: ProbabilisticForecaster + + def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): + """ + Initialize the module and store the evaluation ensemble size. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``ForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + eval_ensemble_size : int or None + Number of ensemble members sampled during validation. ``None`` + uses the forecaster's configured ensemble size. + **kwargs + Keyword arguments forwarded to ``ForecasterModule.__init__`` + (``loss``, ``lr``, ...). + """ + super().__init__(*args, **kwargs) + if eval_ensemble_size is not None and eval_ensemble_size < 1: + raise ValueError( + "eval_ensemble_size must be at least 1, " + f"got {eval_ensemble_size}" + ) + self.eval_ensemble_size = eval_ensemble_size + self.val_metrics = {"ens_mse": []} + + def validation_step(self, batch, batch_idx): + """ + Perform a single ensemble validation step. + + Samples an ensemble from the forecaster and scores its ensemble + mean against the target states on interior nodes. Logs the + root-mean-squared error of the ensemble mean per configured rollout + step and averaged over the rollout, and collects per-variable + ensemble-mean MSE for epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + init_states, target_states, forcing_features, _ = batch + ensemble, _ = self.forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=self.eval_ensemble_size, + ) + ensemble_mean = ensemble.mean(dim=1) + # metrics.mse ignores the std argument, but requires one + std_placeholder = torch.ones( + target_states.shape[-1], device=target_states.device + ) + + time_step_mse = torch.mean( + metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + ), + dim=0, + ) + time_step_rmse = torch.sqrt(time_step_mse) + mean_rmse = torch.mean(time_step_rmse) + self._warn_skipped_val_steps(len(time_step_rmse), "val") + + val_log_dict = { + f"val_loss_unroll{step}": time_step_rmse[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_rmse) + } + val_log_dict["val_mean_loss"] = mean_rmse + self.log_dict( + val_log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + entry_mses = metrics.mse( + ensemble_mean, + target_states, + std_placeholder, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.val_metrics["ens_mse"].append(entry_mses) + + def test_step(self, batch, batch_idx): + """ + Not supported: ensemble test evaluation is not implemented. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + + Raises + ------ + NotImplementedError + Always; only training and ensemble validation are implemented + for probabilistic forecasters. + """ + raise NotImplementedError( + "Ensemble test evaluation is not implemented for " + "probabilistic forecasters." + ) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py new file mode 100644 index 00000000..34458c13 --- /dev/null +++ b/tests/test_probabilistic_forecaster.py @@ -0,0 +1,304 @@ +# Third-party +import pytest +import torch +from torch import nn + +# First-party +from neural_lam import config as nlconfig +from neural_lam import metrics +from neural_lam.models import ( + ARForecaster, + ForecasterModule, + ProbabilisticARForecaster, + ProbabilisticForecasterModule, + StepPredictor, +) +from tests.conftest import init_datastore_example + + +class ZeroStepPredictor(StepPredictor): + """Deterministic predictor always predicting the zero state.""" + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = torch.zeros_like(prev_state) + pred_std = torch.zeros_like(prev_state) if self.output_std else None + return pred_state, pred_std + + +class NoisyStepPredictor(StepPredictor): + """Stochastic predictor sampling a fresh state at every call.""" + + def __init__(self, datastore, **kwargs): + super().__init__(datastore, **kwargs) + self.noise_scale = nn.Parameter(torch.tensor(1.0)) + + def forward(self, prev_state, prev_prev_state, forcing): + pred_state = self.noise_scale * torch.randn_like(prev_state) + return pred_state, None + + +def _example_batch(datastore, B=2, pred_steps=3): + """Create constant example input tensors matching the datastore dims.""" + num_grid_nodes = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + num_past_forcing_steps = 1 + num_future_forcing_steps = 1 + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + target_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) * 5.0 + return init_states, forcing_features, target_states + + +def test_ar_forecaster_training_loss_matches_direct_score(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore) + + init_states, forcing_features, target_states = _example_batch(datastore) + score_fn = metrics.get_metric("mse") + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + per_var_std = torch.ones(d_state) + + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=score_fn, + interior_mask_bool=interior_mask_bool, + per_var_std=per_var_std, + ) + + prediction, _ = forecaster(init_states, forcing_features, target_states) + expected_loss = torch.mean( + score_fn( + prediction, + target_states, + per_var_std, + mask=interior_mask_bool, + ) + ) + + assert batch_loss.shape == () + assert loss_components == {} + torch.testing.assert_close(batch_loss, expected_loss) + + +def test_sample_ensemble_shapes_and_member_variability(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + + # Override masks to test boundary masking behaviour + forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) + forecaster.interior_mask[0, 0] = 1 # One node is interior + forecaster.boundary_mask = 1 - forecaster.interior_mask + + B, pred_steps, num_members = 2, 3, 4 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + num_grid_nodes = datastore.num_grid_points + d_state = target_states.shape[-1] + + torch.manual_seed(42) + ensemble, ensemble_std = forecaster.sample_ensemble( + init_states, + forcing_features, + target_states, + num_members=num_members, + ) + + assert ensemble.shape == ( + B, + num_members, + pred_steps, + num_grid_nodes, + d_state, + ) + assert ensemble_std is None + + # Members carry independent samples on the interior node + assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0]) + # Boundary nodes are overwritten with the true state in every member + assert torch.all(ensemble[:, :, :, 1:] == 5.0) + + # Without an explicit member count the configured ensemble_size is used + default_ensemble, _ = forecaster.sample_ensemble( + init_states, forcing_features, target_states + ) + assert default_ensemble.shape[1] == forecaster.ensemble_size + + +def test_probabilistic_training_loss_gradient_flow(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + d_state = target_states.shape[-1] + + torch.manual_seed(42) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=metrics.get_metric("mse"), + interior_mask_bool=interior_mask_bool, + per_var_std=torch.ones(d_state), + ) + + assert batch_loss.shape == () + assert loss_components == {} + assert torch.isfinite(batch_loss) + + batch_loss.backward() + assert predictor.noise_scale.grad is not None + assert predictor.noise_scale.grad != 0.0 + + +def test_probabilistic_forecaster_rejects_empty_ensemble(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(ValueError, match="ensemble_size"): + ProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + + +def test_module_training_step_delegates_to_forecaster(): + datastore = init_datastore_example("mdp") + predictor = ZeroStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + batch = (init_states, target_states, forcing_features, batch_times) + + batch_loss = model.training_step(batch) + + expected_loss, _ = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + score_fn=model.loss, + interior_mask_bool=model.interior_mask_bool, + per_var_std=model.per_var_std, + ) + + torch.testing.assert_close(batch_loss, expected_loss) + + +class MemberCountRecordingForecaster(ProbabilisticARForecaster): + """ProbabilisticARForecaster recording the requested member count.""" + + def sample_ensemble(self, *args, **kwargs): + self.last_num_members = kwargs.get("num_members") + return super().sample_ensemble(*args, **kwargs) + + +def test_probabilistic_module_validation_scores_ensemble_mean(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, ensemble_size=2 + ) + + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + eval_ensemble_size=3, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(42) + model.validation_step(batch, 0) + + # Validation samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.val_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) + + +def test_probabilistic_module_rejects_empty_eval_ensemble(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + + with pytest.raises(ValueError, match="eval_ensemble_size"): + ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + eval_ensemble_size=0, + ) + + +def test_probabilistic_module_test_step_not_implemented(): + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2 + ) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="mse", + ) + + init_states, forcing_features, target_states = _example_batch(datastore) + batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + batch = (init_states, target_states, forcing_features, batch_times) + + with pytest.raises(NotImplementedError): + model.test_step(batch, 0) From 3f5402d9b57fab41394522620cd8c15c53065b24 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Mon, 6 Jul 2026 13:32:10 +0530 Subject: [PATCH 31/51] Address PR review: clarify scoring-rule wording, rename score_fn to score_metric --- neural_lam/models/forecasters/autoregressive.py | 12 ++++++------ neural_lam/models/forecasters/base.py | 8 ++++---- neural_lam/models/forecasters/probabilistic.py | 17 +++++++++-------- neural_lam/models/module.py | 2 +- tests/test_probabilistic_forecaster.py | 10 +++++----- 5 files changed, 25 insertions(+), 24 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 92a4e938..9a2980c6 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -151,12 +151,12 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the deterministic rollout with the injected scoring rule. + Score the deterministic rollout with the given ``score_metric``. Unrolls a single forecast over the full rollout, scores it against the target states on interior nodes and averages over batch and @@ -182,12 +182,12 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollout. Dims: same as the prediction. - score_fn : Callable + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard @@ -209,7 +209,7 @@ def compute_training_loss( pred_std = per_var_std batch_loss = torch.mean( - score_fn( + score_metric( prediction, target_states, pred_std, diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 4b50c030..ad32de00 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -87,7 +87,7 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: @@ -120,12 +120,12 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during forecasting. Dims: same as the prediction. - score_fn : Callable + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index c6609b7d..0965a4a4 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -81,8 +81,9 @@ class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): state, so the inherited ``ARForecaster.forward`` unrolls one sampled trajectory. This class adds ensemble forecasting on top: unrolling several trajectories and stacking them along an ensemble dimension. - The default training objective scores the ensemble mean with the - injected scoring rule; forecasters with model-specific objectives + The default training objective scores the ensemble mean using the + scoring rule passed to ``compute_training_loss`` (from + ``neural_lam.metrics``); forecasters with model-specific objectives (ensemble scoring rules, variational objectives) override ``compute_training_loss``. """ @@ -186,12 +187,12 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_fn: Callable[..., torch.Tensor], + score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with the injected scoring rule. + Score the ensemble mean with the given ``score_metric``. Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the @@ -217,12 +218,12 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollouts. Dims: same as one ensemble member. - score_fn : Callable + score_metric : Callable The configured scoring rule from ``neural_lam.metrics``, called - as ``score_fn(prediction, target, pred_std, mask=...)``. + as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_fn`` so that only interior + nodes; passed as ``mask`` to ``score_metric`` so that only interior nodes are scored. per_var_std : torch.Tensor or None Shape ``(num_state_vars,)``. Constant per-variable standard @@ -247,7 +248,7 @@ def compute_training_loss( pred_std = per_var_std batch_loss = torch.mean( - score_fn( + score_metric( ensemble_mean, target_states, pred_std, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 62791733..a4ca4d8d 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -381,7 +381,7 @@ def training_step(self, batch): init_states, forcing_features, target_states, - score_fn=self.loss, + score_metric=self.loss, interior_mask_bool=self.interior_mask_bool, per_var_std=self.per_var_std, ) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 34458c13..d22c68dc 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -58,7 +58,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): forecaster = ARForecaster(predictor, datastore) init_states, forcing_features, target_states = _example_batch(datastore) - score_fn = metrics.get_metric("mse") + score_metric = metrics.get_metric("mse") interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] per_var_std = torch.ones(d_state) @@ -67,14 +67,14 @@ def test_ar_forecaster_training_loss_matches_direct_score(): init_states, forcing_features, target_states, - score_fn=score_fn, + score_metric=score_metric, interior_mask_bool=interior_mask_bool, per_var_std=per_var_std, ) prediction, _ = forecaster(init_states, forcing_features, target_states) expected_loss = torch.mean( - score_fn( + score_metric( prediction, target_states, per_var_std, @@ -151,7 +151,7 @@ def test_probabilistic_training_loss_gradient_flow(): init_states, forcing_features, target_states, - score_fn=metrics.get_metric("mse"), + score_metric=metrics.get_metric("mse"), interior_mask_bool=interior_mask_bool, per_var_std=torch.ones(d_state), ) @@ -200,7 +200,7 @@ def test_module_training_step_delegates_to_forecaster(): init_states, forcing_features, target_states, - score_fn=model.loss, + score_metric=model.loss, interior_mask_bool=model.interior_mask_bool, per_var_std=model.per_var_std, ) From 987fecd7fc0bc6d5399754a8ffa3fbf0379f5301 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 09:10:48 +0530 Subject: [PATCH 32/51] Address PR review: move loss and per_var_std onto the Forecaster score_metric/per_var_std were injected into compute_training_loss by ForecasterModule and also used directly for val/test loss reporting, duplicating config the forecaster already needs for its own objective. ARForecaster/ProbabilisticARForecaster now own self.loss and self.per_var_std (computed from an optional config ctor arg), and ForecasterModule reads them off self.forecaster instead. Also trims the CHANGELOG entry for #685 down to one sentence per review feedback. --- CHANGELOG.md | 18 +---- .../models/forecasters/autoregressive.py | 64 ++++++++++----- neural_lam/models/forecasters/base.py | 19 ++--- .../models/forecasters/probabilistic.py | 31 +++---- neural_lam/models/module.py | 80 +++++++------------ neural_lam/models/probabilistic_module.py | 2 +- neural_lam/train_model.py | 9 ++- tests/test_checkpoint.py | 3 +- tests/test_datasets.py | 5 +- tests/test_gnn_layers.py | 2 +- tests/test_gpu_normalization.py | 2 +- tests/test_plotting.py | 8 +- tests/test_prediction_model_classes.py | 15 ++-- tests/test_probabilistic_forecaster.py | 25 +++--- tests/test_training.py | 3 +- 15 files changed, 135 insertions(+), 151 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d5ab6ee8..2809daaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,20 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a general probabilistic forecasting interface: an abstract - `ProbabilisticForecaster` capable of sampling ensemble forecasts - (`sample_ensemble`, members stacked along a new dimension after batch), - its auto-regressive implementation `ProbabilisticARForecaster` (samples - independent trajectories through a stochastic step predictor and by - default trains on the configured scoring rule applied to the ensemble - mean) and a `ProbabilisticForecasterModule` whose validation samples an - ensemble and logs the RMSE of the ensemble mean. Move ownership of the - training objective from `ForecasterModule` onto the `Forecaster`: the - new abstract `Forecaster.compute_training_loss` returns a finished - `(loss, loss_components)` pair and `ForecasterModule.training_step` only - injects the configured scoring rule and interior mask and logs the - result. The deterministic `ARForecaster` training loss is unchanged in - value, only computed by the forecaster itself. +- Add a general probabilistic forecasting interface (`ProbabilisticForecaster`, + `ProbabilisticARForecaster`, `ProbabilisticForecasterModule`) and move + ownership of the training objective, scoring rule and per-variable std + from `ForecasterModule` onto the `Forecaster`. [\#685](https://github.com/mllam/neural-lam/issues/685) @Sir-Sloth-The-Lazy diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 9a2980c6..74f2dc77 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,13 +1,13 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" -# Standard library -from typing import Callable - # Third-party import torch # Local +from ... import metrics +from ...config import NeuralLAMConfig from ...datastore import BaseDatastore +from ...loss_weighting import get_state_feature_weighting from ..step_predictors.base import StepPredictor from .base import Forecaster @@ -19,7 +19,11 @@ class ARForecaster(Forecaster): """ def __init__( - self, predictor: StepPredictor, datastore: BaseDatastore + self, + predictor: StepPredictor, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", ) -> None: """ Initialize the ARForecaster. @@ -30,6 +34,14 @@ def __init__( The predictor to use for each step. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Only required for that case; + forecasters used purely for inference can omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. """ super().__init__() self.predictor = predictor @@ -45,6 +57,31 @@ def __init__( "interior_mask", 1.0 - self.boundary_mask, persistent=False ) + self.loss = metrics.get_metric(loss) + + # Store per_var_std here if the predictor does not output its own std + if not self.predicts_std and config is not None: + da_state_stats = datastore.get_standardization_dataarray( + category="state" + ) + state_feature_weights = get_state_feature_weighting( + config=config, datastore=datastore + ) + diff_std = torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ) + feature_weights_t = torch.tensor( + state_feature_weights, dtype=torch.float32 + ) + self.register_buffer( + "per_var_std", + diff_std / torch.sqrt(feature_weights_t), + persistent=False, + ) + else: + self.per_var_std = None + @property def predicts_std(self) -> bool: """ @@ -138,7 +175,7 @@ def forward( prediction = torch.stack(prediction_list, dim=1) # If predictor outputs std, stack it; otherwise return None so - # ForecasterModule can substitute the constant per_var_std + # callers can substitute the constant per_var_std if pred_std_list: pred_std = torch.stack(pred_std_list, dim=1) else: @@ -151,12 +188,10 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the deterministic rollout with the given ``score_metric``. + Score the deterministic rollout with ``self.loss``. Unrolls a single forecast over the full rollout, scores it against the target states on interior nodes and averages over batch and @@ -182,17 +217,10 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollout. Dims: same as the prediction. - score_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the wrapped predictor does not - output an std, otherwise ``None``. Returns ------- @@ -206,10 +234,10 @@ def compute_training_loss( init_states, forcing_features, target_states ) if pred_std is None: - pred_std = per_var_std + pred_std = self.per_var_std batch_loss = torch.mean( - score_metric( + self.loss( prediction, target_states, pred_std, diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index ad32de00..8dfb0002 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,7 +2,6 @@ # Standard library from abc import ABC, abstractmethod -from typing import Callable # Third-party import torch @@ -87,18 +86,17 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ Compute the training objective for one batch. The forecaster owns its complete training objective: which forecasts to produce from the batch, which loss terms to compute from them and - how to combine those terms into a single scalar. The wrapping - ``ForecasterModule`` only injects the configured scoring rule and - mask, logs the returned components and optimizes the returned loss. + how to combine those terms into a single scalar, using its own + ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The + wrapping ``ForecasterModule`` only injects the interior mask, logs + the returned components and optimizes the returned loss. Parameters ---------- @@ -120,17 +118,10 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during forecasting. Dims: same as the prediction. - score_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the forecaster does not predict its - own std, otherwise ``None``. Returns ------- diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 0965a4a4..2bfd0458 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -2,12 +2,12 @@ # Standard library from abc import abstractmethod -from typing import Callable # Third-party import torch # Local +from ...config import NeuralLAMConfig from ...datastore import BaseDatastore from ..step_predictors.base import StepPredictor from .autoregressive import ARForecaster @@ -93,6 +93,8 @@ def __init__( predictor: StepPredictor, datastore: BaseDatastore, ensemble_size: int, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", ) -> None: """ Initialize the ProbabilisticARForecaster. @@ -107,8 +109,16 @@ def __init__( ensemble_size : int Number of ensemble members to sample when no explicit member count is given, in particular for the training objective. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``per_var_std``). Only required for that case; + forecasters used purely for inference can omit it. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used by + ``compute_training_loss`` and stored as ``self.loss``. """ - super().__init__(predictor, datastore) + super().__init__(predictor, datastore, config=config, loss=loss) if ensemble_size < 1: raise ValueError( f"ensemble_size must be at least 1, got {ensemble_size}" @@ -187,12 +197,10 @@ def compute_training_loss( init_states: torch.Tensor, forcing_features: torch.Tensor, target_states: torch.Tensor, - score_metric: Callable[..., torch.Tensor], interior_mask_bool: torch.Tensor, - per_var_std: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with the given ``score_metric``. + Score the ensemble mean with ``self.loss``. Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the @@ -218,17 +226,10 @@ def compute_training_loss( states at each predicted step, used both as the prediction targets and to overwrite boundary nodes during the rollouts. Dims: same as one ensemble member. - score_metric : Callable - The configured scoring rule from ``neural_lam.metrics``, called - as ``score_metric(prediction, target, pred_std, mask=...)``. interior_mask_bool : torch.Tensor Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``score_metric`` so that only interior + nodes; passed as ``mask`` to ``self.loss`` so that only interior nodes are scored. - per_var_std : torch.Tensor or None - Shape ``(num_state_vars,)``. Constant per-variable standard - deviation to score with when the wrapped predictor does not - output an std, otherwise ``None``. Returns ------- @@ -245,10 +246,10 @@ def compute_training_loss( if ensemble_std is not None: pred_std = ensemble_std.mean(dim=1) else: - pred_std = per_var_std + pred_std = self.per_var_std batch_loss = torch.mean( - score_metric( + self.loss( ensemble_mean, target_states, pred_std, diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index a4ca4d8d..4cebe6ca 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -20,7 +20,6 @@ from .. import metrics, vis from ..config import NeuralLAMConfig from ..datastore import BaseDatastore -from ..loss_weighting import get_state_feature_weighting from ..weather_dataset import WeatherDataset from .forecasters.base import Forecaster @@ -38,7 +37,6 @@ def __init__( forecaster: Forecaster, config: NeuralLAMConfig, datastore: BaseDatastore, - loss: str = "wmse", lr: float = 1e-3, restore_opt: bool = False, n_example_pred: int = 1, @@ -54,13 +52,14 @@ def __init__( Parameters ---------- forecaster : Forecaster - The forecaster model to use for predictions. + The forecaster model to use for predictions. Owns the scoring + rule (``forecaster.loss``) and the constant per-variable std + fallback (``forecaster.per_var_std``) used for training and for + validation/test loss reporting here. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore Datastore providing grid metadata and data access. - loss : str, default "wmse" - The loss function to use. lr : float, default 1e-3 Learning rate for the optimizer. restore_opt : bool, default False @@ -79,7 +78,7 @@ def __init__( args : argparse.Namespace, optional Pre-refactor ``ARModel`` checkpoint hyperparameters. When provided, attributes on ``args`` take precedence over the - corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``, + corresponding explicit kwargs (``lr``, ``restore_opt``, ``n_example_pred``, ``create_gif``, ``val_steps_to_log``, ``metrics_watch``, ``var_leads_metrics_watch``) so legacy checkpoints round-trip through ``load_from_checkpoint`` @@ -90,10 +89,9 @@ def __init__( # inside an argparse Namespace under the single key 'args'. When # Lightning calls __init__ during load_from_checkpoint it would # otherwise drop 'args' (not in the new signature) and silently fall - # back to defaults for loss/lr/create_gif/etc. Unpack the namespace - # here so legacy checkpoints round-trip correctly. + # back to defaults for lr/create_gif/etc. Unpack the namespace here + # so legacy checkpoints round-trip correctly. if args is not None: - loss = getattr(args, "loss", loss) lr = getattr(args, "lr", lr) restore_opt = getattr(args, "restore_opt", restore_opt) n_example_pred = getattr(args, "n_example_pred", n_example_pred) @@ -139,29 +137,6 @@ def __init__( persistent=False, ) - # Store per_var_std here if predictor does not output std - if not self.forecaster.predicts_std: - da_state_stats = datastore.get_standardization_dataarray( - category="state" - ) - state_feature_weights = get_state_feature_weighting( - config=config, datastore=datastore - ) - diff_std = torch.tensor( - da_state_stats.state_diff_std_standardized.values, - dtype=torch.float32, - ) - feature_weights_t = torch.tensor( - state_feature_weights, dtype=torch.float32 - ) - self.register_buffer( - "per_var_std", - diff_std / torch.sqrt(feature_weights_t), - persistent=False, - ) - else: - self.per_var_std = None - # Standardization statistics used to normalize each batch on-device in # `on_after_batch_transfer`. WeatherDataset returns unstandardized # data, so state and forcing are normalized here rather than on CPU. @@ -207,9 +182,6 @@ def __init__( self.forcing_mean = None self.forcing_std = None - # Instantiate loss function - self.loss = metrics.get_metric(loss) - self.val_metrics: dict[str, list] = { "mse": [], } @@ -362,9 +334,10 @@ def training_step(self, batch): """ Perform a single training step. - The training objective is fully assembled by the wrapped forecaster; - this method injects the configured scoring rule and interior mask, - then logs the loss and any loss components the forecaster returns. + The training objective is fully assembled by the wrapped forecaster, + which owns its own scoring rule; this method injects the interior + mask, then logs the loss and any loss components the forecaster + returns. Parameters ---------- @@ -381,9 +354,7 @@ def training_step(self, batch): init_states, forcing_features, target_states, - score_metric=self.loss, interior_mask_bool=self.interior_mask_bool, - per_var_std=self.per_var_std, ) log_dict = { @@ -452,10 +423,10 @@ def validation_step(self, batch, batch_idx): """ prediction, target_states, pred_std, _ = self.common_step(batch) if pred_std is None: - pred_std = self.per_var_std + pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.loss( + self.forecaster.loss( prediction, target_states, pred_std, @@ -532,10 +503,10 @@ def test_step(self, batch, batch_idx): self.test_metrics["output_std"].append(mean_pred_std) if pred_std is None: - pred_std = self.per_var_std + pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.loss( + self.forecaster.loss( prediction, target_states, pred_std, @@ -572,7 +543,7 @@ def test_step(self, batch, batch_idx): ) self.test_metrics[metric_name].append(batch_metric_vals) - spatial_loss = self.loss( + spatial_loss = self.forecaster.loss( prediction, target_states, pred_std, average_grid=False ) log_spatial_losses = spatial_loss[ @@ -980,15 +951,20 @@ def on_load_checkpoint(self, checkpoint): # 1. Broad namespace remap: for pre-refactor checkpoints # The old ``ARModel`` was a flat LightningModule. Everything that # belonged to the predictor needs to be moved to - # 'forecaster.predictor.' + # 'forecaster.predictor.', while 'per_var_std' (now owned by the + # forecaster itself) moves to 'forecaster.per_var_std' and + # 'interior_mask_bool' (still owned by the module) stays as-is. old_keys = list(loaded_state_dict.keys()) for key in old_keys: - if not key.startswith("forecaster.") and key not in ( - "interior_mask_bool", - "per_var_std", - ): - new_key = f"forecaster.predictor.{key}" - loaded_state_dict[new_key] = loaded_state_dict.pop(key) + if key.startswith("forecaster.") or key == "interior_mask_bool": + continue + if key == "per_var_std": + loaded_state_dict["forecaster.per_var_std"] = ( + loaded_state_dict.pop(key) + ) + continue + new_key = f"forecaster.predictor.{key}" + loaded_state_dict[new_key] = loaded_state_dict.pop(key) # 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp # Will be under forecaster.predictor due to the remap above, or diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 24f83049..37a756d7 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -40,7 +40,7 @@ def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): uses the forecaster's configured ensemble size. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` - (``loss``, ``lr``, ...). + (``lr``, ...). """ super().__init__(*args, **kwargs) if eval_ensemble_size is not None and eval_ensemble_size < 1: diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index f98065c4..d5e86536 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -63,7 +63,9 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) return ForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, @@ -457,13 +459,14 @@ def main(input_args=None): mesh_up_gnn_type=args.mesh_up_gnn_type, mesh_down_gnn_type=args.mesh_down_gnn_type, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster( + predictor, datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, lr=args.lr, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 2e5f3148..6f114043 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -50,12 +50,11 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, n_example_pred=1, val_steps_to_log=[1], diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 4b35840e..1941206b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -235,13 +235,14 @@ def _create_graph(): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py index 04c99003..789297d0 100644 --- a/tests/test_gnn_layers.py +++ b/tests/test_gnn_layers.py @@ -73,7 +73,7 @@ def _build_model_and_data( output_clamping_upper=config.training.output_clamping.upper, **gnn_kwargs, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config) B = 2 num_grid_nodes = predictor.num_grid_nodes diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index 8d516bfb..b063d626 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -26,7 +26,7 @@ def _build_module(datastore): ) ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config) return ForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 616d563d..970590be 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -467,13 +467,14 @@ class ModelArgs: output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore=datastore) + forecaster = ARForecaster( + predictor, datastore=datastore, config=config, loss=args.loss + ) model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=args.loss, restore_opt=args.restore_opt, n_example_pred=args.n_example_pred, val_steps_to_log=args.val_steps_to_log, @@ -679,12 +680,11 @@ def _build_metrics_watch_module(datastore, config): output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") return ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 73e2f905..6abc786c 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -97,13 +97,12 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1e-3, restore_opt=False, n_example_pred=1, @@ -193,8 +192,6 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - forecaster = ARForecaster(predictor, datastore) - # Use distinctive non-default values so we can detect silent fallback # to ForecasterModule's defaults during load. saved_loss = "mse" @@ -203,11 +200,14 @@ def test_forecaster_module_old_checkpoint(tmp_path): saved_val_steps = [2] saved_n_example_pred = 7 + forecaster = ARForecaster( + predictor, datastore, config=config, loss=saved_loss + ) + model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss=saved_loss, lr=saved_lr, restore_opt=False, n_example_pred=saved_n_example_pred, @@ -269,7 +269,9 @@ def test_forecaster_module_old_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss=saved_loss + ) # Load from hacked old checkpoint loaded_model = ForecasterModule.load_from_checkpoint( @@ -284,7 +286,6 @@ def test_forecaster_module_old_checkpoint(tmp_path): # Hyperparameters nested in the legacy 'args' namespace must round-trip # rather than silently falling back to ForecasterModule defaults. - assert loaded_model.hparams.loss == saved_loss assert loaded_model.hparams.lr == saved_lr assert loaded_model.hparams.val_steps_to_log == saved_val_steps assert loaded_model.create_gif is saved_create_gif diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index d22c68dc..1a8f9971 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -55,21 +55,21 @@ def _example_batch(datastore, B=2, pred_steps=3): def test_ar_forecaster_training_loss_matches_direct_score(): datastore = init_datastore_example("mdp") predictor = ZeroStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, loss="mse") init_states, forcing_features, target_states = _example_batch(datastore) score_metric = metrics.get_metric("mse") interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] - per_var_std = torch.ones(d_state) + # per_var_std is normally computed from config; override directly since + # this test only cares about the loss computation, not standardization. + forecaster.per_var_std = torch.ones(d_state) batch_loss, loss_components = forecaster.compute_training_loss( init_states, forcing_features, target_states, - score_metric=score_metric, interior_mask_bool=interior_mask_bool, - per_var_std=per_var_std, ) prediction, _ = forecaster(init_states, forcing_features, target_states) @@ -77,7 +77,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): score_metric( prediction, target_states, - per_var_std, + forecaster.per_var_std, mask=interior_mask_bool, ) ) @@ -139,21 +139,22 @@ def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 + predictor, datastore, ensemble_size=2, loss="mse" ) init_states, forcing_features, target_states = _example_batch(datastore) interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] + # per_var_std is normally computed from config; override directly since + # this test only cares about the loss computation, not standardization. + forecaster.per_var_std = torch.ones(d_state) torch.manual_seed(42) batch_loss, loss_components = forecaster.compute_training_loss( init_states, forcing_features, target_states, - score_metric=metrics.get_metric("mse"), interior_mask_bool=interior_mask_bool, - per_var_std=torch.ones(d_state), ) assert batch_loss.shape == () @@ -176,18 +177,17 @@ def test_probabilistic_forecaster_rejects_empty_ensemble(): def test_module_training_step_delegates_to_forecaster(): datastore = init_datastore_example("mdp") predictor = ZeroStepPredictor(datastore=datastore, output_std=False) - forecaster = ARForecaster(predictor, datastore) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", ) init_states, forcing_features, target_states = _example_batch(datastore) @@ -200,9 +200,7 @@ def test_module_training_step_delegates_to_forecaster(): init_states, forcing_features, target_states, - score_metric=model.loss, interior_mask_bool=model.interior_mask_bool, - per_var_std=model.per_var_std, ) torch.testing.assert_close(batch_loss, expected_loss) @@ -232,7 +230,6 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", eval_ensemble_size=3, ) @@ -273,7 +270,6 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", eval_ensemble_size=0, ) @@ -293,7 +289,6 @@ def test_probabilistic_module_test_step_not_implemented(): forecaster=forecaster, config=config, datastore=datastore, - loss="mse", ) init_states, forcing_features, target_states = _example_batch(datastore) diff --git a/tests/test_training.py b/tests/test_training.py index bf1a5884..589e9d89 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -123,13 +123,12 @@ def run_simple_training( output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, ) - forecaster = ARForecaster(predictor, datastore) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") model = ForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - loss="mse", lr=1.0e-3, restore_opt=False, n_example_pred=1, From 511a6d5b493f941230bc1303d859ba38c97ff71f Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 09:47:21 +0530 Subject: [PATCH 33/51] Fail fast when a forecaster is missing per_var_std it needs A Forecaster built without config now silently has per_var_std=None when its predictor doesn't output its own std. Previously per_var_std was always computed by ForecasterModule itself, so this gap didn't exist; now that construction is split across two calls, catch it at ForecasterModule init instead of crashing at the first val/test step. --- neural_lam/models/module.py | 8 ++++++++ tests/test_prediction_model_classes.py | 4 +++- tests/test_probabilistic_forecaster.py | 18 +++++++++--------- 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 4cebe6ca..e7420b52 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -122,6 +122,14 @@ def __init__( self.save_hyperparameters(ignore=["datastore", "forecaster"]) self.datastore = datastore self.forecaster = forecaster + if forecaster.per_var_std is None and not forecaster.predicts_std: + raise ValueError( + "forecaster.per_var_std is None but the forecaster does " + "not predict its own std (forecaster.predicts_std is " + "False), so training/validation/test scoring has no std " + "to use. Pass config to the forecaster's constructor so " + "it can compute the constant per-variable std." + ) self.matched_metrics: set = set() # Compute interior_mask_bool directly from datastore diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 6abc786c..9bc9d9c0 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -132,7 +132,9 @@ def test_forecaster_module_checkpoint(tmp_path): num_future_forcing_steps=1, output_std=False, ) - load_forecaster = ARForecaster(load_predictor, datastore) + load_forecaster = ARForecaster( + load_predictor, datastore, config=config, loss="mse" + ) # Load from checkpoint loaded_model = ForecasterModule.load_from_checkpoint( diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 1a8f9971..515b9232 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -217,15 +217,15 @@ def sample_ensemble(self, *args, **kwargs): def test_probabilistic_module_validation_scores_ensemble_mean(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = MemberCountRecordingForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = MemberCountRecordingForecaster( + predictor, datastore, ensemble_size=2, config=config + ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, @@ -256,14 +256,14 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): def test_probabilistic_module_rejects_empty_eval_ensemble(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2, config=config + ) with pytest.raises(ValueError, match="eval_ensemble_size"): ProbabilisticForecasterModule( @@ -277,14 +277,14 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): def test_probabilistic_module_test_step_not_implemented(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) + forecaster = ProbabilisticARForecaster( + predictor, datastore, ensemble_size=2, config=config + ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, From 56b3d6bd97f5a93b3d2627eefefada529581e644 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 10:47:31 +0530 Subject: [PATCH 34/51] Rename ensemble_std to per_member_std, document mixture semantics Each member's predicted std is its own, not a spread computed across the ensemble, so ensemble_std was a misleading name. Document on ProbabilisticForecaster that a per-member std makes the predictive distribution a mixture of Gaussians, and note in ProbabilisticARForecaster.compute_training_loss that averaging the per-member stds is a simplification of the true mixture variance (which also includes the spread between member means). --- .../models/forecasters/probabilistic.py | 43 +++++++++++++------ tests/test_probabilistic_forecaster.py | 4 +- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 2bfd0458..cfa6169b 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -23,6 +23,14 @@ class ProbabilisticForecaster(Forecaster): members are produced (auto-regressive sampling, diffusion, ...) is left to subclasses; consumers only rely on the shape of the returned ensemble. + + When ``sample_ensemble`` returns a ``per_member_std``, it is each + member's own predicted std, not a std describing the spread across + members. The predictive distribution is then a mixture of ``S`` + Gaussians, one per member: ``p(x) = mean_s N(x; ensemble[:, s], + per_member_std[:, s]**2)``, not a single Gaussian. In particular, the + variance of that mixture is not the average of the per-member + variances: it also includes the spread between the member means. """ @abstractmethod @@ -66,10 +74,12 @@ def sample_ensemble( Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. The sampled forecasts, stacked along the ensemble dimension ``S``. - ensemble_std : torch.Tensor or None - Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` - when the forecaster predicts an std, otherwise ``None``. Dims: - same as ``ensemble``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the forecaster predicts an std, otherwise + ``None``. Dims: same as ``ensemble``. """ @@ -168,10 +178,12 @@ def sample_ensemble( Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. The sampled forecasts, stacked along the ensemble dimension ``S``. - ensemble_std : torch.Tensor or None - Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)`` - when the wrapped predictor outputs an std, otherwise ``None``. - Dims: same as ``ensemble``. + per_member_std : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. + Each member's own predicted std (see the class docstring for + why the ensemble is then a mixture, not this averaged with the + others), when the wrapped predictor outputs an std, otherwise + ``None``. Dims: same as ``ensemble``. """ if num_members is None: num_members = self.ensemble_size @@ -187,10 +199,10 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) - ensemble_std = ( + per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) - return ensemble, ensemble_std + return ensemble, per_member_std def compute_training_loss( self, @@ -205,6 +217,11 @@ def compute_training_loss( Samples an ensemble of ``self.ensemble_size`` forecasts, averages the members into an ensemble mean forecast, scores it against the target states on interior nodes and averages over batch and time. + When members predict their own std, the std passed to ``self.loss`` + is the plain average of the per-member stds; this is a + simplification of the true mixture predictive variance, which + would also include the spread between the member means (see the + ``ProbabilisticForecaster`` class docstring). Parameters ---------- @@ -239,12 +256,12 @@ def compute_training_loss( loss_components : dict of {str: torch.Tensor} Empty; this objective has no separate components. """ - ensemble, ensemble_std = self.sample_ensemble( + ensemble, per_member_std = self.sample_ensemble( init_states, forcing_features, target_states ) ensemble_mean = ensemble.mean(dim=1) - if ensemble_std is not None: - pred_std = ensemble_std.mean(dim=1) + if per_member_std is not None: + pred_std = per_member_std.mean(dim=1) else: pred_std = self.per_var_std diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 515b9232..854f8bff 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -107,7 +107,7 @@ def test_sample_ensemble_shapes_and_member_variability(): d_state = target_states.shape[-1] torch.manual_seed(42) - ensemble, ensemble_std = forecaster.sample_ensemble( + ensemble, per_member_std = forecaster.sample_ensemble( init_states, forcing_features, target_states, @@ -121,7 +121,7 @@ def test_sample_ensemble_shapes_and_member_variability(): num_grid_nodes, d_state, ) - assert ensemble_std is None + assert per_member_std is None # Members carry independent samples on the interior node assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0]) From d64878cfd3d9fbc5c37b3429d77f811724d2d771 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 10:53:28 +0530 Subject: [PATCH 35/51] Update neural_lam/models/forecasters/probabilistic.py Co-authored-by: Joel Oskarsson --- neural_lam/models/forecasters/probabilistic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index cfa6169b..560ae599 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -199,6 +199,7 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) + # After stacking shape of ensemble is (B, S, pred_steps, num_grid_nodes, num_state_vars) per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) From 6fd050a18adda3d9e2cfa7fc3276229a59d7d89b Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:10:06 +0530 Subject: [PATCH 36/51] Leave ProbabilisticARForecaster.compute_training_loss abstract Scoring the ensemble mean with a pointwise metric only rewards the mean being right, giving the model no incentive to keep a calibrated spread, and risks training it to collapse to a point estimate. Redeclare compute_training_loss as abstract on ProbabilisticARForecaster instead of providing that as a default (it would otherwise silently fall back to ARForecaster's single-rollout objective via MRO, not even the ensemble mean). Concrete subclasses must define their own objective. Tests that only need an instantiable forecaster now use a local ConcreteProbabilisticARForecaster example (ensemble-mean scoring, moved out of the library code); a new test locks in that the base class itself cannot be instantiated. --- .../models/forecasters/probabilistic.py | 86 +++++-------------- tests/test_probabilistic_forecaster.py | 60 +++++++++++-- 2 files changed, 74 insertions(+), 72 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 560ae599..ae5a15cb 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -91,11 +91,17 @@ class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): state, so the inherited ``ARForecaster.forward`` unrolls one sampled trajectory. This class adds ensemble forecasting on top: unrolling several trajectories and stacking them along an ensemble dimension. - The default training objective scores the ensemble mean using the - scoring rule passed to ``compute_training_loss`` (from - ``neural_lam.metrics``); forecasters with model-specific objectives - (ensemble scoring rules, variational objectives) override - ``compute_training_loss``. + + ``compute_training_loss`` is intentionally left abstract here (it does + not fall back to ``ARForecaster``'s single-rollout objective, which + would silently train on one stochastic sample). There is no default + objective that fits every stochastic model: scoring the ensemble mean + with a pointwise metric only rewards the mean being right, giving the + model no incentive to keep a calibrated spread, and risks training it + to collapse the ensemble to a point estimate. Concrete subclasses must + define an objective appropriate to how they are meant to be trained + (e.g. an ensemble scoring rule such as CRPS, or a variational + objective). """ def __init__( @@ -199,12 +205,14 @@ def sample_ensemble( member_std_list.append(pred_std) ensemble = torch.stack(member_list, dim=1) - # After stacking shape of ensemble is (B, S, pred_steps, num_grid_nodes, num_state_vars) + # After stacking, ensemble has shape + # (B, S, pred_steps, num_grid_nodes, num_state_vars) per_member_std = ( torch.stack(member_std_list, dim=1) if member_std_list else None ) return ensemble, per_member_std + @abstractmethod def compute_training_loss( self, init_states: torch.Tensor, @@ -213,65 +221,11 @@ def compute_training_loss( interior_mask_bool: torch.Tensor, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Score the ensemble mean with ``self.loss``. - - Samples an ensemble of ``self.ensemble_size`` forecasts, averages - the members into an ensemble mean forecast, scores it against the - target states on interior nodes and averages over batch and time. - When members predict their own std, the std passed to ``self.loss`` - is the plain average of the per-member stds; this is a - simplification of the true mixture predictive variance, which - would also include the spread between the member means (see the - ``ProbabilisticForecaster`` class docstring). - - Parameters - ---------- - init_states : torch.Tensor - Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial - states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims: - ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), - ``num_grid_nodes`` is the number of spatial nodes, and - ``num_state_vars`` is the state feature dimension. - forcing_features : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. - External forcings provided at each predicted step. Dims: ``B`` - is batch size, ``pred_steps`` is the rollout length, - ``num_grid_nodes`` is the number of spatial nodes, and - ``num_forcing_vars`` is the forcing feature dimension (already - concatenated past/current/future windows). - target_states : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True - states at each predicted step, used both as the prediction - targets and to overwrite boundary nodes during the rollouts. - Dims: same as one ensemble member. - interior_mask_bool : torch.Tensor - Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior - nodes; passed as ``mask`` to ``self.loss`` so that only interior - nodes are scored. + Compute the training objective for one batch. - Returns - ------- - batch_loss : torch.Tensor - Scalar. The scoring rule applied to the ensemble mean, averaged - over batch and time. - loss_components : dict of {str: torch.Tensor} - Empty; this objective has no separate components. + Left abstract; see the class docstring for why there is no default + objective. Concrete subclasses typically call ``sample_ensemble`` + and score the resulting members with an objective appropriate to + the model (see ``Forecaster.compute_training_loss`` for the + signature and general contract). """ - ensemble, per_member_std = self.sample_ensemble( - init_states, forcing_features, target_states - ) - ensemble_mean = ensemble.mean(dim=1) - if per_member_std is not None: - pred_std = per_member_std.mean(dim=1) - else: - pred_std = self.per_var_std - - batch_loss = torch.mean( - self.loss( - ensemble_mean, - target_states, - pred_std, - mask=interior_mask_bool, - ) - ) - return batch_loss, {} diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 854f8bff..f7cfe7de 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -37,6 +37,43 @@ def forward(self, prev_state, prev_prev_state, forcing): return pred_state, None +class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster): + """ + Test-only concrete ``ProbabilisticARForecaster``. + + ``ProbabilisticARForecaster`` leaves ``compute_training_loss`` abstract + (no single default objective fits every stochastic model), so tests + that only need a working forecaster to instantiate use this example + ensemble-mean objective rather than the base class directly. + """ + + def compute_training_loss( + self, + init_states, + forcing_features, + target_states, + interior_mask_bool, + ): + ensemble, per_member_std = self.sample_ensemble( + init_states, forcing_features, target_states + ) + ensemble_mean = ensemble.mean(dim=1) + pred_std = ( + per_member_std.mean(dim=1) + if per_member_std is not None + else self.per_var_std + ) + batch_loss = torch.mean( + self.loss( + ensemble_mean, + target_states, + pred_std, + mask=interior_mask_bool, + ) + ) + return batch_loss, {} + + def _example_batch(datastore, B=2, pred_steps=3): """Create constant example input tensors matching the datastore dims.""" num_grid_nodes = datastore.num_grid_points @@ -90,7 +127,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): def test_sample_ensemble_shapes_and_member_variability(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2 ) @@ -138,7 +175,7 @@ def test_sample_ensemble_shapes_and_member_variability(): def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, loss="mse" ) @@ -171,7 +208,18 @@ def test_probabilistic_forecaster_rejects_empty_ensemble(): predictor = NoisyStepPredictor(datastore=datastore, output_std=False) with pytest.raises(ValueError, match="ensemble_size"): - ProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + ConcreteProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + + +def test_probabilistic_ar_forecaster_is_abstract(): + """ProbabilisticARForecaster leaves compute_training_loss abstract, so + it cannot be instantiated directly; only a subclass that defines an + objective can.""" + datastore = init_datastore_example("mdp") + predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + + with pytest.raises(TypeError, match="abstract"): + ProbabilisticARForecaster(predictor, datastore, ensemble_size=2) def test_module_training_step_delegates_to_forecaster(): @@ -206,7 +254,7 @@ def test_module_training_step_delegates_to_forecaster(): torch.testing.assert_close(batch_loss, expected_loss) -class MemberCountRecordingForecaster(ProbabilisticARForecaster): +class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster): """ProbabilisticARForecaster recording the requested member count.""" def sample_ensemble(self, *args, **kwargs): @@ -261,7 +309,7 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, config=config ) @@ -282,7 +330,7 @@ def test_probabilistic_module_test_step_not_implemented(): kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ProbabilisticARForecaster( + forecaster = ConcreteProbabilisticARForecaster( predictor, datastore, ensemble_size=2, config=config ) model = ProbabilisticForecasterModule( From a1350ff21848b0cf6afbb35f94019d091c8a5899 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:18:19 +0530 Subject: [PATCH 37/51] Require an explicit member count instead of a stored ensemble_size Drop ProbabilisticARForecaster's ensemble_size constructor arg and the implicit num_members=None -> self.ensemble_size fallback in sample_ensemble; num_members is now always required. Baking a default member count into the forecaster's state was unnecessary now that compute_training_loss is abstract too (nothing in the shared base class path used it) and just adds an implicit default callers could silently rely on instead of deciding explicitly. The num_members < 1 validation moves from __init__ to sample_ensemble accordingly. ProbabilisticForecasterModule.eval_ensemble_size follows suit: it no longer defaults to None with a forecaster fallback, it's required. Test-only ConcreteProbabilisticARForecaster (used wherever a concrete probabilistic forecaster is needed for testing) gains its own train_num_members for the training objective, since deciding how many members to sample during training is now the concrete subclass's call. --- .../models/forecasters/probabilistic.py | 34 +++++++-------- neural_lam/models/probabilistic_module.py | 9 ++-- tests/test_probabilistic_forecaster.py | 42 +++++++++++-------- 3 files changed, 43 insertions(+), 42 deletions(-) diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index ae5a15cb..1c5aeeb4 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -39,7 +39,7 @@ def sample_ensemble( init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - num_members: int | None = None, + num_members: int, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Sample an ensemble of forecasts. @@ -64,9 +64,8 @@ def sample_ensemble( state values used only to overwrite boundary nodes at each predicted step, identically in every member. Dims: same as one member. - num_members : int or None - Number of ensemble members ``S`` to sample. When ``None``, the - forecaster's configured ensemble size is used. + num_members : int + Number of ensemble members ``S`` to sample. Returns ------- @@ -108,7 +107,6 @@ def __init__( self, predictor: StepPredictor, datastore: BaseDatastore, - ensemble_size: int, config: NeuralLAMConfig | None = None, loss: str = "wmse", ) -> None: @@ -122,9 +120,6 @@ def __init__( fresh sample of the next state. datastore : BaseDatastore The datastore providing grid metadata and boundary masks. - ensemble_size : int - Number of ensemble members to sample when no explicit member - count is given, in particular for the training objective. config : NeuralLAMConfig or None Configuration used to compute the constant per-variable std substituted for ``pred_std`` when ``predictor`` does not output @@ -135,18 +130,13 @@ def __init__( ``compute_training_loss`` and stored as ``self.loss``. """ super().__init__(predictor, datastore, config=config, loss=loss) - if ensemble_size < 1: - raise ValueError( - f"ensemble_size must be at least 1, got {ensemble_size}" - ) - self.ensemble_size = ensemble_size def sample_ensemble( self, init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - num_members: int | None = None, + num_members: int, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Sample an ensemble of forecasts. @@ -174,9 +164,8 @@ def sample_ensemble( Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state values used only to overwrite boundary nodes at each AR step, identically in every member. Dims: same as one member. - num_members : int or None - Number of ensemble members ``S`` to sample. When ``None``, - ``self.ensemble_size`` is used. + num_members : int + Number of ensemble members ``S`` to sample. Returns ------- @@ -190,9 +179,16 @@ def sample_ensemble( why the ensemble is then a mixture, not this averaged with the others), when the wrapped predictor outputs an std, otherwise ``None``. Dims: same as ``ensemble``. + + Raises + ------ + ValueError + If ``num_members`` is less than 1. """ - if num_members is None: - num_members = self.ensemble_size + if num_members < 1: + raise ValueError( + f"num_members must be at least 1, got {num_members}" + ) member_list = [] member_std_list = [] diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 37a756d7..6bda75cd 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -25,7 +25,7 @@ class ProbabilisticForecasterModule(ForecasterModule): # The wrapped forecaster must be able to sample ensemble forecasts forecaster: ProbabilisticForecaster - def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): + def __init__(self, *args, eval_ensemble_size: int, **kwargs): """ Initialize the module and store the evaluation ensemble size. @@ -35,15 +35,14 @@ def __init__(self, *args, eval_ensemble_size: int | None = None, **kwargs): Positional arguments forwarded to ``ForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). - eval_ensemble_size : int or None - Number of ensemble members sampled during validation. ``None`` - uses the forecaster's configured ensemble size. + eval_ensemble_size : int + Number of ensemble members sampled during validation. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` (``lr``, ...). """ super().__init__(*args, **kwargs) - if eval_ensemble_size is not None and eval_ensemble_size < 1: + if eval_ensemble_size < 1: raise ValueError( "eval_ensemble_size must be at least 1, " f"got {eval_ensemble_size}" diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index f7cfe7de..28006671 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -45,8 +45,14 @@ class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster): (no single default objective fits every stochastic model), so tests that only need a working forecaster to instantiate use this example ensemble-mean objective rather than the base class directly. + ``sample_ensemble`` always requires an explicit member count, so this + class takes its own ``train_num_members`` for the training objective. """ + def __init__(self, *args, train_num_members: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.train_num_members = train_num_members + def compute_training_loss( self, init_states, @@ -55,7 +61,10 @@ def compute_training_loss( interior_mask_bool, ): ensemble, per_member_std = self.sample_ensemble( - init_states, forcing_features, target_states + init_states, + forcing_features, + target_states, + num_members=self.train_num_members, ) ensemble_mean = ensemble.mean(dim=1) pred_std = ( @@ -127,9 +136,7 @@ def test_ar_forecaster_training_loss_matches_direct_score(): def test_sample_ensemble_shapes_and_member_variability(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) - forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2 - ) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) # Override masks to test boundary masking behaviour forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask) @@ -165,18 +172,12 @@ def test_sample_ensemble_shapes_and_member_variability(): # Boundary nodes are overwritten with the true state in every member assert torch.all(ensemble[:, :, :, 1:] == 5.0) - # Without an explicit member count the configured ensemble_size is used - default_ensemble, _ = forecaster.sample_ensemble( - init_states, forcing_features, target_states - ) - assert default_ensemble.shape[1] == forecaster.ensemble_size - def test_probabilistic_training_loss_gradient_flow(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, loss="mse" + predictor, datastore, loss="mse", train_num_members=2 ) init_states, forcing_features, target_states = _example_batch(datastore) @@ -203,12 +204,16 @@ def test_probabilistic_training_loss_gradient_flow(): assert predictor.noise_scale.grad != 0.0 -def test_probabilistic_forecaster_rejects_empty_ensemble(): +def test_sample_ensemble_rejects_empty_member_count(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + forecaster = ConcreteProbabilisticARForecaster(predictor, datastore) + init_states, forcing_features, target_states = _example_batch(datastore) - with pytest.raises(ValueError, match="ensemble_size"): - ConcreteProbabilisticARForecaster(predictor, datastore, ensemble_size=0) + with pytest.raises(ValueError, match="num_members"): + forecaster.sample_ensemble( + init_states, forcing_features, target_states, num_members=0 + ) def test_probabilistic_ar_forecaster_is_abstract(): @@ -219,7 +224,7 @@ def test_probabilistic_ar_forecaster_is_abstract(): predictor = NoisyStepPredictor(datastore=datastore, output_std=False) with pytest.raises(TypeError, match="abstract"): - ProbabilisticARForecaster(predictor, datastore, ensemble_size=2) + ProbabilisticARForecaster(predictor, datastore) def test_module_training_step_delegates_to_forecaster(): @@ -272,7 +277,7 @@ def test_probabilistic_module_validation_scores_ensemble_mean(): ) ) forecaster = MemberCountRecordingForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, @@ -310,7 +315,7 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): ) ) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) with pytest.raises(ValueError, match="eval_ensemble_size"): @@ -331,12 +336,13 @@ def test_probabilistic_module_test_step_not_implemented(): ) ) forecaster = ConcreteProbabilisticARForecaster( - predictor, datastore, ensemble_size=2, config=config + predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, datastore=datastore, + eval_ensemble_size=2, ) init_states, forcing_features, target_states = _example_batch(datastore) From 1cbded635ac606a4a4cab8c7a23de739198e5840 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 8 Jul 2026 11:26:10 +0530 Subject: [PATCH 38/51] Implement ProbabilisticForecasterModule.test_step Mirrors validation_step: samples eval_ensemble_size members and scores the ensemble mean, same as validation. Factored the shared sampling + scoring + logging into _ensemble_step(batch, phase) rather than duplicating the block, since validation_step and test_step differ only in their log-key prefix and which metrics dict collects the result. Overrides on_test_epoch_end (rather than inheriting ForecasterModule's) since this module's test_step doesn't populate spatial_loss_maps or plot examples - the inherited version would crash on torch.cat of an empty list. --- neural_lam/models/probabilistic_module.py | 108 ++++++++++++++++------ tests/test_probabilistic_forecaster.py | 27 ++++-- 2 files changed, 99 insertions(+), 36 deletions(-) diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 6bda75cd..453b00c6 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -1,5 +1,8 @@ """Lightning module evaluating probabilistic forecasters as ensembles.""" +# Standard library +import warnings + # Third-party import torch @@ -14,9 +17,9 @@ class ProbabilisticForecasterModule(ForecasterModule): Lightning module for forecasters that sample ensemble forecasts. Training is inherited unchanged from ``ForecasterModule``: the wrapped - forecaster assembles its own training loss. Validation is ensemble - based instead of deterministic: an ensemble is sampled from the - forecaster and scored through its ensemble mean (root-mean-squared + forecaster assembles its own training loss. Validation and testing are + ensemble based instead of deterministic: an ensemble is sampled from + the forecaster and scored through its ensemble mean (root-mean-squared error of the ensemble mean). The module only assumes that the forecaster can sample ensemble forecasts of the correct shape; it makes no assumption on how the members are produced. @@ -36,7 +39,8 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ``ForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). eval_ensemble_size : int - Number of ensemble members sampled during validation. + Number of ensemble members sampled during validation and + testing. **kwargs Keyword arguments forwarded to ``ForecasterModule.__init__`` (``lr``, ...). @@ -49,23 +53,30 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ) self.eval_ensemble_size = eval_ensemble_size self.val_metrics = {"ens_mse": []} + self.test_metrics = {"ens_mse": []} - def validation_step(self, batch, batch_idx): + def _ensemble_step(self, batch, phase: str): """ - Perform a single ensemble validation step. + Sample an ensemble and score its mean against the target states. - Samples an ensemble from the forecaster and scores its ensemble - mean against the target states on interior nodes. Logs the - root-mean-squared error of the ensemble mean per configured rollout - step and averaged over the rollout, and collects per-variable - ensemble-mean MSE for epoch-end aggregation. + Shared by ``validation_step`` and ``test_step``: samples + ``self.eval_ensemble_size`` members, scores the ensemble mean with + plain (unweighted) MSE on interior nodes, logs the root-mean-squared + error per configured rollout step and averaged over the rollout + under the given phase's prefix. Parameters ---------- batch : tuple The batch of data. - batch_idx : int - The index of the batch. + phase : str + Logging phase, either ``"val"`` or ``"test"``. + + Returns + ------- + torch.Tensor + Per-variable ensemble-mean MSE, shape + ``(B, pred_steps, num_state_vars)``, for epoch-end aggregation. """ init_states, target_states, forcing_features, _ = batch ensemble, _ = self.forecaster.sample_ensemble( @@ -91,34 +102,55 @@ def validation_step(self, batch, batch_idx): ) time_step_rmse = torch.sqrt(time_step_mse) mean_rmse = torch.mean(time_step_rmse) - self._warn_skipped_val_steps(len(time_step_rmse), "val") + self._warn_skipped_val_steps(len(time_step_rmse), phase) - val_log_dict = { - f"val_loss_unroll{step}": time_step_rmse[step - 1] + log_dict = { + f"{phase}_loss_unroll{step}": time_step_rmse[step - 1] for step in self.hparams.val_steps_to_log if step <= len(time_step_rmse) } - val_log_dict["val_mean_loss"] = mean_rmse + log_dict[f"{phase}_mean_loss"] = mean_rmse self.log_dict( - val_log_dict, + log_dict, on_step=False, on_epoch=True, sync_dist=True, batch_size=batch[0].shape[0], ) - entry_mses = metrics.mse( + return metrics.mse( ensemble_mean, target_states, std_placeholder, mask=self.interior_mask_bool, sum_vars=False, ) + + def validation_step(self, batch, batch_idx): + """ + Perform a single ensemble validation step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "val") self.val_metrics["ens_mse"].append(entry_mses) def test_step(self, batch, batch_idx): """ - Not supported: ensemble test evaluation is not implemented. + Perform a single ensemble test step. + + Scores the ensemble mean against the target states (see + ``_ensemble_step``) and collects per-variable ensemble-mean MSE for + epoch-end aggregation. Parameters ---------- @@ -126,14 +158,32 @@ def test_step(self, batch, batch_idx): The batch of data. batch_idx : int The index of the batch. + """ + entry_mses = self._ensemble_step(batch, "test") + self.test_metrics["ens_mse"].append(entry_mses) - Raises - ------ - NotImplementedError - Always; only training and ensemble validation are implemented - for probabilistic forecasters. + def on_test_epoch_end(self): """ - raise NotImplementedError( - "Ensemble test evaluation is not implemented for " - "probabilistic forecasters." - ) + Perform actions at the end of the test epoch. + + Aggregates ensemble test metrics. Overrides + ``ForecasterModule.on_test_epoch_end``, which also handles spatial + loss maps and example plots that ``test_step`` here does not + populate. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + if self.trainer.is_global_zero and self.hparams.metrics_watch: + unmatched = set(self.hparams.metrics_watch) - self.matched_metrics + if unmatched: + warnings.warn( + "The following metrics in --metrics_watch " + "were not found during test phase: " + f"{sorted(unmatched)}. Ensure the metric prefix " + "matches the evaluation mode (expected 'test_')." + ) + + self.matched_metrics = set() + + for metric_list in self.test_metrics.values(): + metric_list.clear() diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index 28006671..aecdb09f 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -327,27 +327,40 @@ def test_probabilistic_module_rejects_empty_eval_ensemble(): ) -def test_probabilistic_module_test_step_not_implemented(): +def test_probabilistic_module_test_step_scores_ensemble_mean(): datastore = init_datastore_example("mdp") predictor = NoisyStepPredictor(datastore=datastore, output_std=False) + config = nlconfig.NeuralLAMConfig( datastore=nlconfig.DatastoreSelection( kind=datastore.SHORT_NAME, config_path=datastore.root_path ) ) - forecaster = ConcreteProbabilisticARForecaster( + forecaster = MemberCountRecordingForecaster( predictor, datastore, config=config ) model = ProbabilisticForecasterModule( forecaster=forecaster, config=config, datastore=datastore, - eval_ensemble_size=2, + eval_ensemble_size=3, ) - init_states, forcing_features, target_states = _example_batch(datastore) - batch_times = torch.zeros(init_states.shape[0], target_states.shape[1]) + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, B=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) batch = (init_states, target_states, forcing_features, batch_times) - with pytest.raises(NotImplementedError): - model.test_step(batch, 0) + torch.manual_seed(42) + model.test_step(batch, 0) + + # Test samples the configured number of evaluation members + assert forecaster.last_num_members == 3 + + # Ensemble-mean MSE entries are collected for epoch-end aggregation + d_state = target_states.shape[-1] + (entry_mses,) = model.test_metrics["ens_mse"] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) From 703e53de6ad6155d05aab77df3215dc952137bb0 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 12:27:26 +0530 Subject: [PATCH 39/51] Address PR review: separate ensemble RMSE from validation loss naming Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the training loss, per review feedback. --- neural_lam/models/probabilistic_module.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py index 453b00c6..4a2843bd 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/probabilistic_module.py @@ -63,7 +63,10 @@ def _ensemble_step(self, batch, phase: str): ``self.eval_ensemble_size`` members, scores the ensemble mean with plain (unweighted) MSE on interior nodes, logs the root-mean-squared error per configured rollout step and averaged over the rollout - under the given phase's prefix. + under the given phase's prefix. This RMSE is a diagnostic metric, + not the training loss: it always scores the ensemble mean with + plain MSE, regardless of what objective ``compute_training_loss`` + actually trains on, which is not recomputed here. Parameters ---------- @@ -105,11 +108,11 @@ def _ensemble_step(self, batch, phase: str): self._warn_skipped_val_steps(len(time_step_rmse), phase) log_dict = { - f"{phase}_loss_unroll{step}": time_step_rmse[step - 1] + f"{phase}_ens_rmse_unroll{step}": time_step_rmse[step - 1] for step in self.hparams.val_steps_to_log if step <= len(time_step_rmse) } - log_dict[f"{phase}_mean_loss"] = mean_rmse + log_dict[f"{phase}_mean_ens_rmse"] = mean_rmse self.log_dict( log_dict, on_step=False, From 45ffdeb707874b3d3f52819ddd6d56b1e98b8ba6 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 12:41:17 +0530 Subject: [PATCH 40/51] Address PR review: drop redundant inline comments in probabilistic tests Remove explanatory comments around the per_var_std overrides; the assignments are clear on their own. --- tests/test_probabilistic_forecaster.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index aecdb09f..f5b87d3d 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -107,8 +107,6 @@ def test_ar_forecaster_training_loss_matches_direct_score(): score_metric = metrics.get_metric("mse") interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] - # per_var_std is normally computed from config; override directly since - # this test only cares about the loss computation, not standardization. forecaster.per_var_std = torch.ones(d_state) batch_loss, loss_components = forecaster.compute_training_loss( @@ -183,8 +181,6 @@ def test_probabilistic_training_loss_gradient_flow(): init_states, forcing_features, target_states = _example_batch(datastore) interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) d_state = target_states.shape[-1] - # per_var_std is normally computed from config; override directly since - # this test only cares about the loss computation, not standardization. forecaster.per_var_std = torch.ones(d_state) torch.manual_seed(42) From d020691100e5ace37172f958081a8c66ca191db2 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Fri, 10 Jul 2026 13:08:57 +0530 Subject: [PATCH 41/51] Address PR review: split ForecasterModule into an abstract base + concrete deterministic/probabilistic modules Introduce BaseForecasterModule (abstract) under models/forecasters/ holding shared plumbing (training_step, common_step, batch standardization, checkpoint compatibility, plotting/aggregation helpers), with validation_step, test_step and on_test_epoch_end left abstract since they differ meaningfully between evaluation modes. Rename ForecasterModule to DeterministicForecasterModule and move it, alongside ProbabilisticForecasterModule, into forecasters/ as siblings implementing the shared contract, rather than one subclassing the other. --- neural_lam/models/__init__.py | 5 +- .../models/forecasters/autoregressive.py | 2 +- neural_lam/models/forecasters/base.py | 6 +- .../{module.py => forecasters/base_module.py} | 249 +++------------- .../forecasters/deterministic_module.py | 272 ++++++++++++++++++ .../{ => forecasters}/probabilistic_module.py | 38 +-- neural_lam/train_model.py | 12 +- neural_lam/weather_dataset.py | 2 +- tests/test_checkpoint.py | 8 +- tests/test_datasets.py | 4 +- tests/test_gpu_normalization.py | 15 +- tests/test_plotting.py | 12 +- tests/test_prediction_model_classes.py | 21 +- tests/test_probabilistic_forecaster.py | 4 +- tests/test_train_model_warnings.py | 8 +- tests/test_training.py | 18 +- 16 files changed, 398 insertions(+), 278 deletions(-) rename neural_lam/models/{module.py => forecasters/base_module.py} (79%) create mode 100644 neural_lam/models/forecasters/deterministic_module.py rename neural_lam/models/{ => forecasters}/probabilistic_module.py (82%) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 1bfe9eb5..cbeb1b01 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,12 +3,13 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.base_module import BaseForecasterModule +from .forecasters.deterministic_module import DeterministicForecasterModule from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, ) -from .module import ForecasterModule -from .probabilistic_module import ProbabilisticForecasterModule +from .forecasters.probabilistic_module import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_lam import GraphLAM diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index 74f2dc77..a121bfef 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -141,7 +141,7 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the wrapped predictor outputs an std, otherwise ``None`` (in which - case ``ForecasterModule`` substitutes the constant + case ``DeterministicForecasterModule`` substitutes the constant per-variable std). Dims: same as ``prediction``. """ diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 8dfb0002..1e5b7db1 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -77,7 +77,7 @@ def forward( ``predicts_std`` is True, otherwise ``None``. Per-feature predicted standard deviation; when ``None``, the constant per-variable std is substituted upstream by - ``ForecasterModule``. Dims: same as ``prediction``. + ``DeterministicForecasterModule``. Dims: same as ``prediction``. """ @abstractmethod @@ -95,8 +95,8 @@ def compute_training_loss( to produce from the batch, which loss terms to compute from them and how to combine those terms into a single scalar, using its own ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The - wrapping ``ForecasterModule`` only injects the interior mask, logs - the returned components and optimizes the returned loss. + wrapping ``BaseForecasterModule`` only injects the interior mask, + logs the returned components and optimizes the returned loss. Parameters ---------- diff --git a/neural_lam/models/module.py b/neural_lam/models/forecasters/base_module.py similarity index 79% rename from neural_lam/models/module.py rename to neural_lam/models/forecasters/base_module.py index e7420b52..e32a99b4 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -1,9 +1,10 @@ -"""Lightning module handling training, validation and testing loops.""" +"""Abstract Lightning module shared by deterministic and probabilistic +forecaster modules.""" # Standard library import os import warnings -from typing import Any +from abc import ABC, abstractmethod # Third-party import matplotlib.pyplot as plt @@ -17,17 +18,26 @@ from neural_lam.utils import get_integer_time # Local -from .. import metrics, vis -from ..config import NeuralLAMConfig -from ..datastore import BaseDatastore -from ..weather_dataset import WeatherDataset -from .forecasters.base import Forecaster +from ... import vis +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ...weather_dataset import WeatherDataset +from .base import Forecaster -class ForecasterModule(pl.LightningModule): +class BaseForecasterModule(pl.LightningModule, ABC): """ - Lightning module handling training, validation and testing loops. - Wraps a Forecaster instance which performs the actual prediction. + Abstract Lightning module wrapping a ``Forecaster``. + + Owns everything that does not depend on whether the wrapped forecaster + produces a single deterministic forecast or samples an ensemble: + batch standardization, the training loop, optimizer configuration, + checkpoint compatibility, and the plotting/aggregation helpers used by + validation and testing. ``validation_step``, ``test_step`` and + ``on_test_epoch_end`` differ enough between the two evaluation modes + that they are left abstract; concrete subclasses implement them + independently (see ``DeterministicForecasterModule`` and + ``ProbabilisticForecasterModule``) rather than overriding one another. """ # pylint: disable=arguments-differ @@ -47,7 +57,7 @@ def __init__( args=None, ): """ - Initialize the ForecasterModule. + Initialize the BaseForecasterModule. Parameters ---------- @@ -190,16 +200,6 @@ def __init__( self.forcing_mean = None self.forcing_std = None - self.val_metrics: dict[str, list] = { - "mse": [], - } - self.test_metrics: dict[str, list] = { - "mse": [], - "mae": [], - } - if self.forecaster.predicts_std: - self.test_metrics["output_std"] = [] # Treat as metric - # For making restoring of optimizer state optional self.restore_opt = restore_opt @@ -208,9 +208,6 @@ def __init__( self.create_gif = create_gif self.plotted_examples = 0 - # For storing spatial loss maps during evaluation - self.spatial_loss_maps: list[Any] = [] - # Warn once per phase if val_steps_to_log exceeds the actual rollout self._val_steps_warn_issued = False self._test_steps_warn_issued = False @@ -418,10 +415,15 @@ def _warn_skipped_val_steps(self, pred_steps: int, phase: str) -> None: ) setattr(self, flag, True) + @abstractmethod def validation_step(self, batch, batch_idx): """ Perform a single validation step. + Concrete subclasses must both score the batch and populate + ``self.val_metrics`` for epoch-end aggregation by + ``on_validation_epoch_end``. + Parameters ---------- batch : tuple @@ -429,44 +431,6 @@ def validation_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.forecaster.per_var_std - - time_step_loss = torch.mean( - self.forecaster.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - mean_loss = torch.mean(time_step_loss) - self._warn_skipped_val_steps(len(time_step_loss), "val") - - val_log_dict = { - f"val_loss_unroll{step}": time_step_loss[step - 1] - for step in self.hparams.val_steps_to_log - if step <= len(time_step_loss) - } - val_log_dict["val_mean_loss"] = mean_loss - self.log_dict( - val_log_dict, - on_step=False, - on_epoch=True, - sync_dist=True, - batch_size=batch[0].shape[0], - ) - - entry_mses = metrics.mse( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.val_metrics["mse"].append(entry_mses) def on_validation_epoch_end(self): """ @@ -490,11 +454,15 @@ def on_validation_epoch_end(self): for metric_list in self.val_metrics.values(): metric_list.clear() - # pylint: disable-next=unused-argument + @abstractmethod def test_step(self, batch, batch_idx): """ Perform a single test step. + Concrete subclasses must both score the batch and populate + ``self.test_metrics`` for epoch-end aggregation by + ``on_test_epoch_end``. + Parameters ---------- batch : tuple @@ -502,83 +470,6 @@ def test_step(self, batch, batch_idx): batch_idx : int The index of the batch. """ - prediction, target_states, pred_std, _ = self.common_step(batch) - - if pred_std is not None: - mean_pred_std = torch.mean( - pred_std[..., self.interior_mask_bool, :], dim=-2 - ) - self.test_metrics["output_std"].append(mean_pred_std) - - if pred_std is None: - pred_std = self.forecaster.per_var_std - - time_step_loss = torch.mean( - self.forecaster.loss( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - ), - dim=0, - ) - mean_loss = torch.mean(time_step_loss) - self._warn_skipped_val_steps(len(time_step_loss), "test") - - test_log_dict = { - f"test_loss_unroll{step}": time_step_loss[step - 1] - for step in self.hparams.val_steps_to_log - if step <= len(time_step_loss) - } - test_log_dict["test_mean_loss"] = mean_loss - - self.log_dict( - test_log_dict, - on_step=False, - on_epoch=True, - sync_dist=True, - batch_size=batch[0].shape[0], - ) - - for metric_name in ("mse", "mae"): - metric_func = metrics.get_metric(metric_name) - batch_metric_vals = metric_func( - prediction, - target_states, - pred_std, - mask=self.interior_mask_bool, - sum_vars=False, - ) - self.test_metrics[metric_name].append(batch_metric_vals) - - spatial_loss = self.forecaster.loss( - prediction, target_states, pred_std, average_grid=False - ) - log_spatial_losses = spatial_loss[ - :, - [ - step - 1 - for step in self.hparams.val_steps_to_log - if step <= spatial_loss.shape[1] - ], - ] - self.spatial_loss_maps.append(log_spatial_losses) - - if ( - self.trainer.is_global_zero - and self.plotted_examples < self.n_example_pred - ): - n_additional_examples = min( - prediction.shape[0], - self.n_example_pred - self.plotted_examples, - ) - - self.plot_examples( - batch, - n_additional_examples, - prediction=prediction, - split="test", - ) def plot_examples(self, batch, n_examples, split, prediction): """ @@ -867,82 +758,16 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix): plt.close("all") + @abstractmethod def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - Aggregates and plots test metrics and spatial loss maps. - """ - self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") - - spatial_loss_tensor = self.all_gather_cat( - torch.cat(self.spatial_loss_maps, dim=0) - ) - if self.trainer.is_global_zero: - mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) - - loss_map_figs = [ - vis.plot_spatial_error( - error=loss_map, - datastore=self.datastore, - title=f"Test loss, t={t_i} " - f"({(self.time_step_int * t_i)} {self.time_step_unit})", - ) - for t_i, loss_map in zip( - self.hparams.val_steps_to_log, mean_spatial_loss - ) - ] - - for i, fig in enumerate(loss_map_figs): - key = "test_loss" - if not isinstance(self.logger, pl.loggers.WandbLogger): - key = f"{key}_{i}" - if hasattr(self.logger, "log_image"): - self.logger.log_image(key=key, images=[fig]) - pdf_loss_map_figs = [ - vis.plot_spatial_error(error=loss_map, datastore=self.datastore) - for loss_map in mean_spatial_loss - ] - pdf_loss_maps_dir = os.path.join( - self.logger.save_dir, "spatial_loss_maps" - ) - os.makedirs(pdf_loss_maps_dir, exist_ok=True) - for t_i, fig in zip( - self.hparams.val_steps_to_log, pdf_loss_map_figs - ): - fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) - - torch.save( - mean_spatial_loss.cpu(), - os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), - ) - - if self.hparams.metrics_watch: - unmatched = ( - set(self.hparams.metrics_watch) - self.matched_metrics - ) - if unmatched: - warnings.warn( - "The following metrics in --metrics_watch " - "were not found during test phase: " - f"{sorted(unmatched)}. Ensure the metric prefix " - "matches the evaluation mode (expected 'test_')." - ) - - self.matched_metrics = set() - self.spatial_loss_maps.clear() - - # Clear stored test metrics so repeated `trainer.test()` calls on - # the same model instance start from a clean slate (otherwise the - # tensors accumulate and skew the aggregated metrics). - for metric_list in self.test_metrics.values(): - metric_list.clear() - - # Reset the example-plot counter so example prediction plots are - # generated again on every `trainer.test()` call, not just the - # first one (the guard `plotted_examples < n_example_pred` would - # otherwise stay permanently False). - self.plotted_examples = 0 + Concrete subclasses must at least aggregate and plot + ``self.test_metrics`` (typically via ``aggregate_and_plot_metrics``) + and reset any epoch-scoped state they accumulate during + ``test_step``. + """ def on_load_checkpoint(self, checkpoint): """ diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/forecasters/deterministic_module.py new file mode 100644 index 00000000..0cd1cbe8 --- /dev/null +++ b/neural_lam/models/forecasters/deterministic_module.py @@ -0,0 +1,272 @@ +"""Lightning module evaluating forecasters through a single deterministic +rollout per batch.""" + +# Standard library +import os +import warnings +from typing import Any + +# Third-party +import pytorch_lightning as pl +import torch + +# Local +from ... import metrics, vis +from .base_module import BaseForecasterModule + + +class DeterministicForecasterModule(BaseForecasterModule): + """ + Lightning module for a single deterministic forecast per batch. + + Validation and testing score the forecaster's own single-rollout + prediction directly with ``forecaster.loss``, as opposed to + ``ProbabilisticForecasterModule``, which samples and scores an + ensemble. Training is shared with that module unchanged (see + ``BaseForecasterModule.training_step``). + """ + + def __init__(self, *args, **kwargs): + """ + Initialize the module and its deterministic evaluation metrics. + + Parameters + ---------- + *args + Positional arguments forwarded to + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, + ``datastore``, ...). + **kwargs + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` + (``lr``, ...). + """ + super().__init__(*args, **kwargs) + self.val_metrics: dict[str, list] = { + "mse": [], + } + self.test_metrics: dict[str, list] = { + "mse": [], + "mae": [], + } + if self.forecaster.predicts_std: + self.test_metrics["output_std"] = [] # Treat as metric + + # For storing spatial loss maps during evaluation + self.spatial_loss_maps: list[Any] = [] + + def validation_step(self, batch, batch_idx): + """ + Perform a single validation step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + if pred_std is None: + pred_std = self.forecaster.per_var_std + + time_step_loss = torch.mean( + self.forecaster.loss( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + self._warn_skipped_val_steps(len(time_step_loss), "val") + + val_log_dict = { + f"val_loss_unroll{step}": time_step_loss[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_loss) + } + val_log_dict["val_mean_loss"] = mean_loss + self.log_dict( + val_log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + entry_mses = metrics.mse( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.val_metrics["mse"].append(entry_mses) + + # pylint: disable-next=unused-argument + def test_step(self, batch, batch_idx): + """ + Perform a single test step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ + prediction, target_states, pred_std, _ = self.common_step(batch) + + if pred_std is not None: + mean_pred_std = torch.mean( + pred_std[..., self.interior_mask_bool, :], dim=-2 + ) + self.test_metrics["output_std"].append(mean_pred_std) + + if pred_std is None: + pred_std = self.forecaster.per_var_std + + time_step_loss = torch.mean( + self.forecaster.loss( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + ), + dim=0, + ) + mean_loss = torch.mean(time_step_loss) + self._warn_skipped_val_steps(len(time_step_loss), "test") + + test_log_dict = { + f"test_loss_unroll{step}": time_step_loss[step - 1] + for step in self.hparams.val_steps_to_log + if step <= len(time_step_loss) + } + test_log_dict["test_mean_loss"] = mean_loss + + self.log_dict( + test_log_dict, + on_step=False, + on_epoch=True, + sync_dist=True, + batch_size=batch[0].shape[0], + ) + + for metric_name in ("mse", "mae"): + metric_func = metrics.get_metric(metric_name) + batch_metric_vals = metric_func( + prediction, + target_states, + pred_std, + mask=self.interior_mask_bool, + sum_vars=False, + ) + self.test_metrics[metric_name].append(batch_metric_vals) + + spatial_loss = self.forecaster.loss( + prediction, target_states, pred_std, average_grid=False + ) + log_spatial_losses = spatial_loss[ + :, + [ + step - 1 + for step in self.hparams.val_steps_to_log + if step <= spatial_loss.shape[1] + ], + ] + self.spatial_loss_maps.append(log_spatial_losses) + + if ( + self.trainer.is_global_zero + and self.plotted_examples < self.n_example_pred + ): + n_additional_examples = min( + prediction.shape[0], + self.n_example_pred - self.plotted_examples, + ) + + self.plot_examples( + batch, + n_additional_examples, + prediction=prediction, + split="test", + ) + + def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + Aggregates and plots test metrics and spatial loss maps. + """ + self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") + + spatial_loss_tensor = self.all_gather_cat( + torch.cat(self.spatial_loss_maps, dim=0) + ) + if self.trainer.is_global_zero: + mean_spatial_loss = torch.mean(spatial_loss_tensor, dim=0) + + loss_map_figs = [ + vis.plot_spatial_error( + error=loss_map, + datastore=self.datastore, + title=f"Test loss, t={t_i} " + f"({(self.time_step_int * t_i)} {self.time_step_unit})", + ) + for t_i, loss_map in zip( + self.hparams.val_steps_to_log, mean_spatial_loss + ) + ] + + for i, fig in enumerate(loss_map_figs): + key = "test_loss" + if not isinstance(self.logger, pl.loggers.WandbLogger): + key = f"{key}_{i}" + if hasattr(self.logger, "log_image"): + self.logger.log_image(key=key, images=[fig]) + + pdf_loss_map_figs = [ + vis.plot_spatial_error(error=loss_map, datastore=self.datastore) + for loss_map in mean_spatial_loss + ] + pdf_loss_maps_dir = os.path.join( + self.logger.save_dir, "spatial_loss_maps" + ) + os.makedirs(pdf_loss_maps_dir, exist_ok=True) + for t_i, fig in zip( + self.hparams.val_steps_to_log, pdf_loss_map_figs + ): + fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf")) + + torch.save( + mean_spatial_loss.cpu(), + os.path.join(self.logger.save_dir, "mean_spatial_loss.pt"), + ) + + if self.hparams.metrics_watch: + unmatched = ( + set(self.hparams.metrics_watch) - self.matched_metrics + ) + if unmatched: + warnings.warn( + "The following metrics in --metrics_watch " + "were not found during test phase: " + f"{sorted(unmatched)}. Ensure the metric prefix " + "matches the evaluation mode (expected 'test_')." + ) + + self.matched_metrics = set() + self.spatial_loss_maps.clear() + + # Clear stored test metrics so repeated `trainer.test()` calls on + # the same model instance start from a clean slate (otherwise the + # tensors accumulate and skew the aggregated metrics). + for metric_list in self.test_metrics.values(): + metric_list.clear() + + # Reset the example-plot counter so example prediction plots are + # generated again on every `trainer.test()` call, not just the + # first one (the guard `plotted_examples < n_example_pred` would + # otherwise stay permanently False). + self.plotted_examples = 0 diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/forecasters/probabilistic_module.py similarity index 82% rename from neural_lam/models/probabilistic_module.py rename to neural_lam/models/forecasters/probabilistic_module.py index 4a2843bd..cf8d38b0 100644 --- a/neural_lam/models/probabilistic_module.py +++ b/neural_lam/models/forecasters/probabilistic_module.py @@ -7,22 +7,22 @@ import torch # Local -from .. import metrics -from .forecasters.probabilistic import ProbabilisticForecaster -from .module import ForecasterModule +from ... import metrics +from .base_module import BaseForecasterModule +from .probabilistic import ProbabilisticForecaster -class ProbabilisticForecasterModule(ForecasterModule): +class ProbabilisticForecasterModule(BaseForecasterModule): """ Lightning module for forecasters that sample ensemble forecasts. - Training is inherited unchanged from ``ForecasterModule``: the wrapped - forecaster assembles its own training loss. Validation and testing are - ensemble based instead of deterministic: an ensemble is sampled from - the forecaster and scored through its ensemble mean (root-mean-squared - error of the ensemble mean). The module only assumes that the - forecaster can sample ensemble forecasts of the correct shape; it makes - no assumption on how the members are produced. + Training is inherited unchanged from ``BaseForecasterModule``: the + wrapped forecaster assembles its own training loss. Validation and + testing are ensemble based instead of deterministic: an ensemble is + sampled from the forecaster and scored through its ensemble mean + (root-mean-squared error of the ensemble mean). The module only assumes + that the forecaster can sample ensemble forecasts of the correct shape; + it makes no assumption on how the members are produced. """ # The wrapped forecaster must be able to sample ensemble forecasts @@ -36,13 +36,13 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): ---------- *args Positional arguments forwarded to - ``ForecasterModule.__init__`` (``forecaster``, ``config``, + ``BaseForecasterModule.__init__`` (``forecaster``, ``config``, ``datastore``, ...). eval_ensemble_size : int Number of ensemble members sampled during validation and testing. **kwargs - Keyword arguments forwarded to ``ForecasterModule.__init__`` + Keyword arguments forwarded to ``BaseForecasterModule.__init__`` (``lr``, ...). """ super().__init__(*args, **kwargs) @@ -52,8 +52,8 @@ def __init__(self, *args, eval_ensemble_size: int, **kwargs): f"got {eval_ensemble_size}" ) self.eval_ensemble_size = eval_ensemble_size - self.val_metrics = {"ens_mse": []} - self.test_metrics = {"ens_mse": []} + self.val_metrics: dict[str, list] = {"ens_mse": []} + self.test_metrics: dict[str, list] = {"ens_mse": []} def _ensemble_step(self, batch, phase: str): """ @@ -169,10 +169,10 @@ def on_test_epoch_end(self): """ Perform actions at the end of the test epoch. - Aggregates ensemble test metrics. Overrides - ``ForecasterModule.on_test_epoch_end``, which also handles spatial - loss maps and example plots that ``test_step`` here does not - populate. + Aggregates ensemble test metrics. Implements + ``BaseForecasterModule.on_test_epoch_end`` without the spatial loss + maps and example plots that ``DeterministicForecasterModule`` adds, + since ``test_step`` here does not populate them. """ self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index d5e86536..b0709a89 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -19,7 +19,7 @@ from . import utils from .config import load_config_and_datastore from .gnn_layers import GNN_TYPES -from .models import MODELS, ARForecaster, ForecasterModule +from .models import MODELS, ARForecaster, DeterministicForecasterModule from .weather_dataset import WeatherDataModule @@ -40,8 +40,8 @@ def __init__(self, prog): def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): """ - Reconstruct a ForecasterModule from a checkpoint without requiring the - caller to know the original architecture kwargs. + Reconstruct a DeterministicForecasterModule from a checkpoint without + requiring the caller to know the original architecture kwargs. The checkpoint must have been saved with args in hyper_parameters (i.e. created via train_model.main), so that model class and architecture kwargs @@ -66,7 +66,7 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): forecaster = ARForecaster( predictor, datastore, config=config, loss=args.loss ) - return ForecasterModule.load_from_checkpoint( + return DeterministicForecasterModule.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, @@ -440,7 +440,7 @@ def main(input_args=None): raise ValueError("devices should be 'auto' or a list of integers") # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule predictor_class = MODELS[args.model] predictor = predictor_class( datastore=datastore, @@ -463,7 +463,7 @@ def main(input_args=None): predictor, datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 4168396a..c62e1e5a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -471,7 +471,7 @@ def __getitem__( target states, forcing and batch times. The returned data is unstandardized; normalization is applied on-device - in `ForecasterModule.on_after_batch_transfer`. + in `BaseForecasterModule.on_after_batch_transfer`. Parameters ---------- diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 6f114043..6644aae6 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -8,7 +8,11 @@ # First-party from neural_lam import config as nlconfig from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from tests.dummy_datastore import DummyDatastore @@ -51,7 +55,7 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path): output_clamping_upper=config.training.output_clamping.upper, ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 1941206b..319e3372 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -12,7 +12,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore @@ -239,7 +239,7 @@ def _create_graph(): predictor, datastore=datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py index b063d626..dcc8ba63 100644 --- a/tests/test_gpu_normalization.py +++ b/tests/test_gpu_normalization.py @@ -4,7 +4,11 @@ # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -13,7 +17,8 @@ class _MockStepPredictor(StepPredictor): - """Minimal predictor so a ForecasterModule can be built without a graph.""" + """Minimal predictor so a DeterministicForecasterModule can be built + without a graph.""" def forward(self, prev_state, prev_prev_state, forcing): return torch.zeros_like(prev_state), None @@ -27,7 +32,7 @@ def _build_module(datastore): ) predictor = _MockStepPredictor(datastore=datastore, output_std=False) forecaster = ARForecaster(predictor, datastore, config=config) - return ForecasterModule( + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore ) @@ -111,7 +116,9 @@ def test_safe_std_clamps_near_zero(): eps = torch.finfo(torch.float32).eps with pytest.warns(UserWarning, match="near-zero std"): - std = ForecasterModule._safe_std([0.0, 1.0, 2.0], eps, "state") + std = DeterministicForecasterModule._safe_std( + [0.0, 1.0, 2.0], eps, "state" + ) assert std[0] == eps assert std[1] == 1.0 diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 970590be..f55755e5 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -18,7 +18,11 @@ from neural_lam import config as nlconfig from neural_lam import vis from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models import ARForecaster, ForecasterModule, GraphLAM +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + GraphLAM, +) from neural_lam.weather_dataset import WeatherDataset from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -471,7 +475,7 @@ class ModelArgs: predictor, datastore=datastore, config=config, loss=args.loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -666,7 +670,7 @@ class _SimpleLogger: # Shared ModelArgs for metrics_watch regression tests (issue #302). # Kept at module level to avoid copy-paste duplication across tests. def _build_metrics_watch_module(datastore, config): - """Build a ForecasterModule wired for metrics_watch tests.""" + """Build a DeterministicForecasterModule wired for metrics_watch tests.""" predictor = GraphLAM( datastore=datastore, graph_name="1level", @@ -681,7 +685,7 @@ def _build_metrics_watch_module(datastore, config): output_clamping_upper=config.training.output_clamping.upper, ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - return ForecasterModule( + return DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 9bc9d9c0..b2d901e7 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -7,7 +7,11 @@ # First-party from neural_lam import config as nlconfig -from neural_lam.models import ARForecaster, ForecasterModule, StepPredictor +from neural_lam.models import ( + ARForecaster, + DeterministicForecasterModule, + StepPredictor, +) from tests.conftest import init_datastore_example from tests.dummy_datastore import DummyDatastore @@ -81,7 +85,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS @@ -99,7 +103,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -137,7 +141,7 @@ def test_forecaster_module_checkpoint(tmp_path): ) # Load from checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -195,7 +199,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): output_std=False, ) # Use distinctive non-default values so we can detect silent fallback - # to ForecasterModule's defaults during load. + # to DeterministicForecasterModule's defaults during load. saved_loss = "mse" saved_lr = 0.123 saved_create_gif = True @@ -206,7 +210,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): predictor, datastore, config=config, loss=saved_loss ) - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -276,7 +280,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): ) # Load from hacked old checkpoint - loaded_model = ForecasterModule.load_from_checkpoint( + loaded_model = DeterministicForecasterModule.load_from_checkpoint( ckpt_path, datastore=datastore, forecaster=load_forecaster, @@ -287,7 +291,8 @@ def test_forecaster_module_old_checkpoint(tmp_path): assert loaded_model.forecaster.predictor.__class__.__name__ == "GraphLAM" # Hyperparameters nested in the legacy 'args' namespace must round-trip - # rather than silently falling back to ForecasterModule defaults. + # rather than silently falling back to DeterministicForecasterModule + # defaults. assert loaded_model.hparams.lr == saved_lr assert loaded_model.hparams.val_steps_to_log == saved_val_steps assert loaded_model.create_gif is saved_create_gif diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py index f5b87d3d..d98cc345 100644 --- a/tests/test_probabilistic_forecaster.py +++ b/tests/test_probabilistic_forecaster.py @@ -8,7 +8,7 @@ from neural_lam import metrics from neural_lam.models import ( ARForecaster, - ForecasterModule, + DeterministicForecasterModule, ProbabilisticARForecaster, ProbabilisticForecasterModule, StepPredictor, @@ -233,7 +233,7 @@ def test_module_training_step_delegates_to_forecaster(): ) ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a..6bb0654f 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -44,7 +44,8 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): def test_create_gif_forwarded_to_forecaster_module(): - """--create_gif must be forwarded to ForecasterModule.__init__.""" + """--create_gif must be forwarded to + DeterministicForecasterModule.__init__.""" mock_args = MagicMock() mock_args.eval = None mock_args.load = None @@ -76,7 +77,8 @@ def capture_init(_self, **kwargs): patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), patch("neural_lam.train_model.ARForecaster"), patch( - "neural_lam.models.module.ForecasterModule.__init__", + "neural_lam.models.forecasters.deterministic_module." + "DeterministicForecasterModule.__init__", capture_init, ), pytest.raises(SystemExit), @@ -85,5 +87,5 @@ def capture_init(_self, **kwargs): assert ( "create_gif" in captured_kwargs - ), "create_gif was not forwarded to ForecasterModule" + ), "create_gif was not forwarded to DeterministicForecasterModule" assert captured_kwargs["create_gif"] is True diff --git a/tests/test_training.py b/tests/test_training.py index 589e9d89..aca432c5 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -13,7 +13,7 @@ from neural_lam.create_graph import create_graph_from_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore -from neural_lam.models import ForecasterModule +from neural_lam.models import DeterministicForecasterModule from neural_lam.weather_dataset import WeatherDataModule from tests.conftest import init_datastore_example @@ -105,7 +105,7 @@ def run_simple_training( ) # Build predictor and forecaster externally, then inject into - # ForecasterModule + # DeterministicForecasterModule # First-party from neural_lam.models import MODELS, ARForecaster @@ -125,7 +125,7 @@ def run_simple_training( ) forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") - model = ForecasterModule( + model = DeterministicForecasterModule( forecaster=forecaster, config=config, datastore=datastore, @@ -175,9 +175,9 @@ def all_gather(self, tensor_to_gather, sync_grads=False): return tensor_to_gather module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) # Simulate a 3D metric tensor: (N_eval, pred_steps, d_f) @@ -206,9 +206,9 @@ def all_gather(self, tensor, sync_grads=False): return torch.stack([tensor, tensor], dim=0) module = MockModule() - # Bind the real ForecasterModule.all_gather_cat to our mock - module.all_gather_cat = ForecasterModule.all_gather_cat.__get__( - module, MockModule + # Bind the real DeterministicForecasterModule.all_gather_cat to our mock + module.all_gather_cat = ( + DeterministicForecasterModule.all_gather_cat.__get__(module, MockModule) ) tensor = torch.randn(4, 3, 5) # (N_eval, pred_steps, d_f) From f5c5bc088af9a8257e444bcffb7371dac17a85be Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 11 Jul 2026 08:36:07 +0530 Subject: [PATCH 42/51] docs: clarify GNN-type flag scope in train_model.py Note that g2m/m2g flags apply to Graph-EFM too, while mesh_up/mesh_down flags only affect Hi-LAM since Graph-EFM hard-codes those GNN types. --- neural_lam/train_model.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index f98065c4..3abe4860 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -181,29 +181,34 @@ def main(input_args=None): type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for grid-to-mesh encoding", + help="GNN type for grid-to-mesh encoding. Applies to all models, " + "including the probabilistic Graph-EFM model", ) arch_group.add_argument( "--m2g_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for mesh-to-grid decoding", + help="GNN type for mesh-to-grid decoding. Applies to all models, " + "including the probabilistic Graph-EFM model", ) arch_group.add_argument( "--mesh_up_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for upward mesh message passing in hierarchical models", + help="GNN type for upward mesh message passing in hierarchical " + "models. Only affects Hi-LAM; the probabilistic Graph-EFM model " + "hard-codes its mesh-up GNN types", ) arch_group.add_argument( "--mesh_down_gnn_type", type=str, default="InteractionNet", choices=list(GNN_TYPES.keys()), - help="GNN type for downward mesh message passing in " - "hierarchical models", + help="GNN type for downward mesh message passing in hierarchical " + "models. Only affects Hi-LAM; the probabilistic Graph-EFM model " + "hard-codes its mesh-down GNN type", ) # Training options From 961d1f293b51f85e70dbf447bf18f115b53982b3 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 11 Jul 2026 08:58:13 +0530 Subject: [PATCH 43/51] refactor: build the latent prior from BaseGraphEFM.__init__ Move prior construction fully into the base constructor instead of each subclass calling self.build_prior() itself: BaseGraphEFM.__init__ now takes latent_dim/learn_prior/prior_dist/prior_layers/g2m_gnn_type, validates the loaded graph against a new expects_hierarchical class attribute, derives num_mesh_nodes generically, and builds prior_model directly. GraphEFM and GraphEFMMultiScale forward these to super() and drop their duplicated graph-type check, num_mesh_nodes assignment and build_prior call, reusing self.latent_dim instead. --- .../models/step_predictors/graph/graph_efm.py | 165 +++++++++++------- 1 file changed, 103 insertions(+), 62 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 4f54c94e..24d5e1d3 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -41,13 +41,20 @@ class BaseGraphEFM(StepPredictor): 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: their - constructors verify the loaded graph is of the expected type, build the - mesh embedders and the prior/encoder/decoder latent modules, and they - implement :meth:`embedd_mesh`. See :class:`GraphEFM` (hierarchical - graph) and :class:`GraphEFMMultiScale` (flat graph). + graph type: it loads the graph, verifies it matches the type declared + by the subclass's ``expects_hierarchical`` class attribute, and builds + the prior (via :meth:`build_prior`, delegating to the subclass's + :meth:`build_learnable_prior` when learned). Concrete subclasses build + the mesh embedders and the encoder/decoder latent modules, and + implement :meth:`embedd_mesh` and :meth:`build_learnable_prior`. See + :class:`GraphEFM` (hierarchical graph) and :class:`GraphEFMMultiScale` + (flat graph). """ + #: Set by concrete subclasses: whether they require a hierarchical + #: (True) or flat (False) mesh graph. + expects_hierarchical: bool + def __init__( self, config: NeuralLAMConfig, @@ -55,6 +62,11 @@ def __init__( graph_name: str, hidden_dim: int = 64, hidden_layers: int = 1, + latent_dim: Optional[int] = None, + learn_prior: bool = True, + prior_dist: str = "isotropic", + prior_layers: int = 2, + g2m_gnn_type: str = "InteractionNet", num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, @@ -65,9 +77,9 @@ def __init__( 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. + embedders, the constant per-variable std and the prior. Building + the mesh embedders and the encoder/decoder latent modules is left + to the subclass constructor. Parameters ---------- @@ -84,6 +96,25 @@ def __init__( Dimensionality of internal node and edge representations. hidden_layers : int Number of hidden layers in internal MLPs. + latent_dim : int, optional + Dimensionality of the latent variable at each latent-carrying + mesh node (top level for hierarchical graphs, all mesh nodes + for flat graphs); defaults to ``hidden_dim`` when None. The + resolved value is stored as ``self.latent_dim`` for the + subclass to reuse when building its encoder/decoder. + learn_prior : bool + If True, the prior is the graph-type specific learnable encoder + built by :meth:`build_learnable_prior`, conditioned on the + previous state; if False, a constant ``Normal(0, 1)`` prior is + used. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or + ``"diagonal"``. + prior_layers : int + Number of on-mesh GNN layers in the learnable prior. + g2m_gnn_type : str + GNN type for the grid-to-mesh step of the prior (key in + ``gnn_layers.GNN_TYPES``). num_past_forcing_steps : int Number of past forcing steps included in the input window. num_future_forcing_steps : int @@ -108,6 +139,23 @@ def __init__( self.hierarchical = utils.load_and_register_graph( self, datastore, graph_name ) + if self.hierarchical != self.expects_hierarchical: + expected_kind = ( + "hierarchical" if self.expects_hierarchical else "flat" + ) + actual_kind = "hierarchical" if self.hierarchical else "flat" + raise ValueError( + f"{type(self).__name__} requires a {expected_kind} mesh " + f"graph, but graph '{graph_name}' is {actual_kind}" + ) + + # The latent variable lives on the top mesh level for hierarchical + # graphs, and on every mesh node for flat graphs. + self.num_mesh_nodes = ( + self.mesh_static_features[-1].shape[0] + if self.hierarchical + else self.mesh_static_features.shape[0] + ) # Specify dimensions of data self.num_state_vars = datastore.get_num_data_vars(category="state") @@ -166,6 +214,18 @@ def __init__( # inert -- accepted for interface parity with other StepPredictors. self.prepare_clamping_params(datastore) + # Prior over the latent variable. + self.latent_dim = latent_dim if latent_dim is not None else hidden_dim + self.prior_model = self.build_prior( + learn_prior=learn_prior, + latent_dim=self.latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_layers, + ) + def build_prior( self, learn_prior, @@ -182,7 +242,7 @@ def build_prior( When ``learn_prior`` is True the (graph-type specific) learnable prior is delegated to :meth:`build_learnable_prior`; otherwise the constant ``Normal(0, 1)`` prior, which is identical for every graph type, is - built here. Must be called after ``self.num_mesh_nodes`` is set. + built here. Parameters ---------- @@ -598,6 +658,8 @@ class GraphEFM(BaseGraphEFM): decoder is a ``HiGraphLatentDecoder``. """ + expects_hierarchical = True + def __init__( self, config: NeuralLAMConfig, @@ -620,7 +682,8 @@ def __init__( output_clamping_upper: Optional[Dict[str, float]] = None, ): """ - Build the mesh embedders and hierarchical latent modules. + Build the mesh embedders and the hierarchical encoder/decoder + latent modules. The prior is built by the base class. Parameters ---------- @@ -639,9 +702,12 @@ def __init__( 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. + node; defaults to ``hidden_dim`` when None. Forwarded to the + base class, which resolves the default and stores it as + ``self.latent_dim``. prior_intra_level_layers : int Number of intra-level GNN layers in the (learned) prior. + Forwarded to the base class as ``prior_layers``. encoder_intra_level_layers : int Number of intra-level GNN layers in the variational encoder. decoder_intra_level_layers : int @@ -649,10 +715,10 @@ def __init__( 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 is used. Forwarded to the base class. prior_dist : str Output distribution of the prior: ``"isotropic"`` or - ``"diagonal"``. + ``"diagonal"``. Forwarded to the base class. num_past_forcing_steps : int Number of past forcing steps included in the input window. num_future_forcing_steps : int @@ -678,6 +744,11 @@ def __init__( graph_name=graph_name, hidden_dim=hidden_dim, hidden_layers=hidden_layers, + latent_dim=latent_dim, + learn_prior=learn_prior, + prior_dist=prior_dist, + prior_layers=prior_intra_level_layers, + g2m_gnn_type=g2m_gnn_type, num_past_forcing_steps=num_past_forcing_steps, num_future_forcing_steps=num_future_forcing_steps, output_std=output_std, @@ -685,17 +756,9 @@ def __init__( output_clamping_upper=output_clamping_upper, ) - if not self.hierarchical: - raise ValueError( - f"{type(self).__name__} requires a hierarchical mesh graph, " - f"but graph '{graph_name}' is flat" - ) - 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): @@ -754,22 +817,9 @@ def __init__( ] ) - latent_dim = latent_dim if latent_dim is not None else hidden_dim - - # Prior (constant prior shared via the base class) - self.prior_model = self.build_prior( - learn_prior=learn_prior, - latent_dim=latent_dim, - hidden_dim=hidden_dim, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - prior_dist=prior_dist, - prior_layers=prior_intra_level_layers, - ) - # Encoder (variational posterior) + Decoder self.encoder = HiGraphLatentEncoder( - latent_dim=latent_dim, + latent_dim=self.latent_dim, g2m_edge_index=self.g2m_edge_index, m2m_edge_index=self.m2m_edge_index, mesh_up_edge_index=self.mesh_up_edge_index, @@ -786,7 +836,7 @@ def __init__( 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, + latent_dim=self.latent_dim, num_state_vars=self.num_state_vars, intra_level_layers=decoder_intra_level_layers, hidden_layers=hidden_layers, @@ -900,6 +950,8 @@ class GraphEFMMultiScale(BaseGraphEFM): ``GraphLatentDecoder``. """ + expects_hierarchical = False + def __init__( self, config: NeuralLAMConfig, @@ -922,7 +974,8 @@ def __init__( output_clamping_upper: Optional[Dict[str, float]] = None, ): """ - Build the mesh embedders and flat-graph latent modules. + Build the mesh embedders and the flat-graph encoder/decoder latent + modules. The prior is built by the base class. Parameters ---------- @@ -941,9 +994,12 @@ def __init__( 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. + defaults to ``hidden_dim`` when None. Forwarded to the base + class, which resolves the default and stores it as + ``self.latent_dim``. prior_m2m_layers : int Number of on-mesh (m2m) GNN layers in the (learned) prior. + Forwarded to the base class as ``prior_layers``. encoder_m2m_layers : int Number of on-mesh (m2m) GNN layers in the variational encoder. decoder_m2m_layers : int @@ -951,10 +1007,10 @@ def __init__( 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. + used. Forwarded to the base class. prior_dist : str Output distribution of the prior: ``"isotropic"`` or - ``"diagonal"``. + ``"diagonal"``. Forwarded to the base class. num_past_forcing_steps : int Number of past forcing steps included in the input window. num_future_forcing_steps : int @@ -980,6 +1036,11 @@ def __init__( graph_name=graph_name, hidden_dim=hidden_dim, hidden_layers=hidden_layers, + latent_dim=latent_dim, + learn_prior=learn_prior, + prior_dist=prior_dist, + prior_layers=prior_m2m_layers, + g2m_gnn_type=g2m_gnn_type, num_past_forcing_steps=num_past_forcing_steps, num_future_forcing_steps=num_future_forcing_steps, output_std=output_std, @@ -987,13 +1048,6 @@ def __init__( output_clamping_upper=output_clamping_upper, ) - if self.hierarchical: - raise ValueError( - f"{type(self).__name__} requires a flat mesh graph, " - f"but graph '{graph_name}' is hierarchical" - ) - - 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 " @@ -1008,22 +1062,9 @@ def __init__( 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 (constant prior shared via the base class) - self.prior_model = self.build_prior( - learn_prior=learn_prior, - latent_dim=latent_dim, - hidden_dim=hidden_dim, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - prior_dist=prior_dist, - prior_layers=prior_m2m_layers, - ) - # Encoder (variational posterior) + Decoder self.encoder = GraphLatentEncoder( - latent_dim=latent_dim, + latent_dim=self.latent_dim, g2m_edge_index=self.g2m_edge_index, m2m_edge_index=self.m2m_edge_index, hidden_dim=hidden_dim, @@ -1037,7 +1078,7 @@ def __init__( m2m_edge_index=self.m2m_edge_index, m2g_edge_index=self.m2g_edge_index, hidden_dim=hidden_dim, - latent_dim=latent_dim, + latent_dim=self.latent_dim, num_state_vars=self.num_state_vars, m2m_layers=decoder_m2m_layers, hidden_layers=hidden_layers, From c7b784812fc5e20464349a8e9fbb9b6d1addbd79 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Wed, 15 Jul 2026 02:25:54 +0530 Subject: [PATCH 44/51] refactor: strip likelihood/ELBO assembly out of GraphEFMPredictor Removes estimate_likelihood/compute_step_loss and the per_var_std buffer they existed to feed, plus the now-unused config constructor arg (its only use was building per_var_std). Per the objective now living on the Forecaster (#700), the predictor stays a pure network construct: encoder/prior/decoder plus the forward sampling path. --- .../models/step_predictors/graph/graph_efm.py | 228 +----------------- tests/test_graph_efm_predictor.py | 130 +--------- 2 files changed, 20 insertions(+), 338 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 24d5e1d3..0c287fff 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -3,7 +3,7 @@ graphs.""" # Standard library -from typing import Callable, Dict, Optional +from typing import Dict, Optional # Third-party import torch @@ -11,9 +11,7 @@ # 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, @@ -34,11 +32,11 @@ class BaseGraphEFM(StepPredictor): 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. + ``StepPredictor`` directly. ``forward`` samples a single step from the + prior. The encoder (variational posterior) and prior are exposed as + ``self.encoder``/``self.prior_model`` for a forecaster to condition on + the target and assemble a training objective from; the predictor itself + does not compute any loss. This base class sets up everything that is independent of the mesh graph type: it loads the graph, verifies it matches the type declared @@ -57,7 +55,6 @@ class BaseGraphEFM(StepPredictor): def __init__( self, - config: NeuralLAMConfig, datastore: BaseDatastore, graph_name: str, hidden_dim: int = 64, @@ -77,15 +74,11 @@ def __init__( Set up the graph-type independent parts of the predictor. Loads the graph, builds the grid embedders, the grid-mesh edge - embedders, the constant per-variable std and the prior. Building - the mesh embedders and the encoder/decoder latent modules is left - to the subclass constructor. + embedders and the prior. Building the mesh embedders and the + 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. @@ -121,8 +114,7 @@ def __init__( 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. + mean; if False, ``forward`` returns ``None`` for the std. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional @@ -184,31 +176,6 @@ def __init__( 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. @@ -435,167 +402,6 @@ def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): 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, @@ -662,7 +468,6 @@ class GraphEFM(BaseGraphEFM): def __init__( self, - config: NeuralLAMConfig, datastore: BaseDatastore, graph_name: str = "hierarchical", hidden_dim: int = 64, @@ -687,9 +492,6 @@ def __init__( 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. @@ -731,15 +533,13 @@ def __init__( ``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. + mean; if False, ``forward`` returns ``None`` for the std. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional Upper clamping limits per output variable. """ super().__init__( - config=config, datastore=datastore, graph_name=graph_name, hidden_dim=hidden_dim, @@ -954,7 +754,6 @@ class GraphEFMMultiScale(BaseGraphEFM): def __init__( self, - config: NeuralLAMConfig, datastore: BaseDatastore, graph_name: str = "multiscale", hidden_dim: int = 64, @@ -979,9 +778,6 @@ def __init__( 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. @@ -1023,15 +819,13 @@ def __init__( ``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. + mean; if False, ``forward`` returns ``None`` for the std. output_clamping_lower : dict of str to float, optional Lower clamping limits per output variable. output_clamping_upper : dict of str to float, optional Upper clamping limits per output variable. """ super().__init__( - config=config, datastore=datastore, graph_name=graph_name, hidden_dim=hidden_dim, diff --git a/tests/test_graph_efm_predictor.py b/tests/test_graph_efm_predictor.py index 3633c3fb..a18580e0 100644 --- a/tests/test_graph_efm_predictor.py +++ b/tests/test_graph_efm_predictor.py @@ -3,8 +3,7 @@ These mirror the smoke-test pattern used for the deterministic predictors (see ``tests/test_gnn_layers.py``): build the flat (GraphEFMMultiScale) and hierarchical (GraphEFM) variants on the real example datastore with a freshly -created graph, then exercise ``forward``, ``compute_step_loss`` and the -sampling helpers on synthetic tensors. +created graph, then exercise ``forward`` on synthetic tensors. """ # Standard library @@ -15,10 +14,7 @@ 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, GraphEFMMultiScale, @@ -29,15 +25,9 @@ NUM_FUTURE_FORCING_STEPS = 1 -def _datastore_and_config_with_graph(graph_name): +def _datastore_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 @@ -54,11 +44,11 @@ def _datastore_and_config_with_graph(graph_name): hierarchical=hierarchical, n_max_levels=n_max_levels, ) - return datastore, config + return datastore def _build_predictor(graph_name, output_std=False): - datastore, config = _datastore_and_config_with_graph(graph_name) + datastore = _datastore_with_graph(graph_name) if graph_name == "hierarchical": predictor_class = GraphEFM layer_kwargs = { @@ -74,7 +64,6 @@ def _build_predictor(graph_name, output_std=False): "decoder_m2m_layers": 1, } predictor = predictor_class( - config=config, datastore=datastore, graph_name=graph_name, hidden_dim=4, @@ -87,7 +76,7 @@ def _build_predictor(graph_name, output_std=False): output_std=output_std, **layer_kwargs, ) - return predictor, datastore, config + return predictor, datastore def _make_inputs(predictor, datastore, batch_size=2): @@ -107,7 +96,7 @@ def _make_inputs(predictor, datastore, batch_size=2): 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) + predictor, datastore = _build_predictor(graph_name) prev_state, prev_prev_state, forcing, d_state = _make_inputs( predictor, datastore ) @@ -122,7 +111,7 @@ def test_forward_shapes_and_no_std(graph_name): 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) + predictor, datastore = _build_predictor(graph_name, output_std=True) prev_state, prev_prev_state, forcing, d_state = _make_inputs( predictor, datastore ) @@ -136,84 +125,11 @@ def test_forward_output_std_returns_std(graph_name): 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) + 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) @@ -222,33 +138,6 @@ def test_forward_member_stochasticity(graph_name): 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"), (GraphEFMMultiScale, "hierarchical")], @@ -256,10 +145,9 @@ def test_per_var_std_none_when_output_std(): def test_graph_type_mismatch_raises(predictor_class, graph_name): """GraphEFM requires a hierarchical graph and GraphEFMMultiScale a flat one; constructing with the wrong graph type raises ValueError.""" - datastore, config = _datastore_and_config_with_graph(graph_name) + datastore = _datastore_with_graph(graph_name) with pytest.raises(ValueError, match="mesh graph"): predictor_class( - config=config, datastore=datastore, graph_name=graph_name, hidden_dim=4, From 98ab692ad6da66fd08c2df76467c2314a4357373 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 21:04:59 +0530 Subject: [PATCH 45/51] Ignore .idea directory in .gitignore Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries. --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index a2c5bb26..37242198 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,7 @@ tags # Coc configuration directory .vim .vscode +.idea # macos .DS_Store From f0e01b1efccea28e1d257911afba2429bfd5178d Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:15:14 +0530 Subject: [PATCH 46/51] Address PR review: delegate validation/test loss computation to Forecaster DeterministicForecasterModule.validation_step/test_step read self.forecaster.loss and self.forecaster.per_var_std directly, so the Module still knew how to compute a loss from a prediction. Add an abstract Forecaster.score (implemented on ARForecaster) that resolves the pred_std fallback and applies a scoring rule internally; the Module now only calls forecaster.score(...) and never touches loss/per_var_std itself. --- .../models/forecasters/autoregressive.py | 63 +++++++++++++++++++ neural_lam/models/forecasters/base.py | 61 +++++++++++++++++- neural_lam/models/forecasters/base_module.py | 9 +-- .../forecasters/deterministic_module.py | 20 +++--- tests/test_prediction_model_classes.py | 46 ++++++++++++++ 5 files changed, 180 insertions(+), 19 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index a121bfef..ca8166cd 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,5 +1,8 @@ """Forecaster that uses an auto-regressive strategy to unroll a forecast.""" +# Standard library +from typing import Callable, Optional + # Third-party import torch @@ -245,3 +248,63 @@ def compute_training_loss( ) ) return batch_loss, {} + + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Substitutes ``self.per_var_std`` for ``pred_std`` when the latter is + ``None`` (predictor does not output its own std), then applies + ``metric`` (defaulting to ``self.loss``, the configured scoring + rule). + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``; ``None`` when + the wrapped predictor does not output one, in which case + ``self.per_var_std`` is substituted. + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to ``self.loss``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + """ + if pred_std is None: + pred_std = self.per_var_std + metric_fn = self.loss if metric is None else metric + return metric_fn( + prediction, + target_states, + pred_std, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index 1e5b7db1..da63869b 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -2,6 +2,7 @@ # Standard library from abc import ABC, abstractmethod +from typing import Callable, Optional # Third-party import torch @@ -75,9 +76,10 @@ def forward( pred_std : torch.Tensor or None Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when ``predicts_std`` is True, otherwise ``None``. Per-feature - predicted standard deviation; when ``None``, the constant - per-variable std is substituted upstream by - ``DeterministicForecasterModule``. Dims: same as ``prediction``. + predicted standard deviation; when ``None``, the forecaster's + own constant per-variable std fallback is substituted by + ``compute_training_loss``/``score``, not by the caller. Dims: + same as ``prediction``. """ @abstractmethod @@ -134,3 +136,56 @@ def compute_training_loss( the training phase. Empty when the objective has no separate components worth logging. """ + + @abstractmethod + def score( + self, + prediction: torch.Tensor, + target_states: torch.Tensor, + pred_std: Optional[torch.Tensor], + metric: Optional[Callable[..., torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ) -> torch.Tensor: + """ + Score an already-produced prediction for reporting (not training). + + Wrapping ``BaseForecasterModule`` subclasses use this for + validation/test logging and diagnostics instead of computing a loss + themselves: the forecaster owns both its scoring rule and its + ``pred_std`` fallback, so it is the only place that knows how to + turn a raw ``pred_std`` (possibly ``None``) into a valid one and + apply a metric to it. + + Parameters + ---------- + prediction : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. Forecast to + score. + target_states : torch.Tensor + Shape ``(..., num_grid_nodes, num_state_vars)``. True states to + score against. Dims: same as ``prediction``. + pred_std : torch.Tensor or None + Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. + Predicted standard deviation for ``prediction``, as returned + alongside it by ``forward``. When ``None``, implementations + substitute their own constant per-variable std fallback. + metric : callable or None, optional + Scoring function with the ``neural_lam.metrics`` signature + ``(pred, target, pred_std, mask=None, average_grid=True, + sum_vars=True) -> torch.Tensor``. Defaults to the forecaster's + own configured scoring rule when ``None``. + mask : torch.Tensor or None, optional + Shape ``(num_grid_nodes,)``, boolean. Forwarded to ``metric``. + average_grid : bool, optional + Forwarded to ``metric``. + sum_vars : bool, optional + Forwarded to ``metric``. + + Returns + ------- + torch.Tensor + The metric's output; shape depends on ``average_grid`` and + ``sum_vars`` (see ``neural_lam.metrics``). + """ diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/forecasters/base_module.py index e32a99b4..85faa451 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -62,10 +62,11 @@ def __init__( Parameters ---------- forecaster : Forecaster - The forecaster model to use for predictions. Owns the scoring - rule (``forecaster.loss``) and the constant per-variable std - fallback (``forecaster.per_var_std``) used for training and for - validation/test loss reporting here. + The forecaster model to use for predictions. Owns the training + objective (``compute_training_loss``) and validation/test + scoring (``score``); this module and its subclasses never + compute a loss themselves, only inject shared inputs (e.g. the + interior mask) and log what the forecaster returns. config : NeuralLAMConfig Configuration object for the neural LAM model. datastore : BaseDatastore diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/forecasters/deterministic_module.py index 0cd1cbe8..217d4937 100644 --- a/neural_lam/models/forecasters/deterministic_module.py +++ b/neural_lam/models/forecasters/deterministic_module.py @@ -20,7 +20,7 @@ class DeterministicForecasterModule(BaseForecasterModule): Lightning module for a single deterministic forecast per batch. Validation and testing score the forecaster's own single-rollout - prediction directly with ``forecaster.loss``, as opposed to + prediction via ``forecaster.score``, as opposed to ``ProbabilisticForecasterModule``, which samples and scores an ensemble. Training is shared with that module unchanged (see ``BaseForecasterModule.training_step``). @@ -66,11 +66,9 @@ def validation_step(self, batch, batch_idx): The index of the batch. """ prediction, target_states, pred_std, _ = self.common_step(batch) - if pred_std is None: - pred_std = self.forecaster.per_var_std time_step_loss = torch.mean( - self.forecaster.loss( + self.forecaster.score( prediction, target_states, pred_std, @@ -95,10 +93,11 @@ def validation_step(self, batch, batch_idx): batch_size=batch[0].shape[0], ) - entry_mses = metrics.mse( + entry_mses = self.forecaster.score( prediction, target_states, pred_std, + metric=metrics.mse, mask=self.interior_mask_bool, sum_vars=False, ) @@ -124,11 +123,8 @@ def test_step(self, batch, batch_idx): ) self.test_metrics["output_std"].append(mean_pred_std) - if pred_std is None: - pred_std = self.forecaster.per_var_std - time_step_loss = torch.mean( - self.forecaster.loss( + self.forecaster.score( prediction, target_states, pred_std, @@ -155,17 +151,17 @@ def test_step(self, batch, batch_idx): ) for metric_name in ("mse", "mae"): - metric_func = metrics.get_metric(metric_name) - batch_metric_vals = metric_func( + batch_metric_vals = self.forecaster.score( prediction, target_states, pred_std, + metric=metrics.get_metric(metric_name), mask=self.interior_mask_bool, sum_vars=False, ) self.test_metrics[metric_name].append(batch_metric_vals) - spatial_loss = self.forecaster.loss( + spatial_loss = self.forecaster.score( prediction, target_states, pred_std, average_grid=False ) log_spatial_losses = spatial_loss[ diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index b2d901e7..51df3261 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -7,6 +7,7 @@ # First-party from neural_lam import config as nlconfig +from neural_lam import metrics from neural_lam.models import ( ARForecaster, DeterministicForecasterModule, @@ -75,6 +76,51 @@ def test_ar_forecaster_unroll(): assert torch.all(prediction[:, :, 1:, :] == 5.0) +def test_ar_forecaster_score(): + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = MockStepPredictor(datastore=datastore, output_std=False) + forecaster = ARForecaster(predictor, datastore, config=config, loss="mse") + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + mask = torch.ones(num_grid_nodes, dtype=torch.bool) + + # pred_std=None falls back to forecaster.per_var_std and applies the + # forecaster's own configured scoring rule (self.loss) + scored = forecaster.score(prediction, target, None, mask=mask) + expected = forecaster.loss( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored, expected) + + # An explicit metric overrides self.loss, still substituting the + # per_var_std fallback + scored_mse = forecaster.score( + prediction, target, None, metric=metrics.mse, mask=mask + ) + expected_mse = metrics.mse( + prediction, target, forecaster.per_var_std, mask=mask + ) + assert torch.equal(scored_mse, expected_mse) + + # An explicit pred_std is used as-is, not overridden by per_var_std + explicit_std = torch.full((d_state,), 2.0) + scored_explicit = forecaster.score( + prediction, target, explicit_std, mask=mask + ) + expected_explicit = forecaster.loss( + prediction, target, explicit_std, mask=mask + ) + assert torch.equal(scored_explicit, expected_explicit) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") From 8b3ab331bc883c87f9aa2d6f0021524190cc7f35 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:29:42 +0530 Subject: [PATCH 47/51] Address PR review: move per_var_std/config validation into the Forecaster BaseForecasterModule.__init__ raised if forecaster.per_var_std was None and the forecaster didn't predict its own std -- a check on the Forecaster's configuration living in the wrong class. Move it into ARForecaster via a new _resolve_pred_std helper, shared by score() and compute_training_loss(): construction with config=None now always succeeds (a valid state for forecasters only ever used for inference), and the ValueError instead fires from the Forecaster itself, only once scoring is attempted without any std to use. --- .../models/forecasters/autoregressive.py | 71 ++++++++++++++++--- neural_lam/models/forecasters/base_module.py | 8 --- .../models/forecasters/probabilistic.py | 7 +- tests/test_prediction_model_classes.py | 42 +++++++++++ 4 files changed, 107 insertions(+), 21 deletions(-) diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index ca8166cd..e7add53c 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -40,8 +40,11 @@ def __init__( config : NeuralLAMConfig or None Configuration used to compute the constant per-variable std substituted for ``pred_std`` when ``predictor`` does not output - its own (see ``per_var_std``). Only required for that case; - forecasters used purely for inference can omit it. + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. loss : str, default "wmse" The scoring rule (from ``neural_lam.metrics``) used by ``compute_training_loss`` and stored as ``self.loss``. @@ -232,12 +235,18 @@ def compute_training_loss( batch and time. loss_components : dict of {str: torch.Tensor} Empty; the deterministic objective has no separate components. + + Raises + ------ + ValueError + If the predictor does not output its own std and no + ``per_var_std`` fallback is available; see + ``_resolve_pred_std``. """ prediction, pred_std = self( init_states, forcing_features, target_states ) - if pred_std is None: - pred_std = self.per_var_std + pred_std = self._resolve_pred_std(pred_std) batch_loss = torch.mean( self.loss( @@ -249,6 +258,41 @@ def compute_training_loss( ) return batch_loss, {} + def _resolve_pred_std( + self, pred_std: Optional[torch.Tensor] + ) -> torch.Tensor: + """ + Return ``pred_std``, or the constant ``per_var_std`` fallback. + + Parameters + ---------- + pred_std : torch.Tensor or None + Predicted standard deviation as returned by ``forward``, + possibly ``None``. + + Returns + ------- + torch.Tensor + ``pred_std`` unchanged when given; otherwise ``self.per_var_std``. + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available (``predictor.predicts_std`` is False and this + forecaster was constructed without ``config``). + """ + if pred_std is not None: + return pred_std + if self.per_var_std is None: + raise ValueError( + "No pred_std available for scoring: predictor.predicts_std " + "is False and this forecaster has no per_var_std fallback " + "(it was constructed without config). Pass config to the " + "constructor, or use a predictor that outputs its own std." + ) + return self.per_var_std + def score( self, prediction: torch.Tensor, @@ -262,10 +306,9 @@ def score( """ Score an already-produced prediction for reporting (not training). - Substitutes ``self.per_var_std`` for ``pred_std`` when the latter is - ``None`` (predictor does not output its own std), then applies - ``metric`` (defaulting to ``self.loss``, the configured scoring - rule). + Resolves ``pred_std`` via ``_resolve_pred_std`` (substituting + ``self.per_var_std`` when ``None``), then applies ``metric`` + (defaulting to ``self.loss``, the configured scoring rule). Parameters ---------- @@ -279,7 +322,8 @@ def score( Shape ``(..., num_grid_nodes, num_state_vars)``, or ``None``. Predicted standard deviation for ``prediction``; ``None`` when the wrapped predictor does not output one, in which case - ``self.per_var_std`` is substituted. + ``self.per_var_std`` is substituted (see ``_resolve_pred_std`` + for when this raises instead). metric : callable or None, optional Scoring function with the ``neural_lam.metrics`` signature ``(pred, target, pred_std, mask=None, average_grid=True, @@ -296,9 +340,14 @@ def score( torch.Tensor The metric's output; shape depends on ``average_grid`` and ``sum_vars`` (see ``neural_lam.metrics``). + + Raises + ------ + ValueError + If ``pred_std`` is ``None`` and no ``per_var_std`` fallback is + available; see ``_resolve_pred_std``. """ - if pred_std is None: - pred_std = self.per_var_std + pred_std = self._resolve_pred_std(pred_std) metric_fn = self.loss if metric is None else metric return metric_fn( prediction, diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/forecasters/base_module.py index 85faa451..c77de238 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/forecasters/base_module.py @@ -133,14 +133,6 @@ def __init__( self.save_hyperparameters(ignore=["datastore", "forecaster"]) self.datastore = datastore self.forecaster = forecaster - if forecaster.per_var_std is None and not forecaster.predicts_std: - raise ValueError( - "forecaster.per_var_std is None but the forecaster does " - "not predict its own std (forecaster.predicts_std is " - "False), so training/validation/test scoring has no std " - "to use. Pass config to the forecaster's constructor so " - "it can compute the constant per-variable std." - ) self.matched_metrics: set = set() # Compute interior_mask_bool directly from datastore diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py index 1c5aeeb4..57526d1c 100644 --- a/neural_lam/models/forecasters/probabilistic.py +++ b/neural_lam/models/forecasters/probabilistic.py @@ -123,8 +123,11 @@ def __init__( config : NeuralLAMConfig or None Configuration used to compute the constant per-variable std substituted for ``pred_std`` when ``predictor`` does not output - its own (see ``per_var_std``). Only required for that case; - forecasters used purely for inference can omit it. + its own (see ``per_var_std``). Required in that case for + ``score``/``compute_training_loss`` to work (they raise + ``ValueError`` via ``_resolve_pred_std`` otherwise); forecasters + used purely for inference (``forward``/``sample_ensemble``) can + omit it. loss : str, default "wmse" The scoring rule (from ``neural_lam.metrics``) used by ``compute_training_loss`` and stored as ``self.loss``. diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index 51df3261..85bb06d1 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -2,6 +2,7 @@ from argparse import Namespace # Third-party +import pytest import pytorch_lightning as pl import torch @@ -121,6 +122,47 @@ def test_ar_forecaster_score(): assert torch.equal(scored_explicit, expected_explicit) +def test_ar_forecaster_without_config_raises_on_use_not_construction(): + """A predictor that doesn't output std plus no config is a valid, + unambiguous state at construction time (the forecaster may only ever + be used for inference), so ARForecaster must not raise there. It + should only raise once scoring is actually attempted and has no + std to use, and the error should come from the forecaster itself, not + a wrapping module.""" + datastore = init_datastore_example("mdp") + predictor = MockStepPredictor(datastore=datastore, output_std=False) + + # Construction succeeds even though predicts_std=False and config=None + forecaster = ARForecaster(predictor, datastore) + assert forecaster.per_var_std is None + + B, num_grid_nodes = 2, predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + prediction = torch.zeros(B, num_grid_nodes, d_state) + target = torch.ones(B, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.score(prediction, target, None) + + num_past_forcing_steps = 1 + num_future_forcing_steps = 1 + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + pred_steps = 3 + init_states = torch.ones(B, 2, num_grid_nodes, d_state) + forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing) + true_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) + + with pytest.raises(ValueError, match="per_var_std fallback"): + forecaster.compute_training_loss( + init_states, + forcing_features, + true_states, + interior_mask_bool=torch.ones(num_grid_nodes, dtype=torch.bool), + ) + + def test_forecaster_module_checkpoint(tmp_path): datastore = init_datastore_example("mdp") From 7dc0b906d2615e3cfdcf56e3e3c0a386066b7fdf Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 18 Jul 2026 22:43:39 +0530 Subject: [PATCH 48/51] Address PR review: move ForecasterModules out of forecasters/ into modules/ BaseForecasterModule, DeterministicForecasterModule and ProbabilisticForecasterModule (Lightning wrappers around a Forecaster) were mixed in with the Forecaster classes themselves under neural_lam/models/forecasters/. Move them into their own neural_lam/models/modules/ package (base.py, deterministic.py, probabilistic.py) so the two concerns: what a forecaster is vs how it's trained/evaluated by Lightning live in separate directories. Pure move: internal imports and neural_lam/models/__init__.py updated accordingly, no behavioural change. --- neural_lam/models/__init__.py | 6 +++--- neural_lam/models/modules/__init__.py | 8 ++++++++ .../{forecasters/base_module.py => modules/base.py} | 2 +- .../deterministic_module.py => modules/deterministic.py} | 2 +- .../probabilistic_module.py => modules/probabilistic.py} | 4 ++-- tests/test_train_model_warnings.py | 2 +- 6 files changed, 16 insertions(+), 8 deletions(-) create mode 100644 neural_lam/models/modules/__init__.py rename neural_lam/models/{forecasters/base_module.py => modules/base.py} (99%) rename neural_lam/models/{forecasters/deterministic_module.py => modules/deterministic.py} (99%) rename neural_lam/models/{forecasters/probabilistic_module.py => modules/probabilistic.py} (98%) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cbeb1b01..be27cfec 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,13 +3,13 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster -from .forecasters.base_module import BaseForecasterModule -from .forecasters.deterministic_module import DeterministicForecasterModule from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, ) -from .forecasters.probabilistic_module import ProbabilisticForecasterModule +from .modules.base import BaseForecasterModule +from .modules.deterministic import DeterministicForecasterModule +from .modules.probabilistic import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_lam import GraphLAM diff --git a/neural_lam/models/modules/__init__.py b/neural_lam/models/modules/__init__.py new file mode 100644 index 00000000..b2e70700 --- /dev/null +++ b/neural_lam/models/modules/__init__.py @@ -0,0 +1,8 @@ +""" +Lightning modules wrapping forecasters for training and evaluation. +""" + +# Local +from .base import BaseForecasterModule +from .deterministic import DeterministicForecasterModule +from .probabilistic import ProbabilisticForecasterModule diff --git a/neural_lam/models/forecasters/base_module.py b/neural_lam/models/modules/base.py similarity index 99% rename from neural_lam/models/forecasters/base_module.py rename to neural_lam/models/modules/base.py index c77de238..0f239ce3 100644 --- a/neural_lam/models/forecasters/base_module.py +++ b/neural_lam/models/modules/base.py @@ -22,7 +22,7 @@ from ...config import NeuralLAMConfig from ...datastore import BaseDatastore from ...weather_dataset import WeatherDataset -from .base import Forecaster +from ..forecasters.base import Forecaster class BaseForecasterModule(pl.LightningModule, ABC): diff --git a/neural_lam/models/forecasters/deterministic_module.py b/neural_lam/models/modules/deterministic.py similarity index 99% rename from neural_lam/models/forecasters/deterministic_module.py rename to neural_lam/models/modules/deterministic.py index 217d4937..2842c8f2 100644 --- a/neural_lam/models/forecasters/deterministic_module.py +++ b/neural_lam/models/modules/deterministic.py @@ -12,7 +12,7 @@ # Local from ... import metrics, vis -from .base_module import BaseForecasterModule +from .base import BaseForecasterModule class DeterministicForecasterModule(BaseForecasterModule): diff --git a/neural_lam/models/forecasters/probabilistic_module.py b/neural_lam/models/modules/probabilistic.py similarity index 98% rename from neural_lam/models/forecasters/probabilistic_module.py rename to neural_lam/models/modules/probabilistic.py index cf8d38b0..bcb61395 100644 --- a/neural_lam/models/forecasters/probabilistic_module.py +++ b/neural_lam/models/modules/probabilistic.py @@ -8,8 +8,8 @@ # Local from ... import metrics -from .base_module import BaseForecasterModule -from .probabilistic import ProbabilisticForecaster +from ..forecasters.probabilistic import ProbabilisticForecaster +from .base import BaseForecasterModule class ProbabilisticForecasterModule(BaseForecasterModule): diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index 6bb0654f..44a9f9cb 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -77,7 +77,7 @@ def capture_init(_self, **kwargs): patch("neural_lam.train_model.MODELS", {"graph_lam": MagicMock()}), patch("neural_lam.train_model.ARForecaster"), patch( - "neural_lam.models.forecasters.deterministic_module." + "neural_lam.models.modules.deterministic." "DeterministicForecasterModule.__init__", capture_init, ), From e027608e34a6124a16f1d98a6dbf96068add8c15 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Tue, 21 Jul 2026 04:29:51 +0530 Subject: [PATCH 49/51] Address PR review: inline build_prior, move graph-type check to subclasses Three changes to BaseGraphEFM per review: - Inline build_prior's body into __init__ (it was only ever called once, from there) and remove the method. - Remove the expects_hierarchical class attribute and the generic hierarchical-vs-flat branch it drove in the base class; replace with an abstract latent_spatial_dim property that each subclass implements from its own knowledge of its graph shape (top mesh level for GraphEFM, all mesh nodes for GraphEFMMultiScale), used to size the constant prior. - Move the graph-type ValueError back out of the base class into the subclasses via a new check_graph_type hook, so the "hierarchical or flat" concern lives only in GraphEFM/GraphEFMMultiScale, not the base class. check_graph_type is a method the base class calls (not code left for the subclass constructor to run after super().__init__() returns): since the base class now builds the prior -- calling build_learnable_prior, which assumes edge_index tensors of the shape the subclass expects -- the check has to happen before that runs, not after. Doing it the latter way was tried first and left test_graph_type_mismatch_raises[GraphEFMMultiScale- hierarchical] hitting a RuntimeError deep in gnn_layers.py's InteractionNet (BufferList indexed as if it were a single edge_index tensor) instead of the intended ValueError, because build_learnable_prior ran during super().__init__(), before the subclass's post-super() check ever got a chance to fire. --- .../models/step_predictors/graph/graph_efm.py | 226 +++++++++++------- 1 file changed, 134 insertions(+), 92 deletions(-) diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 46dab6b9..9fc59c2e 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -39,20 +39,18 @@ class BaseGraphEFM(StepPredictor): does not compute any loss. This base class sets up everything that is independent of the mesh - graph type: it loads the graph, verifies it matches the type declared - by the subclass's ``expects_hierarchical`` class attribute, and builds - the prior (via :meth:`build_prior`, delegating to the subclass's - :meth:`build_learnable_prior` when learned). Concrete subclasses build - the mesh embedders and the encoder/decoder latent modules, and - implement :meth:`embedd_mesh` and :meth:`build_learnable_prior`. See - :class:`GraphEFM` (hierarchical graph) and :class:`GraphEFMMultiScale` - (flat graph). + graph type: it loads the graph, calls the subclass's + :meth:`check_graph_type` to verify the loaded graph matches what it + requires, then builds the prior, delegating to the subclass's + :meth:`build_learnable_prior` when learned and to + :attr:`latent_spatial_dim` for the constant prior's node count. + Concrete subclasses build the mesh embedders and the encoder/decoder + latent modules, and implement :meth:`embedd_mesh`, + :meth:`check_graph_type`, :meth:`build_learnable_prior` and + :attr:`latent_spatial_dim`. See :class:`GraphEFM` (hierarchical graph) + and :class:`GraphEFMMultiScale` (flat graph). """ - #: Set by concrete subclasses: whether they require a hierarchical - #: (True) or flat (False) mesh graph. - expects_hierarchical: bool - def __init__( self, datastore: BaseDatastore, @@ -139,23 +137,12 @@ def __init__( graph_name, mesh_node_features_scaling=grid_xy_max_span, ) - if self.hierarchical != self.expects_hierarchical: - expected_kind = ( - "hierarchical" if self.expects_hierarchical else "flat" - ) - actual_kind = "hierarchical" if self.hierarchical else "flat" - raise ValueError( - f"{type(self).__name__} requires a {expected_kind} mesh " - f"graph, but graph '{graph_name}' is {actual_kind}" - ) - - # The latent variable lives on the top mesh level for hierarchical - # graphs, and on every mesh node for flat graphs. - self.num_mesh_nodes = ( - self.mesh_static_features[-1].shape[0] - if self.hierarchical - else self.mesh_static_features.shape[0] - ) + # Delegated to the subclass, which knows whether it requires a + # hierarchical or flat mesh graph. Must run before anything below + # that assumes a specific graph shape (build_learnable_prior, + # latent_spatial_dim), so it cannot wait until the subclass + # constructor resumes after this call returns. + self.check_graph_type(graph_name) # Specify dimensions of data self.num_state_vars = datastore.get_num_data_vars(category="state") @@ -189,74 +176,66 @@ def __init__( # inert -- accepted for interface parity with other StepPredictors. self.prepare_clamping_params(datastore) - # Prior over the latent variable. + # Prior over the latent variable. When learn_prior is True the + # (graph-type specific) learnable prior is delegated to + # build_learnable_prior; otherwise the constant Normal(0, 1) prior, + # identical for every graph type, is built directly here. self.latent_dim = latent_dim if latent_dim is not None else hidden_dim - self.prior_model = self.build_prior( - learn_prior=learn_prior, - latent_dim=self.latent_dim, - hidden_dim=hidden_dim, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - prior_dist=prior_dist, - prior_layers=prior_layers, - ) + if learn_prior: + self.prior_model = self.build_learnable_prior( + latent_dim=self.latent_dim, + hidden_dim=hidden_dim, + hidden_layers=hidden_layers, + g2m_gnn_type=g2m_gnn_type, + prior_dist=prior_dist, + prior_layers=prior_layers, + ) + else: + self.prior_model = ConstantLatentEncoder( + latent_dim=self.latent_dim, + num_mesh_nodes=self.latent_spatial_dim, + output_dist=prior_dist, + ) - def build_prior( - self, - learn_prior, - latent_dim, - hidden_dim, - hidden_layers, - g2m_gnn_type, - prior_dist, - prior_layers, - ): + def check_graph_type(self, graph_name: str) -> None: """ - Build the prior over the latent variable. + Verify the loaded graph (``self.hierarchical``) is of the type this + predictor requires. - When ``learn_prior`` is True the (graph-type specific) learnable prior - is delegated to :meth:`build_learnable_prior`; otherwise the constant - ``Normal(0, 1)`` prior, which is identical for every graph type, is - built here. + Implemented by the concrete subclass, which is the only place that + knows whether it requires a hierarchical or flat mesh graph. Called + by this base class right after loading the graph, before anything + that assumes a specific graph shape. Parameters ---------- - learn_prior : bool - If True, build a learnable prior conditioned on the previous - state; if False, build a constant (input-independent) prior. - latent_dim : int - Dimensionality of the latent variable at each mesh node. - hidden_dim : int - Dimensionality of internal node and edge representations. - hidden_layers : int - Number of hidden layers in internal MLPs. - g2m_gnn_type : str - GNN type for the grid-to-mesh step of the learnable prior. - prior_dist : str - Output distribution of the prior: ``"isotropic"`` or - ``"diagonal"``. - prior_layers : int - Number of on-mesh GNN layers in the learnable prior. + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If ``self.hierarchical`` does not match what this predictor + requires. + """ + raise NotImplementedError("check_graph_type not implemented") + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes the latent variable lives on. + + Implemented by the concrete subclass, which knows the mesh graph + type: the top mesh level for hierarchical graphs, or every mesh + node for flat graphs. Returns ------- - torch.nn.Module - The prior latent encoder. + int + Number of latent-carrying mesh nodes. """ - if learn_prior: - return self.build_learnable_prior( - latent_dim=latent_dim, - hidden_dim=hidden_dim, - hidden_layers=hidden_layers, - g2m_gnn_type=g2m_gnn_type, - prior_dist=prior_dist, - prior_layers=prior_layers, - ) - return ConstantLatentEncoder( - latent_dim=latent_dim, - num_mesh_nodes=self.num_mesh_nodes, - output_dist=prior_dist, - ) + raise NotImplementedError("latent_spatial_dim not implemented") def build_learnable_prior( self, @@ -472,8 +451,6 @@ class GraphEFM(BaseGraphEFM): decoder is a ``HiGraphLatentDecoder``. """ - expects_hierarchical = True - def __init__( self, datastore: BaseDatastore, @@ -653,6 +630,40 @@ def __init__( output_std=bool(output_std), ) + def check_graph_type(self, graph_name: str) -> None: + """ + Verify the loaded graph is hierarchical. + + Parameters + ---------- + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If the loaded graph is flat. + """ + if not self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a hierarchical mesh " + f"graph, but graph '{graph_name}' is flat" + ) + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes on the top mesh level, where the latent + variable lives. + + Returns + ------- + int + Number of top-level mesh nodes. + """ + return self.mesh_static_features[-1].shape[0] + def build_learnable_prior( self, latent_dim, @@ -758,8 +769,6 @@ class GraphEFMMultiScale(BaseGraphEFM): ``GraphLatentDecoder``. """ - expects_hierarchical = False - def __init__( self, datastore: BaseDatastore, @@ -852,8 +861,8 @@ def __init__( 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)" + f"{self.num_grid_nodes + self.latent_spatial_dim} nodes " + f"({self.num_grid_nodes} grid, {self.latent_spatial_dim} mesh)" ) # Embedders @@ -889,6 +898,39 @@ def __init__( output_std=bool(output_std), ) + def check_graph_type(self, graph_name: str) -> None: + """ + Verify the loaded graph is flat. + + Parameters + ---------- + graph_name : str + Name of the graph directory that was loaded, for the error + message. + + Raises + ------ + ValueError + If the loaded graph is hierarchical. + """ + if self.hierarchical: + raise ValueError( + f"{type(self).__name__} requires a flat mesh graph, " + f"but graph '{graph_name}' is hierarchical" + ) + + @property + def latent_spatial_dim(self) -> int: + """ + Number of mesh nodes, where the latent variable lives. + + Returns + ------- + int + Number of mesh nodes. + """ + return len(self.mesh_static_features) + def build_learnable_prior( self, latent_dim, From 780941f7f3b4a510c658b703e2a23bc7c749ec07 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Tue, 21 Jul 2026 05:44:42 +0530 Subject: [PATCH 50/51] Address PR review: remove now-unnecessary TYPE_CHECKING guard on BaseDatastore utils/graph.py's TYPE_CHECKING-guarded BaseDatastore import predates the utils/ package split (#682); this checks whether it's still needed now that it's merged. It still was, but not for the reason the guard's comment implied. The cycle isn't the monolithic-vs-package structure -- it's datastore/mdp.py importing log_on_rank_zero from the utils *package* (`from ..utils import log_on_rank_zero`), which only resolves once utils/__init__.py has fully run. utils/__init__.py imports .graph before .logging, so an eager BaseDatastore import in graph.py (which pulls in the datastore package, which imports mdp.py) hits log_on_rank_zero before it's bound. Fix: import log_on_rank_zero from the utils.logging *submodule* directly in mdp.py, not the package. That resolves independently of utils/__init__.py's progress, breaking the cycle without relying on import order (unlike reordering utils/__init__.py, which also works but is fragile and silently reintroducible). The TYPE_CHECKING guard is no longer needed. --- neural_lam/datastore/mdp.py | 2 +- neural_lam/utils/graph.py | 12 ++++-------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 7cad45d7..24a508c9 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -18,7 +18,7 @@ from numpy import ndarray # Local -from ..utils import log_on_rank_zero +from ..utils.logging import log_on_rank_zero from .base import BaseRegularGridDatastore, CartesianGridShape diff --git a/neural_lam/utils/graph.py b/neural_lam/utils/graph.py index a843afc3..009e747f 100644 --- a/neural_lam/utils/graph.py +++ b/neural_lam/utils/graph.py @@ -4,7 +4,7 @@ import os import warnings from pathlib import Path -from typing import TYPE_CHECKING, Any, Union +from typing import Any, Union # Third-party import torch @@ -12,13 +12,9 @@ from torch import nn # Local +from ..datastore import BaseDatastore from .buffer_list import BufferList -if TYPE_CHECKING: - # Imported only for type checking to avoid a runtime import cycle - # Local - from ..datastore import BaseDatastore - LEGACY_GRAPH_SPEC_VERSION = "legacy" @@ -428,7 +424,7 @@ def load_graph_spec_version() -> str: def load_and_register_graph( module: nn.Module, - datastore: "BaseDatastore", + datastore: BaseDatastore, graph_name: str, mesh_node_features_scaling: float, ) -> bool: @@ -472,7 +468,7 @@ def load_and_register_graph( def compute_grid_input_dim( - datastore: "BaseDatastore", + datastore: BaseDatastore, num_past_forcing_steps: int, num_future_forcing_steps: int, ) -> int: From d18a09b0d90b9341bec0db23e5ecc61f9e5f9973 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sat, 25 Jul 2026 09:15:44 +0530 Subject: [PATCH 51/51] feat: instantiate Graph-EFM model (ELBO forecaster + probabilistic config/CLI) Add GraphEFMForecaster, which trains the GraphEFM/GraphEFMMultiScale latent step predictors via their ELBO through the predictors' step_distributions interface, with the static graph embedding hoisted out of the rollout. Add a probabilistic config section (ProbabilisticConfig) and config-aware train_model wiring registering graph_efm (hierarchical) and graph_efm_ms (flat). Move the decoder residual into the predictor so the predicted mean is clamped to the configured output range like the deterministic models, and use an empty m2m placeholder when no intra-level layers are configured. --- neural_lam/config.py | 51 +- neural_lam/models/__init__.py | 16 +- neural_lam/models/forecasters/__init__.py | 1 + neural_lam/models/forecasters/graph_efm.py | 208 ++++++++ neural_lam/models/latent/base_decoder.py | 29 +- neural_lam/models/latent/hi_graph_decoder.py | 4 +- .../models/step_predictors/graph/graph_efm.py | 157 ++++-- neural_lam/train_model.py | 201 ++++++-- tests/test_graph_efm_model.py | 469 ++++++++++++++++++ tests/test_latent_modules.py | 53 +- 10 files changed, 1044 insertions(+), 145 deletions(-) create mode 100644 neural_lam/models/forecasters/graph_efm.py create mode 100644 tests/test_graph_efm_model.py diff --git a/neural_lam/config.py b/neural_lam/config.py index 1da43fff..37f39714 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -3,7 +3,7 @@ # Standard library import dataclasses from pathlib import Path -from typing import Dict, Union +from typing import Dict, Optional, Union # Third-party import dataclass_wizard @@ -116,6 +116,49 @@ class TrainingConfig: ) +@dataclasses.dataclass +class ProbabilisticConfig: + """ + Configuration for the probabilistic (Graph-EFM) models. + + Only used when training or evaluating a probabilistic model + (``--model graph_efm``/``graph_efm_ms``); ignored otherwise. Every field + has a default, so the ``probabilistic`` config section may be omitted + entirely for deterministic models. + + Attributes + ---------- + latent_dim : int, optional + Dimensionality of the latent variable at each latent-carrying mesh + node. Defaults to the model's ``hidden_dim`` when None. + prior_layers : int + Number of on-mesh GNN layers in the prior. + encoder_layers : int + Number of on-mesh GNN layers in the variational encoder. + decoder_layers : int + Number of on-mesh GNN layers in the latent decoder. + learn_prior : bool + If True, the prior is a learned encoder conditioned on the previous + state; if False, a constant ``Normal(0, 1)`` prior is used. + prior_dist : str + Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. + kl_beta : float + Weight of the KL term in the ELBO. When 0, the prior and KL are not + computed (pure auto-encoder training). + eval_ensemble_size : int + Number of ensemble members sampled during validation and testing. + """ + + latent_dim: Optional[int] = None + prior_layers: int = 2 + encoder_layers: int = 2 + decoder_layers: int = 4 + learn_prior: bool = True + prior_dist: str = "isotropic" + kl_beta: float = 1.0 + eval_ensemble_size: int = 5 + + @dataclasses.dataclass class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): """ @@ -133,10 +176,16 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): training : TrainingConfig Configuration for training the model, including loss function and feature-weighting strategy. Defaults to ``TrainingConfig()``. + probabilistic : ProbabilisticConfig + Configuration for the probabilistic (Graph-EFM) models. Defaults to + ``ProbabilisticConfig()`` and is ignored by deterministic models. """ datastore: DatastoreSelection training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig) + probabilistic: ProbabilisticConfig = dataclasses.field( + default_factory=ProbabilisticConfig + ) class _(dataclass_wizard.JSONWizard.Meta): """ diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index c93c1e9c..dadcee46 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,6 +3,7 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.graph_efm import GraphEFMForecaster from .forecasters.probabilistic import ( ProbabilisticARForecaster, ProbabilisticForecaster, @@ -18,15 +19,16 @@ from .step_predictors.graph.hi_lam_parallel import HiLAMParallel from .step_predictors.graph.hierarchical import BaseHiGraphModel -# NOTE: GraphEFM/GraphEFMMultiScale 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. +# Graph-EFM models are probabilistic: train_model.py builds them with a +# config-aware, probabilistic assembly path (GraphEFMForecaster wrapped in a +# ProbabilisticForecasterModule), distinct from the deterministic models +# above. ``PROBABILISTIC_MODELS`` marks which entries take that path. MODELS = { "graph_lam": GraphLAM, "hi_lam": HiLAM, "hi_lam_parallel": HiLAMParallel, + "graph_efm": GraphEFM, + "graph_efm_ms": GraphEFMMultiScale, } + +PROBABILISTIC_MODELS = {"graph_efm", "graph_efm_ms"} diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index 254c4ba0..3c2093de 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -5,4 +5,5 @@ # Local from .autoregressive import ARForecaster from .base import Forecaster +from .graph_efm import GraphEFMForecaster from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster diff --git a/neural_lam/models/forecasters/graph_efm.py b/neural_lam/models/forecasters/graph_efm.py new file mode 100644 index 00000000..84a71b69 --- /dev/null +++ b/neural_lam/models/forecasters/graph_efm.py @@ -0,0 +1,208 @@ +"""Probabilistic forecaster training a Graph-EFM predictor via its ELBO.""" + +# Third-party +import torch + +# Local +from ...config import NeuralLAMConfig +from ...datastore import BaseDatastore +from ..step_predictors.graph.graph_efm import BaseGraphEFM +from .probabilistic import ProbabilisticARForecaster + + +class GraphEFMForecaster(ProbabilisticARForecaster): + """ + Auto-regressive ensemble forecaster for Graph-EFM step predictors. + + Wraps a :class:`BaseGraphEFM` predictor (hierarchical ``GraphEFM`` or + flat ``GraphEFMMultiScale``), a latent-variable model consisting of a + conditional prior, a variational encoder and a latent decoder. Forecast + sampling (``forward``, ``sample_ensemble``) is inherited from + :class:`ProbabilisticARForecaster`: each predictor call samples the + prior and decodes, so unrolling produces one stochastic trajectory and + stacking several gives an ensemble. + + This class supplies the training objective the base class leaves + abstract: the evidence lower bound (ELBO). ``compute_training_loss`` + runs a variational rollout in which, at each step, the latent is drawn + from the encoder (conditioned on the target), and accumulates a + reconstruction likelihood term and a KL term between the encoder and the + prior. The scoring rule ``self.loss``, the constant per-variable std + fallback ``self.per_var_std`` and the boundary/interior masks all live + on the forecaster (set up by :class:`ARForecaster`); the predictor only + provides the network building blocks. + """ + + def __init__( + self, + predictor: BaseGraphEFM, + datastore: BaseDatastore, + config: NeuralLAMConfig | None = None, + loss: str = "wmse", + kl_beta: float = 1.0, + ) -> None: + """ + Initialize the GraphEFMForecaster. + + Parameters + ---------- + predictor : BaseGraphEFM + The Graph-EFM step predictor to use for each step. Each + ``forward`` call samples the prior and decodes; the encoder, + prior and decoder sub-models are also used directly by + ``compute_training_loss`` to assemble the ELBO. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + config : NeuralLAMConfig or None + Configuration used to compute the constant per-variable std + substituted for ``pred_std`` when ``predictor`` does not output + its own (see ``ARForecaster.per_var_std``). Required for + training when the predictor's ``output_std`` is False, since the + likelihood term then needs the fallback std. + loss : str, default "wmse" + The scoring rule (from ``neural_lam.metrics``) used for the + reconstruction likelihood term and stored as ``self.loss``. + kl_beta : float, default 1.0 + Weight of the KL term in the ELBO. When ``0`` the prior and KL + are not computed at all (pure auto-encoder training); the prior + network then receives no gradient. + """ + super().__init__(predictor, datastore, config=config, loss=loss) + self.kl_beta = kl_beta + + def compute_training_loss( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + interior_mask_bool: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the ELBO training objective for one batch. + + Unrolls a variational rollout over the full forecast: at each step + the grid, graph and target are embedded, the latent is drawn from + the encoder (variational posterior), the decoder reconstructs the + next-state mean, and a reconstruction likelihood term is + accumulated. When ``kl_beta`` is positive a KL term between the + encoder and the prior is accumulated too. Both terms are summed over + the rollout and averaged over the batch; the loss is + ``-likelihood + kl_beta * kl``. The rollout advances on the + predicted mean with boundary nodes overwritten by the true state. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the rollout length, + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already + concatenated past/current/future windows). + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + states at each predicted step, used as the encoder conditioning + and reconstruction target and to overwrite boundary nodes during + the rollout. Dims: same as the prediction. + interior_mask_bool : torch.Tensor + Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior + nodes; passed as ``mask`` to ``self.loss`` so only interior + nodes contribute to the likelihood. + + Returns + ------- + batch_loss : torch.Tensor + Scalar. The negative ELBO for the batch, to take gradients of. + loss_components : dict of {str: torch.Tensor} + Scalar ELBO diagnostics to log: ``"elbo_likelihood"`` always, + plus ``"elbo_kl"`` and ``"elbo"`` when ``kl_beta > 0``. + + Raises + ------ + ValueError + If the predictor does not output its own std and no + ``per_var_std`` fallback is available (this forecaster was + constructed without ``config``); see ``_resolve_pred_std``. + """ + predictor = self.predictor + + prev_prev_state = init_states[:, 0] + prev_state = init_states[:, 1] + pred_steps = forcing_features.shape[1] + compute_kl = self.kl_beta > 0 + + # The graph embedding depends only on static features, so it is + # constant across the rollout and computed once here. + graph_emb = predictor.embedd_graph(init_states.shape[0]) + + likelihood_terms = [] + kl_terms = [] + + for i in range(pred_steps): + forcing = forcing_features[:, i] + target_state = target_states[:, i] + + # Posterior latent (conditioned on target), reconstruction, and + # -- when a KL term is needed -- the prior, in one predictor call. + prior_dist, posterior_dist, pred_mean, pred_std = ( + predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=compute_kl, + ) + ) + pred_std = self._resolve_pred_std(pred_std) + + # Reconstruction likelihood, summed over interior grid and vars + entry_likelihoods = -self.loss( + pred_mean, + target_state, + pred_std, + mask=interior_mask_bool, + average_grid=False, + sum_vars=False, + ) # (B, num_interior_grid_nodes, num_state_vars) + likelihood_terms.append(torch.sum(entry_likelihoods, dim=(1, 2))) + + if compute_kl: + kl_terms.append( + torch.sum( + torch.distributions.kl_divergence( + posterior_dist, prior_dist + ), + dim=(1, 2), + ) + ) # (B,) + + # Advance the rollout on the predicted mean, boundary overwritten + new_state = ( + self.boundary_mask * target_state + + self.interior_mask * pred_mean + ) + prev_prev_state = prev_state + prev_state = new_state + + # Sum each term over the rollout, then average over the batch + mean_likelihood = torch.mean( + torch.sum(torch.stack(likelihood_terms, dim=1), dim=1) + ) + loss_components = {"elbo_likelihood": mean_likelihood} + + if compute_kl: + mean_kl = torch.mean(torch.sum(torch.stack(kl_terms, dim=1), dim=1)) + batch_loss = -mean_likelihood + self.kl_beta * mean_kl + loss_components["elbo_kl"] = mean_kl + loss_components["elbo"] = mean_likelihood - mean_kl + else: + batch_loss = -mean_likelihood + + return batch_loss, loss_components diff --git a/neural_lam/models/latent/base_decoder.py b/neural_lam/models/latent/base_decoder.py index bd19b111..dbd27e58 100644 --- a/neural_lam/models/latent/base_decoder.py +++ b/neural_lam/models/latent/base_decoder.py @@ -10,12 +10,12 @@ 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. + mesh to the next-state increment (and optionally std) on the grid. Subclasses implement :meth:`combine_with_latent`, which fuses the latent representation with the grid representation. The resulting features are - mapped to either ``num_state_vars`` outputs (mean only) or - ``2 * num_state_vars`` outputs (mean, std) depending on + mapped to either ``num_state_vars`` outputs (mean increment only) or + ``2 * num_state_vars`` outputs (mean increment, std) depending on ``output_std``. """ @@ -107,14 +107,15 @@ def combine_with_latent( """ raise NotImplementedError("combine_with_latent not implemented") - def forward(self, grid_rep, latent_samples, last_state, graph_emb): + def forward(self, grid_rep, latent_samples, graph_emb): """ - Predict mean (and optionally std) of the next weather state. + Predict the next-state increment (and optionally std). The latent samples are embedded to the internal dimensionality and fused with the grid representation by ``combine_with_latent``; the - result is mapped to distribution parameters. The mean is predicted - as a residual on top of ``last_state``. + result is mapped to distribution parameters. The predicted mean is an + increment relative to the current state, which the caller adds onto + the current state (and clamps). Parameters ---------- @@ -125,10 +126,6 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): 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``, @@ -136,9 +133,9 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): Returns ------- - pred_mean : torch.Tensor - Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted mean - of the next state. + mean_delta : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted + increment of the next state relative to the current state. pred_std : torch.Tensor or None Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is True, otherwise None. Predicted std of the @@ -171,6 +168,4 @@ def forward(self, grid_rep, latent_samples, last_state, graph_emb): mean_delta = state_params pred_std = None - pred_mean = last_state + mean_delta - - return pred_mean, pred_std + return mean_delta, pred_std diff --git a/neural_lam/models/latent/hi_graph_decoder.py b/neural_lam/models/latent/hi_graph_decoder.py index 63e8344d..83a9fede 100644 --- a/neural_lam/models/latent/hi_graph_decoder.py +++ b/neural_lam/models/latent/hi_graph_decoder.py @@ -230,10 +230,10 @@ def combine_with_latent( ) ): new_mesh_rep = current_mesh_rep - new_m2m_rep = graph_emb["m2m"][level] + new_m2m_rep = None if self.intra_up_gnns is not None: new_mesh_rep, new_m2m_rep = self.intra_up_gnns[level]( - new_mesh_rep, new_m2m_rep + new_mesh_rep, graph_emb["m2m"][level] ) # Saved for residual connections in the downward pass diff --git a/neural_lam/models/step_predictors/graph/graph_efm.py b/neural_lam/models/step_predictors/graph/graph_efm.py index 9fc59c2e..55424b9f 100644 --- a/neural_lam/models/step_predictors/graph/graph_efm.py +++ b/neural_lam/models/step_predictors/graph/graph_efm.py @@ -171,9 +171,8 @@ def __init__( self.g2m_embedder = utils.make_mlp([g2m_dim] + self.mlp_blueprint_end) self.m2g_embedder = utils.make_mlp([m2g_dim] + self.mlp_blueprint_end) - # Compute indices and define 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. + # Compute indices and define the clamping functions applied to the + # predicted next-state mean in step_distributions. self.prepare_clamping_params(datastore) # Prior over the latent variable. When learn_prior is True the @@ -338,9 +337,9 @@ def embedd_mesh(self, batch_size): """ raise NotImplementedError("embedd_mesh not implemented") - def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): + def embedd_grid(self, prev_state, prev_prev_state, forcing): """ - Embed the grid (states up to t-1) and the full graph. + Embed the grid representation of the states up to t-1. Parameters ---------- @@ -353,10 +352,8 @@ def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): Returns ------- - grid_emb : torch.Tensor + 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] @@ -370,13 +367,27 @@ def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): dim=-1, ) # (B, num_grid_nodes, grid_dim) - grid_emb = self.grid_prev_embedder(grid_features) + return 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. + def embedd_graph(self, batch_size): + """ + Embed the static grid-mesh edge and mesh graph features. + + The embedding depends only on static graph features and is therefore + constant across an autoregressive rollout, unlike :meth:`embedd_grid`. + + Parameters + ---------- + batch_size : int + Batch size to expand the embeddings to. + + Returns + ------- + dict + Edge/mesh embeddings, each entry of shape ``(B, *, d_h)`` (``g2m``, + ``m2g`` and the entries added by :meth:`embedd_mesh`). + """ graph_emb = { "g2m": self.expand_to_batch( self.g2m_embedder(self.g2m_features), batch_size @@ -387,7 +398,90 @@ def embedd_grid_and_graph(self, prev_state, prev_prev_state, forcing): } graph_emb.update(self.embedd_mesh(batch_size)) - return grid_emb, graph_emb + return graph_emb + + def step_distributions( + self, + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=None, + compute_prior=True, + ): + """ + Compute the latent distributions and next-state prediction for a step. + + Embeds the grid, then draws the latent either from the variational + posterior (when ``target_state`` is given -- the training path) or + from the prior (inference), and decodes that sample into the + next-state mean (and, when ``output_std``, std) as a residual on + ``prev_state``. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, d_forcing)``. + graph_emb : dict + Static graph embedding as returned by :meth:`embedd_graph`. + target_state : torch.Tensor, optional + Shape ``(B, num_grid_nodes, d_state)``. Target ``X_{t+1}``. When + given, the latent is sampled from the variational posterior + conditioned on it; when None, from the prior. + compute_prior : bool + On the posterior (training) path, whether to also compute the + prior distribution (needed for a KL term). Ignored on the + inference path, where the prior is always computed since it is the + sampling distribution. + + Returns + ------- + prior_dist : torch.distributions.Normal or None + The prior over the latent, or None on the posterior path when + ``compute_prior`` is False. + posterior_dist : torch.distributions.Normal or None + The variational posterior over the latent, or None on the + inference path (``target_state`` is None). + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, d_state)``. Decoder mean of + ``X_{t+1}``. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, d_state)`` when ``output_std`` is True, + otherwise None. + """ + grid_prev_emb = self.embedd_grid(prev_state, prev_prev_state, forcing) + + prior_dist = None + posterior_dist = None + if target_state is not None: + grid_current_emb = self.embedd_grid_with_target( + prev_state, prev_prev_state, forcing, target_state + ) + posterior_dist = self.encoder(grid_current_emb, graph_emb=graph_emb) + latent_samples = posterior_dist.rsample() + if compute_prior: + prior_dist = self.prior_model( + grid_prev_emb, graph_emb=graph_emb + ) + else: + prior_dist = self.prior_model(grid_prev_emb, graph_emb=graph_emb) + latent_samples = prior_dist.rsample() + # (B, num_mesh_nodes, d_latent) + + # Decode the latent into a state increment, then add it onto prev_state + # (X_t) and clamp to the valid range (a no-op when no clamping limits + # are configured), as for the deterministic models. + mean_delta, pred_std = self.decoder( + grid_prev_emb, latent_samples, graph_emb + ) + pred_mean = self.get_clamped_new_state(mean_delta, prev_state) + # (B, num_grid_nodes, d_state) + + return prior_dist, posterior_dist, pred_mean, pred_std def forward( self, @@ -396,10 +490,11 @@ def forward( 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. + Sample one time step prediction from the prior. + + Embeds the graph and grid, samples the latent from the prior, decodes + it and returns the predicted next state. The prediction is stochastic + only through the latent sample; no observation noise is added. Parameters ---------- @@ -419,26 +514,10 @@ def forward( 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 + graph_emb = self.embedd_graph(prev_state.shape[0]) + _, _, pred_mean, pred_std = self.step_distributions( + prev_state, prev_prev_state, forcing, graph_emb, target_state=None ) - - # 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 @@ -753,8 +832,8 @@ def embedd_mesh(self, 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) + # No intra-level GNNs consume these, so no embedding is produced + mesh_emb["m2m"] = [] return mesh_emb diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 04999cb3..80590688 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -19,7 +19,14 @@ from . import utils from .config import load_config_and_datastore from .gnn_layers import GNN_TYPES -from .models import MODELS, ARForecaster, DeterministicForecasterModule +from .models import ( + MODELS, + PROBABILISTIC_MODELS, + ARForecaster, + DeterministicForecasterModule, + GraphEFMForecaster, + ProbabilisticForecasterModule, +) from .weather_dataset import WeatherDataModule @@ -38,39 +45,155 @@ def __init__(self, prog): ) -def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): +def build_predictor(args, config, datastore): """ - Reconstruct a DeterministicForecasterModule from a checkpoint without - requiring the caller to know the original architecture kwargs. - - The checkpoint must have been saved with args in hyper_parameters (i.e. - created via train_model.main), so that model class and architecture kwargs - can be recovered automatically. + Construct the step predictor for ``args.model``. + + Graph-EFM models (see ``PROBABILISTIC_MODELS``) are latent-variable + predictors whose constructor differs from the deterministic graph + models: they take ``latent_dim``/``learn_prior``/``prior_dist`` and + per-graph-type latent layer counts instead of ``processor_layers`` and + ``mesh_aggr``. This function supplies the right kwargs for each. + + Parameters + ---------- + args : argparse.Namespace + Parsed command-line arguments (see ``main``). + config : NeuralLAMConfig + Loaded neural-lam configuration, for the output-clamping limits. + datastore : BaseDatastore + Datastore providing static features and variable counts. + + Returns + ------- + StepPredictor + The constructed step predictor. """ - ckpt = torch.load(ckpt_path, weights_only=False) - args = ckpt["hyper_parameters"]["args"] predictor_class = MODELS[args.model] - predictor = predictor_class( + common_kwargs = dict( datastore=datastore, graph_name=args.graph, hidden_dim=args.hidden_dim, hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, num_past_forcing_steps=args.num_past_forcing_steps, num_future_forcing_steps=args.num_future_forcing_steps, output_std=args.output_std, output_clamping_lower=config.training.output_clamping.lower, output_clamping_upper=config.training.output_clamping.upper, + g2m_gnn_type=args.g2m_gnn_type, + m2g_gnn_type=args.m2g_gnn_type, + ) + + if args.model in PROBABILISTIC_MODELS: + prob = config.probabilistic + # Graph-EFM latent layer counts are named per graph type: + # intra-level for the hierarchical model, m2m for the flat one. + if args.model == "graph_efm": + layer_kwargs = dict( + prior_intra_level_layers=prob.prior_layers, + encoder_intra_level_layers=prob.encoder_layers, + decoder_intra_level_layers=prob.decoder_layers, + ) + else: + layer_kwargs = dict( + prior_m2m_layers=prob.prior_layers, + encoder_m2m_layers=prob.encoder_layers, + decoder_m2m_layers=prob.decoder_layers, + ) + return predictor_class( + **common_kwargs, + latent_dim=prob.latent_dim, + learn_prior=prob.learn_prior, + prior_dist=prob.prior_dist, + **layer_kwargs, + ) + + return predictor_class( + **common_kwargs, + processor_layers=args.processor_layers, + mesh_aggr=args.mesh_aggr, + mesh_up_gnn_type=args.mesh_up_gnn_type, + mesh_down_gnn_type=args.mesh_down_gnn_type, ) + + +def build_forecaster_module(args, config, datastore, predictor): + """ + Wrap a predictor in its forecaster and pick the Lightning module class. + + Graph-EFM models are trained probabilistically (``GraphEFMForecaster`` + optimizing the ELBO, evaluated as an ensemble by + ``ProbabilisticForecasterModule``); the other models use the + deterministic ``ARForecaster``/``DeterministicForecasterModule`` path. + The best-checkpoint monitor differs accordingly. + + Parameters + ---------- + args : argparse.Namespace + Parsed command-line arguments (see ``main``). + config : NeuralLAMConfig + Loaded neural-lam configuration. + datastore : BaseDatastore + Datastore providing grid metadata and boundary masks. + predictor : StepPredictor + The step predictor to wrap, as built by ``build_predictor``. + + Returns + ------- + forecaster : Forecaster + The forecaster wrapping ``predictor``. + module_class : type + The ``BaseForecasterModule`` subclass to instantiate. + module_kwargs : dict + Extra keyword arguments for ``module_class`` beyond the shared ones + (e.g. ``eval_ensemble_size`` for the probabilistic module). + val_monitor : str + Name of the validation metric to monitor for the best checkpoint. + """ + if args.model in PROBABILISTIC_MODELS: + forecaster = GraphEFMForecaster( + predictor, + datastore, + config=config, + loss=args.loss, + kl_beta=config.probabilistic.kl_beta, + ) + return ( + forecaster, + ProbabilisticForecasterModule, + {"eval_ensemble_size": config.probabilistic.eval_ensemble_size}, + "val_mean_ens_rmse", + ) + forecaster = ARForecaster( predictor, datastore, config=config, loss=args.loss ) - return DeterministicForecasterModule.load_from_checkpoint( + return forecaster, DeterministicForecasterModule, {}, "val_mean_loss" + + +def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): + """ + Reconstruct a forecaster module from a checkpoint without requiring the + caller to know the original architecture kwargs. + + The checkpoint must have been saved with args in hyper_parameters (i.e. + created via train_model.main), so that model class and architecture kwargs + can be recovered automatically. Deterministic and Graph-EFM + (probabilistic) checkpoints are both supported; the correct forecaster + and module class are chosen from ``args.model``. + """ + ckpt = torch.load(ckpt_path, weights_only=False) + args = ckpt["hyper_parameters"]["args"] + predictor = build_predictor(args, config, datastore) + forecaster, module_class, module_kwargs, _ = build_forecaster_module( + args, config, datastore, predictor + ) + return module_class.load_from_checkpoint( ckpt_path, forecaster=forecaster, datastore=datastore, weights_only=False, + **module_kwargs, ) @@ -213,6 +336,11 @@ def main(input_args=None): "hard-codes its mesh-down GNN type", ) + # Probabilistic / Graph-EFM hyperparameters (latent_dim, kl_beta, + # eval_ensemble_size, ...) are not CLI flags: they live in the + # ``probabilistic`` section of the neural-lam config (see + # ``ProbabilisticConfig``), read by build_predictor/build_forecaster_module. + # Training options train_group = parser.add_argument_group("Training Options") train_group.add_argument( @@ -461,31 +589,15 @@ def main(input_args=None): except ValueError: raise ValueError("devices should be 'auto' or a list of integers") - # Build predictor and forecaster externally, then inject into - # DeterministicForecasterModule - predictor_class = MODELS[args.model] - predictor = predictor_class( - datastore=datastore, - graph_name=args.graph, - hidden_dim=args.hidden_dim, - hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, - num_past_forcing_steps=args.num_past_forcing_steps, - num_future_forcing_steps=args.num_future_forcing_steps, - output_std=args.output_std, - output_clamping_lower=config.training.output_clamping.lower, - output_clamping_upper=config.training.output_clamping.upper, - g2m_gnn_type=args.g2m_gnn_type, - m2g_gnn_type=args.m2g_gnn_type, - mesh_up_gnn_type=args.mesh_up_gnn_type, - mesh_down_gnn_type=args.mesh_down_gnn_type, - ) - forecaster = ARForecaster( - predictor, datastore, config=config, loss=args.loss + # Build predictor and forecaster externally, then inject into the + # forecaster module. Graph-EFM models take a probabilistic assembly path + # (see build_predictor/build_forecaster_module); the others deterministic. + predictor = build_predictor(args, config, datastore) + forecaster, module_class, module_kwargs, val_monitor = ( + build_forecaster_module(args, config, datastore, predictor) ) - model = DeterministicForecasterModule( + model = module_class( forecaster=forecaster, config=config, datastore=datastore, @@ -498,6 +610,7 @@ def main(input_args=None): metrics_watch=args.metrics_watch, var_leads_metrics_watch=args.var_leads_metrics_watch, args=args, + **module_kwargs, ) if args.eval: @@ -525,8 +638,8 @@ def main(input_args=None): # checkpoint instead of losing all progress since the last validation. val_checkpoint = pl.callbacks.ModelCheckpoint( dirpath=os.path.join(run_dir, "checkpoints"), - filename="min_val_loss", - monitor="val_mean_loss", + filename=f"min_{val_monitor}", + monitor=val_monitor, mode="min", save_top_k=1, save_on_train_epoch_end=False, @@ -540,11 +653,19 @@ def main(input_args=None): save_on_train_epoch_end=True, enable_version_counter=False, ) + # With kl_beta == 0 the Graph-EFM prior network is never used in the loss, + # so multi-device DDP must be told to expect unused parameters. + strategy = ( + "ddp_find_unused_parameters_true" + if args.model in PROBABILISTIC_MODELS + and config.probabilistic.kl_beta == 0 + else "auto" + ) trainer = pl.Trainer( max_epochs=args.epochs, deterministic=True, default_root_dir=run_dir, - strategy="auto", + strategy=strategy, accelerator=device_name, num_nodes=args.num_nodes, devices=devices, diff --git a/tests/test_graph_efm_model.py b/tests/test_graph_efm_model.py new file mode 100644 index 00000000..e90a8036 --- /dev/null +++ b/tests/test_graph_efm_model.py @@ -0,0 +1,469 @@ +"""Integration tests for the full Graph-EFM model. + +Exercises the pieces that turn the Graph-EFM single-step predictors into a +trainable model: the ``GraphEFMForecaster`` ELBO objective, its inherited +ensemble sampling, the ``ProbabilisticForecasterModule`` wrapping, and the +config-aware assembly path in ``train_model`` (``build_predictor`` / +``build_forecaster_module``). Predictors are built on the real example +datastore with a freshly created graph, mirroring +``tests/test_graph_efm_predictor.py``. +""" + +# Standard library +from argparse import Namespace +from pathlib import Path + +# Third-party +import pytest +import torch + +# First-party +from neural_lam import config as nlconfig +from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.models import ( + GraphEFM, + GraphEFMForecaster, + GraphEFMMultiScale, + ProbabilisticForecasterModule, +) +from neural_lam.train_model import build_forecaster_module, build_predictor +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 + + +def _datastore_and_config(graph_name): + """ + Build the example datastore + config and ensure ``graph_name`` exists. + + Parameters + ---------- + graph_name : str + Graph directory name; ``"hierarchical"`` builds a multi-level graph, + anything else a flat one. + + Returns + ------- + datastore : BaseDatastore + The example ``mdp`` datastore. + config : NeuralLAMConfig + A configuration selecting that datastore. + """ + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + + hierarchical = graph_name == "hierarchical" + n_max_levels = 3 if hierarchical else 1 + graph_dir_path = Path(datastore.root_path) / "graph" / graph_name + if not graph_dir_path.exists(): + create_graph_from_datastore( + datastore=datastore, + output_root_path=str(graph_dir_path), + hierarchical=hierarchical, + n_max_levels=n_max_levels, + ) + return datastore, config + + +def _build_predictor(datastore, graph_name, output_std=False): + """ + Construct a small Graph-EFM predictor for ``graph_name``. + + Parameters + ---------- + datastore : BaseDatastore + Datastore to build the predictor on. + graph_name : str + Graph directory name selecting the flat vs hierarchical variant. + output_std : bool, default False + Whether the decoder outputs its own std. + + Returns + ------- + BaseGraphEFM + The constructed predictor. + """ + if graph_name == "hierarchical": + predictor_class = GraphEFM + layer_kwargs = { + "prior_intra_level_layers": 1, + "encoder_intra_level_layers": 1, + "decoder_intra_level_layers": 1, + } + else: + predictor_class = GraphEFMMultiScale + layer_kwargs = { + "prior_m2m_layers": 1, + "encoder_m2m_layers": 1, + "decoder_m2m_layers": 1, + } + return predictor_class( + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + learn_prior=True, + prior_dist="isotropic", + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + output_std=output_std, + **layer_kwargs, + ) + + +def _example_batch(datastore, predictor, batch_size=2, pred_steps=3): + """ + Build a synthetic ``(init_states, forcing_features, target_states)`` batch. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing variable counts. + predictor : BaseGraphEFM + Predictor providing the grid node count. + batch_size : int, default 2 + Number of samples in the batch. + pred_steps : int, default 3 + Rollout length. + + Returns + ------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. + """ + num_grid_nodes = predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + init_states = torch.randn(batch_size, 2, num_grid_nodes, d_state) + forcing_features = torch.randn( + batch_size, pred_steps, num_grid_nodes, d_forcing + ) + target_states = torch.randn(batch_size, pred_steps, num_grid_nodes, d_state) + return init_states, forcing_features, target_states + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_forecaster_forward_and_ensemble_shapes(graph_name): + """The forecaster unrolls a prior-sampled rollout and stacks members into + an ensemble of the documented shape.""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster(predictor, datastore, config=config) + + B, pred_steps, num_members = 2, 3, 4 + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=B, pred_steps=pred_steps + ) + d_state = target_states.shape[-1] + num_grid_nodes = predictor.num_grid_nodes + + prediction, pred_std = forecaster( + init_states, forcing_features, target_states + ) + assert prediction.shape == (B, pred_steps, num_grid_nodes, d_state) + assert pred_std is None # output_std=False predictor + + ensemble, per_member_std = forecaster.sample_ensemble( + init_states, forcing_features, target_states, num_members=num_members + ) + assert ensemble.shape == ( + B, + num_members, + pred_steps, + num_grid_nodes, + d_state, + ) + assert per_member_std is None + # Members carry independent latent samples + assert not torch.allclose(ensemble[:, 0], ensemble[:, 1]) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_predictor_step_distributions_contract(graph_name): + """step_distributions reuses the shared graph embedding and returns the + prior on the inference path, and the posterior (with the prior gated on + compute_prior) on the training path.""" + datastore, _ = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=2, pred_steps=1 + ) + prev_prev_state, prev_state = init_states[:, 0], init_states[:, 1] + forcing, target_state = forcing_features[:, 0], target_states[:, 0] + d_state = target_state.shape[-1] + num_grid_nodes = predictor.num_grid_nodes + + graph_emb = predictor.embedd_graph(2) + assert {"g2m", "m2g", "mesh"} <= set(graph_emb) + + # Inference path: latent from the prior, no posterior + prior, posterior, pred_mean, pred_std = predictor.step_distributions( + prev_state, prev_prev_state, forcing, graph_emb, target_state=None + ) + assert prior is not None and posterior is None + assert pred_mean.shape == (2, num_grid_nodes, d_state) + assert pred_std is None + + # Training path with KL: both distributions present + prior, posterior, _, _ = predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=True, + ) + assert prior is not None and posterior is not None + + # Training path without KL: prior skipped + prior, posterior, _, _ = predictor.step_distributions( + prev_state, + prev_prev_state, + forcing, + graph_emb, + target_state=target_state, + compute_prior=False, + ) + assert prior is None and posterior is not None + + +def test_predictor_clamps_predicted_mean(): + """With output clamping configured for a feature, Graph-EFM keeps the + predicted mean for that feature within the configured bounds, clamping it + like the deterministic models do.""" + datastore, _ = _datastore_and_config("1level") + state_names = datastore.get_vars_names(category="state") + lower, upper = -0.5, 0.5 + predictor = GraphEFMMultiScale( + datastore=datastore, + graph_name="1level", + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + prior_m2m_layers=1, + encoder_m2m_layers=1, + decoder_m2m_layers=1, + output_clamping_lower={state_names[0]: lower}, + output_clamping_upper={state_names[0]: upper}, + ) + # The first state feature has a two-sided (sigmoid) clamp registered + assert predictor.clamp_lower_upper_idx.tolist() == [0] + lower_n = (lower - predictor.state_mean[0]) / predictor.state_std[0] + upper_n = (upper - predictor.state_mean[0]) / predictor.state_std[0] + + B = 2 + num_grid_nodes = predictor.num_grid_nodes + d_state = len(state_names) + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + prev_state = torch.randn(B, num_grid_nodes, d_state) + # The current value of the clamped feature must be within its bounds so + # the inverse clamp is finite; the midpoint is a safe choice. + prev_state[..., 0] = (lower_n + upper_n) / 2 + prev_prev_state = torch.randn(B, num_grid_nodes, d_state) + forcing = torch.randn(B, num_grid_nodes, d_forcing) + + pred_mean, _ = predictor(prev_state, prev_prev_state, forcing) + + clamped_feature = pred_mean[..., 0] + assert torch.all(clamped_feature > lower_n) + assert torch.all(clamped_feature < upper_n) + + +def test_hierarchical_zero_intra_level_layers_runs(): + """A hierarchical GraphEFM with no intra-level layers uses an empty m2m + placeholder; forward must still run (regression for m2m handling).""" + datastore, _ = _datastore_and_config("hierarchical") + predictor = GraphEFM( + datastore=datastore, + graph_name="hierarchical", + hidden_dim=4, + hidden_layers=1, + latent_dim=4, + prior_intra_level_layers=0, + encoder_intra_level_layers=0, + decoder_intra_level_layers=0, + ) + assert not predictor.embedd_m2m + + B = 2 + num_grid_nodes = predictor.num_grid_nodes + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + prev_state = torch.randn(B, num_grid_nodes, d_state) + prev_prev_state = torch.randn(B, num_grid_nodes, d_state) + forcing = torch.randn(B, num_grid_nodes, d_forcing) + + pred_mean, _ = predictor(prev_state, prev_prev_state, forcing) + assert pred_mean.shape == (B, num_grid_nodes, d_state) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_elbo_training_loss_gradient_flow(graph_name): + """compute_training_loss returns a finite scalar ELBO with likelihood/KL + components, and gradients flow back into the predictor.""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=1.0 + ) + + init_states, forcing_features, target_states = _example_batch( + datastore, predictor + ) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + + torch.manual_seed(0) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + assert batch_loss.shape == () + assert torch.isfinite(batch_loss) + assert set(loss_components) == {"elbo_likelihood", "elbo_kl", "elbo"} + assert (loss_components["elbo_kl"] >= 0).all() + + batch_loss.backward() + grads = [p.grad for p in predictor.parameters() if p.grad is not None] + assert grads, "no gradients reached the predictor" + assert any(torch.any(g != 0) for g in grads) + + +@pytest.mark.parametrize("graph_name", ["1level", "hierarchical"]) +def test_elbo_kl_beta_zero_skips_kl(graph_name): + """With kl_beta=0 the loss is the negative likelihood alone and no KL + component is reported (pure auto-encoder training).""" + datastore, config = _datastore_and_config(graph_name) + predictor = _build_predictor(datastore, graph_name) + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=0.0 + ) + + init_states, forcing_features, target_states = _example_batch( + datastore, predictor + ) + interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool) + + torch.manual_seed(0) + batch_loss, loss_components = forecaster.compute_training_loss( + init_states, + forcing_features, + target_states, + interior_mask_bool=interior_mask_bool, + ) + + assert set(loss_components) == {"elbo_likelihood"} + torch.testing.assert_close(batch_loss, -loss_components["elbo_likelihood"]) + batch_loss.backward() # still differentiable + + +def test_module_training_and_validation_steps(): + """The ProbabilisticForecasterModule delegates training to the forecaster + ELBO and scores an ensemble mean during validation.""" + datastore, config = _datastore_and_config("1level") + predictor = _build_predictor(datastore, "1level") + forecaster = GraphEFMForecaster( + predictor, datastore, config=config, loss="mse", kl_beta=1.0 + ) + model = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + eval_ensemble_size=2, + ) + + B, pred_steps = 2, 3 + init_states, forcing_features, target_states = _example_batch( + datastore, predictor, batch_size=B, pred_steps=pred_steps + ) + batch_times = torch.zeros(B, pred_steps) + batch = (init_states, target_states, forcing_features, batch_times) + + torch.manual_seed(0) + train_loss = model.training_step(batch) + assert train_loss.shape == () + assert torch.isfinite(train_loss) + + model.validation_step(batch, 0) + (entry_mses,) = model.val_metrics["ens_mse"] + d_state = target_states.shape[-1] + assert entry_mses.shape == (B, pred_steps, d_state) + assert torch.all(torch.isfinite(entry_mses)) + + +@pytest.mark.parametrize( + "model_name, predictor_class, graph_name", + [ + ("graph_efm", GraphEFM, "hierarchical"), + ("graph_efm_ms", GraphEFMMultiScale, "1level"), + ], +) +def test_train_model_assembly_selects_probabilistic_path( + model_name, predictor_class, graph_name +): + """train_model's build_predictor/build_forecaster_module route the + graph_efm* models through the probabilistic assembly, reading the + Graph-EFM hyperparameters from the ``probabilistic`` config section and + producing the right predictor, forecaster, module class and checkpoint + monitor.""" + datastore, _ = _datastore_and_config(graph_name) + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ), + probabilistic=nlconfig.ProbabilisticConfig( + latent_dim=4, + prior_layers=1, + encoder_layers=1, + decoder_layers=1, + kl_beta=0.5, + eval_ensemble_size=3, + ), + ) + args = Namespace( + model=model_name, + graph=graph_name, + hidden_dim=4, + hidden_layers=1, + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + output_std=False, + g2m_gnn_type="InteractionNet", + m2g_gnn_type="InteractionNet", + loss="mse", + ) + + predictor = build_predictor(args, config, datastore) + assert isinstance(predictor, predictor_class) + # Hyperparameters came from config.probabilistic, not the CLI args + assert predictor.latent_dim == 4 + + forecaster, module_class, module_kwargs, val_monitor = ( + build_forecaster_module(args, config, datastore, predictor) + ) + assert isinstance(forecaster, GraphEFMForecaster) + assert forecaster.kl_beta == 0.5 + assert module_class is ProbabilisticForecasterModule + assert module_kwargs == {"eval_ensemble_size": 3} + assert val_monitor == "val_mean_ens_rmse" diff --git a/tests/test_latent_modules.py b/tests/test_latent_modules.py index 377182d1..405688c7 100644 --- a/tests/test_latent_modules.py +++ b/tests/test_latent_modules.py @@ -205,21 +205,15 @@ def test_graph_decoder_shapes_with_output_std( 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 - ) + mean_delta, pred_std = dec(grid_rep, latent_samples, flat_graph_emb) expected_shape = (B, flat_dims["num_grid"], flat_dims["num_state_vars"]) - assert pred_mean.shape == expected_shape + assert mean_delta.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() + (mean_delta.sum() + pred_std.sum()).backward() _assert_every_param_has_grad(dec) @@ -242,14 +236,8 @@ def test_graph_decoder_no_output_std_returns_none( 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 == ( + mean_delta, pred_std = dec(grid_rep, latent_samples, flat_graph_emb) + assert mean_delta.shape == ( B, flat_dims["num_grid"], flat_dims["num_state_vars"], @@ -296,11 +284,8 @@ def test_flat_modules_zero_m2m_layers_skip_processing( 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 == ( + mean_delta, _ = dec(grid_rep, latent_samples, flat_graph_emb) + assert mean_delta.shape == ( B, flat_dims["num_grid"], flat_dims["num_state_vars"], @@ -400,13 +385,9 @@ def test_hi_graph_decoder_shape_back_to_grid(hi_dims, hi_edges, hi_graph_emb): 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 - ) + mean_delta, pred_std = dec(grid_rep, latent_samples, hi_graph_emb) expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) - assert pred_mean.shape == expected_shape + assert mean_delta.shape == expected_shape assert pred_std.shape == expected_shape assert (pred_std > 0).all() @@ -465,13 +446,11 @@ def test_hi_graph_decoder_three_levels(): 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) + mean_delta, pred_std = dec(grid_rep, latent_samples, graph_emb) + assert mean_delta.shape == (B, num_grid, num_state_vars) assert pred_std.shape == (B, num_grid, num_state_vars) - (pred_mean.sum() + pred_std.sum()).backward() + (mean_delta.sum() + pred_std.sum()).backward() _assert_every_param_has_grad(dec) @@ -549,13 +528,9 @@ def test_hi_graph_decoder_zero_intra_layers(hi_dims, hi_edges, hi_graph_emb): 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 - ) + mean_delta, pred_std = dec(grid_rep, latent_samples, hi_graph_emb) expected_shape = (B, hi_dims["num_grid"], hi_dims["num_state_vars"]) - assert pred_mean.shape == expected_shape + assert mean_delta.shape == expected_shape assert pred_std.shape == expected_shape