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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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 2cf6150670d3a5c6d738fc67f8ee6fe560afbb60 Mon Sep 17 00:00:00 2001 From: Jeevant Prakhar Singh Date: Sun, 21 Jun 2026 15:56:15 +0530 Subject: [PATCH 24/24] feat: add probabilistic forecaster interface decoupled from Graph-EFM Introduce the latent-variable forecasting interface as three layers with a clean boundary: the predictor produces distributions only, the forecaster owns the rollout and the KL between those distributions, and the module assembles the ELBO. - Add LatentStepPredictor (step_predictors/probabilistic.py): abstract step_distributions -> (prior, posterior, pred_mean, pred_std); forward is the prior-sampling special case. Likelihood/KL/ELBO live outside it. - Add ProbabilisticARForecaster: training_rollout (posterior-conditioned, reduces per-step KL) and sample_trajectories (prior ensemble). - Add ProbabilisticForecasterModule: assembles the beta-weighted negative ELBO in training_step, substituting the constant per_var_std fallback. - Add crps_ens metric for ensemble CRPS. - Add interface tests driven by a graph-free dummy LatentStepPredictor. Graph-EFM is not yet wired onto this interface; that lands in a follow-up. --- neural_lam/metrics.py | 74 +++++ neural_lam/models/__init__.py | 3 + .../models/forecasters/probabilistic.py | 156 +++++++++++ neural_lam/models/probabilistic_module.py | 100 +++++++ .../models/step_predictors/probabilistic.py | 114 ++++++++ tests/test_probabilistic_interface.py | 256 ++++++++++++++++++ 6 files changed, 703 insertions(+) create mode 100644 neural_lam/models/forecasters/probabilistic.py create mode 100644 neural_lam/models/probabilistic_module.py create mode 100644 neural_lam/models/step_predictors/probabilistic.py create mode 100644 tests/test_probabilistic_interface.py diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 1eb1d526..c95b8bbf 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -387,6 +387,80 @@ def crps_gauss( ) +def crps_ens( + pred: torch.Tensor, + target: torch.Tensor, + pred_std: Optional[torch.Tensor] = None, + mask: Optional[torch.Tensor] = None, + average_grid: bool = True, + sum_vars: bool = True, + ens_dim: int = 1, +) -> torch.Tensor: + """ + (Negative) Continuous Ranked Probability Score from ensemble samples. + + Unbiased ("fair") estimator of the CRPS computed directly from a set of + ensemble members, following Zamo and Naveau / WeatherBench 2. + + Parameters + ---------- + pred : torch.Tensor + Shape ``(..., M, ..., N, num_variables)``. Ensemble predictions, with + the ``M`` ensemble members laid out along ``ens_dim``. + target : torch.Tensor + Shape ``(..., N, num_variables)``. Ground-truth target (no ensemble + dimension). + pred_std : torch.Tensor or None, optional + Unused; accepted so the signature matches the other metrics. + mask : torch.Tensor or None, optional + Shape ``(N,)``. Boolean mask selecting grid nodes. ``None`` uses all. + average_grid : bool, optional + If True, average over the grid dimension ``N``. + sum_vars : bool, optional + If True, sum over the variable dimension ``num_variables``. + ens_dim : int, optional + Batch dimension along which the ensemble members are laid out and + reduced. Default ``1``. + + Returns + ------- + torch.Tensor + Reduced CRPS values; shape depends on ``average_grid`` and + ``sum_vars``. + """ + del pred_std # unused; present for a uniform metric signature + num_ens = pred.shape[ens_dim] + if num_ens == 1: + # With a single sample the CRPS reduces to the MAE + return mae( + pred.squeeze(ens_dim), + target, + None, + mask=mask, + average_grid=average_grid, + sum_vars=sum_vars, + ) + + mean_mae = torch.mean( + torch.abs(pred - target.unsqueeze(ens_dim)), dim=ens_dim + ) # (..., N, num_variables) + if num_ens == 2: + # Cheap exact estimator for a two-member ensemble + pair_diffs_term = -0.5 * torch.abs( + pred.select(ens_dim, 0) - pred.select(ens_dim, 1) + ) + else: + # Rank-based estimator, O(M log M) compute and O(M) memory. + # Ranks start at 1; two argsorts give the per-entry ranks. + ranks = pred.argsort(dim=ens_dim).argsort(ens_dim) + 1 + pair_diffs_term = (1 / (num_ens - 1)) * torch.mean( + (num_ens + 1 - 2 * ranks) * pred, dim=ens_dim + ) + crps_estimator = mean_mae + pair_diffs_term # (..., N, num_variables) + + return mask_and_reduce_metric(crps_estimator, mask, average_grid, sum_vars) + + DEFINED_METRICS = { "mse": mse, "mae": mae, diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 68e2a4e1..2c8920bb 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -3,7 +3,9 @@ # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster +from .forecasters.probabilistic import ProbabilisticARForecaster from .module import ForecasterModule +from .probabilistic_module import ProbabilisticForecasterModule from .step_predictors.base import StepPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_efm import GraphEFM, GraphEFMMS @@ -11,6 +13,7 @@ from .step_predictors.graph.hi_lam import HiLAM from .step_predictors.graph.hi_lam_parallel import HiLAMParallel from .step_predictors.graph.hierarchical import BaseHiGraphModel +from .step_predictors.probabilistic import LatentStepPredictor # NOTE: GraphEFM/GraphEFMMS are intentionally NOT registered in MODELS yet. # The shared construction call in train_model.py instantiates the chosen diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py new file mode 100644 index 00000000..120f67a9 --- /dev/null +++ b/neural_lam/models/forecasters/probabilistic.py @@ -0,0 +1,156 @@ +"""Autoregressive forecaster for Graph-EFM latent-variable predictors.""" + +# Standard library +from typing import Optional + +# Third-party +import torch + +# Local +from ..step_predictors.probabilistic import LatentStepPredictor +from .autoregressive import ARForecaster + + +class ProbabilisticARForecaster(ARForecaster): + """ + Autoregressive forecaster for latent-variable step predictors. + + The inherited :meth:`ARForecaster.forward` is the evaluation rollout: each + step calls the predictor (which samples its prior), conditions + autoregressively on the predicted state, and overwrites boundary nodes + with the true value. This subclass adds :meth:`training_rollout`, the + posterior-conditioned rollout used for training: at each step it asks the + predictor for the prior, the variational posterior (conditioned on the + true target) and the decoded reconstruction, autoregresses on the + predicted mean, and reduces the per-step KL between posterior and prior. + The KL is purely between the model's own distributions, so it is reduced + here; the likelihood and the ELBO assembly are left to the module. + """ + + # The wrapped predictor must expose per-step latent distributions + predictor: LatentStepPredictor + + def training_rollout( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + target_states: torch.Tensor, + compute_kl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Posterior-conditioned rollout producing the per-step ELBO components. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. + target_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. The true next + states, used to condition the posterior and to overwrite the + boundary nodes. + compute_kl : bool, optional + If ``False`` (KL weight zero), skip the prior and return + ``kl_terms`` as ``None``. + + Returns + ------- + pred_means : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. Raw decoder + means (not boundary-overwritten), for the likelihood. + pred_stds : torch.Tensor or None + Shape ``(B, pred_steps, num_grid_nodes, d_state)`` when the + predictor emits its own std, otherwise ``None``. + kl_terms : torch.Tensor or None + Shape ``(B, pred_steps)``. Per-step KL of posterior from prior, + or ``None`` when ``compute_kl`` is ``False``. + """ + prev_prev_state = init_states[:, 0] + prev_state = init_states[:, 1] + mean_list = [] + std_list = [] + kl_list = [] + + for i in range(forcing_features.shape[1]): + target = target_states[:, i] + prior_dist, posterior_dist, pred_mean, pred_std = ( + self.predictor.step_distributions( + prev_state, + prev_prev_state, + forcing_features[:, i], + target=target, + compute_prior=compute_kl, + ) + ) + + mean_list.append(pred_mean) # raw, masked to interior later + if pred_std is not None: + std_list.append(pred_std) + if posterior_dist is not None and prior_dist is not None: + kl = torch.distributions.kl_divergence( + posterior_dist, prior_dist + ) + # Sum over all latent dimensions, leaving the batch dim + kl_list.append(kl.sum(dim=tuple(range(1, kl.ndim)))) + + # Autoregress on the predicted mean; boundary overwritten by truth + new_state = ( + self.boundary_mask * target + self.interior_mask * pred_mean + ) + prev_prev_state = prev_state + prev_state = new_state + + pred_means = torch.stack(mean_list, dim=1) + pred_stds = torch.stack(std_list, dim=1) if std_list else None + kl_terms = torch.stack(kl_list, dim=1) if kl_list else None + return pred_means, pred_stds, kl_terms + + def sample_trajectories( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + boundary_states: torch.Tensor, + num_traj: int, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Sample an ensemble of prior trajectories. + + Runs the evaluation rollout (:meth:`ARForecaster.forward`, which + samples the prior each step) ``num_traj`` times and stacks the + members along a new ensemble dimension. Used for the CRPS term and + for ensemble evaluation. + + Parameters + ---------- + init_states : torch.Tensor + Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. + boundary_states : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, d_state)``. True states + used to overwrite boundary nodes at each step. + num_traj : int + Number of trajectories ``S`` to sample. + + Returns + ------- + traj_means : torch.Tensor + Shape ``(B, S, pred_steps, num_grid_nodes, d_state)``. + traj_stds : torch.Tensor or None + Shape ``(B, S, pred_steps, num_grid_nodes, d_state)`` when the + predictor emits its own std, otherwise ``None``. + """ + mean_list = [] + std_list = [] + for _ in range(num_traj): + prediction, pred_std = self( + init_states, forcing_features, boundary_states + ) + mean_list.append(prediction) + if pred_std is not None: + std_list.append(pred_std) + + traj_means = torch.stack(mean_list, dim=1) # (B, S, T, N, d_state) + traj_stds = torch.stack(std_list, dim=1) if std_list else None + return traj_means, traj_stds diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py new file mode 100644 index 00000000..dcbedb01 --- /dev/null +++ b/neural_lam/models/probabilistic_module.py @@ -0,0 +1,100 @@ +"""Lightning module that trains latent-variable forecasters with the ELBO.""" + +# Local +from .forecasters.probabilistic import ProbabilisticARForecaster +from .module import ForecasterModule + + +class ProbabilisticForecasterModule(ForecasterModule): + """ + Lightning module for latent-variable forecasters. + + Extends :class:`ForecasterModule` by assembling the training loss as the + (beta-weighted) negative ELBO. The forecaster's ``training_rollout`` + returns the per-step decoder reconstruction and the per-step KL; this + module computes the likelihood from the reconstruction and the target + through ``self.loss`` (substituting the constant ``per_var_std`` when the + predictor emits no std), sums the likelihood and KL over the rollout, + weights the KL by ``kl_beta`` and averages over the batch. Evaluation + reuses the inherited single-member rollout. + """ + + # The wrapped forecaster must roll out a latent-variable predictor + forecaster: ProbabilisticARForecaster + + def __init__(self, *args, kl_beta: float = 1.0, **kwargs): + """ + Initialize the module and store the KL weight. + + Parameters + ---------- + *args + Positional arguments forwarded to + :meth:`ForecasterModule.__init__` (``forecaster``, ``config``, + ``datastore``, ...). + kl_beta : float, optional + Weight of the KL term in the ELBO. ``0`` trains a pure + autoencoder (the prior and KL are skipped). Default ``1.0``. + **kwargs + Keyword arguments forwarded to + :meth:`ForecasterModule.__init__` (``loss``, ``lr``, ...). + """ + super().__init__(*args, **kwargs) + self.kl_beta = kl_beta + + def training_step(self, batch): + """ + Assemble the beta-weighted negative ELBO over the rollout. + + Parameters + ---------- + batch : tuple + ``(init_states, target_states, forcing_features, batch_times)``. + + Returns + ------- + torch.Tensor + Scalar training loss. + """ + init_states, target_states, forcing_features, _ = batch + compute_kl = self.kl_beta > 0 + + pred_means, pred_stds, kl_terms = self.forecaster.training_rollout( + init_states, + forcing_features, + target_states, + compute_kl=compute_kl, + ) + pred_std = pred_stds if pred_stds is not None else self.per_var_std + + # Per-entry log-likelihood is the negative loss (exactly the + # log-likelihood for the nll loss); sum over interior grid + vars. + entry_log_lik = -self.loss( + pred_means, + target_states, + pred_std, + mask=self.interior_mask_bool, + average_grid=False, + sum_vars=False, + ) # (B, pred_steps, num_interior_nodes, d_state) + likelihood_terms = entry_log_lik.sum(dim=(2, 3)) # (B, pred_steps) + + # Sum the per-step terms over the rollout, mean over the batch + per_sample_likelihood = likelihood_terms.sum(dim=1) # (B,) + loss = -per_sample_likelihood.mean() + log_dict = {"elbo_likelihood": per_sample_likelihood.mean()} + if kl_terms is not None: + per_sample_kl = kl_terms.sum(dim=1) # (B,) + loss = loss + self.kl_beta * per_sample_kl.mean() + log_dict["elbo_kl"] = per_sample_kl.mean() + log_dict["train_loss"] = loss + + self.log_dict( + log_dict, + prog_bar=True, + on_step=True, + on_epoch=True, + sync_dist=True, + batch_size=init_states.shape[0], + ) + return loss diff --git a/neural_lam/models/step_predictors/probabilistic.py b/neural_lam/models/step_predictors/probabilistic.py new file mode 100644 index 00000000..80d9a007 --- /dev/null +++ b/neural_lam/models/step_predictors/probabilistic.py @@ -0,0 +1,114 @@ +"""Interface for latent-variable step predictors.""" + +# Standard library +from abc import abstractmethod +from typing import Optional, Tuple + +# Third-party +import torch + +# Local +from .base import StepPredictor + + +class LatentStepPredictor(StepPredictor): + """ + Abstract step predictor for latent-variable models. + + Extends :class:`StepPredictor` for predictors that advance the state by + sampling a latent variable and decoding it. Subclasses implement + :meth:`step_distributions`, which returns the latent prior, the + variational posterior and the decoded reconstruction for one time step. + The evaluation :meth:`forward` is provided here as the prior-sampling + special case. The predictor produces only distributions and a + reconstruction; likelihood, KL and the ELBO are assembled outside it. + """ + + @abstractmethod + def step_distributions( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + target: Optional[torch.Tensor] = None, + compute_prior: bool = True, + ) -> Tuple[ + Optional[torch.distributions.Distribution], + Optional[torch.distributions.Distribution], + torch.Tensor, + Optional[torch.Tensor], + ]: + """ + Latent distributions and decoded reconstruction for one time step. + + The latent sample is drawn from the variational posterior when a + ``target`` is given (training path) and from the prior otherwise + (evaluation path). This is the only entry point that takes the + target. The predictor computes no likelihood or KL. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings + for this step. + target : torch.Tensor or None, optional + Shape ``(B, num_grid_nodes, num_state_vars)``. The true next state + ``X_{t+1}``. When given, the latent is sampled from the + variational posterior conditioned on it; when ``None``, from the + prior. + compute_prior : bool, optional + If ``False``, skip the prior on the training path (``prior_dist`` + is returned as ``None``); used when the KL weight is zero. On the + evaluation path the prior is always computed. + + Returns + ------- + prior_dist : torch.distributions.Distribution or None + Prior over the latent variable, or ``None`` when skipped. + posterior_dist : torch.distributions.Distribution or None + Variational posterior over the latent variable, or ``None`` on the + evaluation path. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Decoded mean given + the latent sample. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` + is True, otherwise ``None``. + """ + + def forward( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Advance one step by sampling the latent prior and decoding. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings + for this step. + + Returns + ------- + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Predicted next state + ``X_{t+1}``. + pred_std : torch.Tensor or None + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` + is True, otherwise ``None``. + """ + _, _, pred_mean, pred_std = self.step_distributions( + prev_state, prev_prev_state, forcing, target=None + ) + return pred_mean, pred_std diff --git a/tests/test_probabilistic_interface.py b/tests/test_probabilistic_interface.py new file mode 100644 index 00000000..6366f900 --- /dev/null +++ b/tests/test_probabilistic_interface.py @@ -0,0 +1,256 @@ +"""Interface-level checks for the latent-variable forecasting stack. + +Drives a minimal, graph-free dummy ``LatentStepPredictor`` through +``ProbabilisticARForecaster`` and ``ProbabilisticForecasterModule`` to verify +the predictor -> forecaster -> module -> ELBO contract independently of any +concrete predictor implementation. The Graph-EFM wiring is tested separately. +""" + +# Standard library +import warnings + +# Third-party +import pytest +import pytorch_lightning as pl +import torch +import wandb +from torch import nn + +# First-party +from neural_lam import config as nlconfig +from neural_lam.models import ( + LatentStepPredictor, + ProbabilisticARForecaster, + ProbabilisticForecasterModule, +) +from neural_lam.weather_dataset import WeatherDataModule +from tests.conftest import init_datastore_example + +NUM_PAST_FORCING_STEPS = 1 +NUM_FUTURE_FORCING_STEPS = 1 +NUM_MESH_NODES = 3 +LATENT_DIM = 4 + + +class _DummyLatentPredictor(LatentStepPredictor): + """Minimal graph-free latent predictor exercising the interface.""" + + def __init__(self, datastore): + """ + Build a graph-free latent predictor over ``datastore``. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing grid metadata and state-variable counts. + """ + super().__init__(datastore=datastore, output_std=False) + d_state = datastore.get_num_data_vars(category="state") + self.encode_latent = nn.Linear(d_state, LATENT_DIM) + self.decode_latent = nn.Linear(LATENT_DIM, d_state) + + def _latent_dist(self, state): + """ + Build a diagonal-Normal latent distribution from a grid state. + + Parameters + ---------- + state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. State summarised + into latent parameters by mean-pooling over the grid. + + Returns + ------- + torch.distributions.Normal + Shape ``(B, NUM_MESH_NODES, LATENT_DIM)``. Unit-variance Normal. + """ + mean = self.encode_latent(state.mean(dim=1)) # (B, LATENT_DIM) + mean = mean.unsqueeze(1).expand(-1, NUM_MESH_NODES, -1) + return torch.distributions.Normal(mean, torch.ones_like(mean)) + + def step_distributions( + self, + prev_state, + prev_prev_state, + forcing, + target=None, + compute_prior=True, + ): + """ + Latent distributions and reconstruction for one step. + + Parameters + ---------- + prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_t``. + prev_prev_state : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. ``X_{t-1}``; + unused by this dummy. + forcing : torch.Tensor + Shape ``(B, num_grid_nodes, num_forcing_vars)``; unused. + target : torch.Tensor or None, optional + Shape ``(B, num_grid_nodes, num_state_vars)``. True ``X_{t+1}``. + compute_prior : bool, optional + Whether to build the prior on the training path. + + Returns + ------- + prior_dist : torch.distributions.Normal or None + Prior over the latent, or ``None`` when skipped. + posterior_dist : torch.distributions.Normal or None + Posterior over the latent, or ``None`` on the eval path. + pred_mean : torch.Tensor + Shape ``(B, num_grid_nodes, num_state_vars)``. Decoded mean. + pred_std : None + This dummy never emits an std. + """ + if target is not None: + posterior_dist = self._latent_dist(target) + latent = posterior_dist.rsample() + prior_dist = ( + self._latent_dist(prev_state) if compute_prior else None + ) + else: + posterior_dist = None + prior_dist = self._latent_dist(prev_state) + latent = prior_dist.rsample() + + delta = self.decode_latent(latent.mean(dim=1)).unsqueeze(1) + pred_mean = prev_state + delta + return prior_dist, posterior_dist, pred_mean, None + + +def _build_module(kl_beta=1.0): + datastore = init_datastore_example("mdp") + config = nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind=datastore.SHORT_NAME, config_path=datastore.root_path + ) + ) + predictor = _DummyLatentPredictor(datastore) + forecaster = ProbabilisticARForecaster(predictor, datastore) + module = ProbabilisticForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss="nll", + kl_beta=kl_beta, + ) + return datastore, module + + +def _make_batch(datastore, pred_steps=3, batch_size=2): + num_grid = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * ( + NUM_PAST_FORCING_STEPS + NUM_FUTURE_FORCING_STEPS + 1 + ) + torch.manual_seed(0) + init_states = torch.randn(batch_size, 2, num_grid, d_state) + target_states = torch.randn(batch_size, pred_steps, num_grid, d_state) + forcing = torch.randn(batch_size, pred_steps, num_grid, d_forcing) + return init_states, target_states, forcing + + +@pytest.mark.parametrize("kl_beta", [1.0, 0.0]) +def test_training_rollout_terms(kl_beta): + """training_rollout returns per-step means/kl; kl is None when β == 0.""" + datastore, module = _build_module(kl_beta=kl_beta) + init_states, target_states, forcing = _make_batch(datastore) + + pred_means, pred_stds, kl_terms = module.forecaster.training_rollout( + init_states, + forcing, + target_states, + compute_kl=kl_beta > 0, + ) + b, t = init_states.shape[0], forcing.shape[1] + assert pred_means.shape == ( + b, + t, + datastore.num_grid_points, + datastore.get_num_data_vars(category="state"), + ) + assert pred_stds is None # output_std=False + if kl_beta > 0: + assert kl_terms.shape == (b, t) + else: + assert kl_terms is None + + +def test_sample_trajectories_shape(): + """sample_trajectories stacks an ensemble dimension of prior rollouts.""" + datastore, module = _build_module() + init_states, target_states, forcing = _make_batch(datastore) + + num_traj = 4 + traj_means, traj_stds = module.forecaster.sample_trajectories( + init_states, forcing, target_states, num_traj + ) + b, t = init_states.shape[0], forcing.shape[1] + assert traj_means.shape == ( + b, + num_traj, + t, + datastore.num_grid_points, + datastore.get_num_data_vars(category="state"), + ) + assert traj_stds is None # output_std=False + + +def test_elbo_loss_is_finite_and_differentiable(): + """The assembled ELBO loss is finite and gradients reach the predictor.""" + datastore, module = _build_module() + init_states, target_states, forcing = _make_batch(datastore) + + pred_means, _, kl_terms = module.forecaster.training_rollout( + init_states, + forcing, + target_states, + compute_kl=True, + ) + entry_log_lik = -module.loss( + pred_means, + target_states, + module.per_var_std, + mask=module.interior_mask_bool, + average_grid=False, + sum_vars=False, + ) + likelihood = entry_log_lik.sum(dim=(2, 3)).sum(dim=1) # (B,) + loss = -likelihood.mean() + module.kl_beta * kl_terms.sum(dim=1).mean() + + assert torch.isfinite(loss) + loss.backward() + assert any( + p.grad is not None and torch.any(p.grad != 0) + for p in module.forecaster.predictor.parameters() + ) + + +def test_end_to_end_probabilistic_training_runs(): + """A full Trainer epoch (train + val) runs through the whole pipeline.""" + datastore, module = _build_module() + + data_module = WeatherDataModule( + datastore=datastore, + ar_steps_train=3, + ar_steps_eval=5, + batch_size=2, + num_workers=1, + num_past_forcing_steps=NUM_PAST_FORCING_STEPS, + num_future_forcing_steps=NUM_FUTURE_FORCING_STEPS, + ) + + trainer = pl.Trainer( + max_epochs=1, + accelerator="cpu", + devices=1, + log_every_n_steps=1, + enable_checkpointing=False, + logger=False, + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + wandb.init(mode="disabled") + trainer.fit(module, data_module)