neural_lam.models.step_predictors.graph.graph_efm

neural_lam.models.step_predictors.graph.graph_efm#

Graph-based Ensemble Forecasting Model (Graph-EFM) single-step predictors, for hierarchical (GraphEFM) and flat (GraphEFMMultiScale) mesh graphs.

Classes#

`BaseGraphEFM`	Base class for Graph-based Ensemble Forecasting Model single-step
`GraphEFM`	Graph-based Ensemble Forecasting Model on a hierarchical mesh graph.
`GraphEFMMultiScale`	Graph-based Ensemble Forecasting Model on a flat mesh graph

Module Contents#

class neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM(datastore: neural_lam.datastore.BaseDatastore, graph_name: str, hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: int | None = None, learn_prior: bool = True, prior_dist: str = 'isotropic', prior_layers: int = 2, g2m_gnn_type: str = 'InteractionNet', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, output_clamping_lower: Dict[str, float] | None = None, output_clamping_upper: Dict[str, float] | None = None)[source]#

Bases: neural_lam.models.step_predictors.base.StepPredictor

Base class for Graph-based Ensemble Forecasting Model single-step predictors.

A latent-variable step predictor consisting of a conditional prior, a variational encoder and a latent decoder, each of which carries its own grid-to-mesh, on-mesh and mesh-to-grid GNNs. The encode-process-decode backbone of BaseGraphModel therefore does not apply – this extends StepPredictor directly. forward samples a single step from the prior. The encoder (variational posterior) and prior are exposed as self.encoder/self.prior_model for a forecaster to condition on the target and assemble a training objective from; the predictor itself does not compute any loss.

This base class sets up everything that is independent of the mesh graph type: it loads the graph, calls the subclass’s check_graph_type() to verify the loaded graph matches what it requires, then builds the prior, delegating to the subclass’s build_learnable_prior() when learned and to latent_spatial_dim for the constant prior’s node count. Concrete subclasses build the mesh embedders and the encoder/decoder latent modules, and implement embedd_mesh(), check_graph_type(), build_learnable_prior() and latent_spatial_dim. See GraphEFM (hierarchical graph) and GraphEFMMultiScale (flat graph).

Set up the graph-type independent parts of the predictor.

Loads the graph, builds the grid embedders, the grid-mesh edge embedders and the prior. Building the mesh embedders and the encoder/decoder latent modules is left to the subclass constructor.

Parameters:

datastore (BaseDatastore) – Datastore providing static features, standardization statistics and variable counts.
graph_name (str) – Name of the graph directory (under <root>/graph) to load. Must be of the graph type required by the concrete subclass.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
latent_dim (int, optional) – Dimensionality of the latent variable at each latent-carrying mesh node (top level for hierarchical graphs, all mesh nodes for flat graphs); defaults to hidden_dim when None. The resolved value is stored as self.latent_dim for the subclass to reuse when building its encoder/decoder.
learn_prior (bool) – If True, the prior is the graph-type specific learnable encoder built by build_learnable_prior(), conditioned on the previous state; if False, a constant Normal(0, 1) prior is used.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal".
prior_layers (int) – Number of on-mesh GNN layers in the learnable prior.
g2m_gnn_type (str) – GNN type for the grid-to-mesh step of the prior (key in gnn_layers.GNN_TYPES).
num_past_forcing_steps (int) – Number of past forcing steps included in the input window.
num_future_forcing_steps (int) – Number of future forcing steps included in the input window.
output_std (bool) – If True, the decoder outputs a per-variable std alongside the mean; if False, forward returns None for the std.
output_clamping_lower (dict of str to float, optional) – Lower clamping limits per output variable.
output_clamping_upper (dict of str to float, optional) – Upper clamping limits per output variable.

abstract build_learnable_prior(latent_dim: int, hidden_dim: int, hidden_layers: int, g2m_gnn_type: str, prior_dist: str, prior_layers: int) → torch.nn.Module[source]#

Build the graph-type specific learnable prior encoder.

Implemented by the concrete subclass, which knows the mesh graph type and therefore the appropriate latent encoder class.

Parameters:

latent_dim (int) – Dimensionality of the latent variable at each mesh node.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
g2m_gnn_type (str) – GNN type for the grid-to-mesh step of the prior.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal".
prior_layers (int) – Number of on-mesh GNN layers in the prior.

Returns:

The learnable prior latent encoder.

Return type:

torch.nn.Module

abstract check_graph_type(graph_name: str) → None[source]#

Verify the loaded graph (self.hierarchical) is of the type this predictor requires.

Implemented by the concrete subclass, which is the only place that knows whether it requires a hierarchical or flat mesh graph. Called by this base class right after loading the graph, before anything that assumes a specific graph shape.

Parameters:: graph_name (str) – Name of the graph directory that was loaded, for the error message.
Raises:: ValueError – If self.hierarchical does not match what this predictor requires.

embedd_grid_and_graph(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) → tuple[torch.Tensor, dict[str, Any]][source]#

Embed the grid (states up to t-1) and the full graph.

Parameters:

prev_state (torch.Tensor) – Shape (B, num_grid_nodes, d_state). X_t.
prev_prev_state (torch.Tensor) – Shape (B, num_grid_nodes, d_state). X_{t-1}.
forcing (torch.Tensor) – Shape (B, num_grid_nodes, d_forcing).

Returns:

grid_emb (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Grid embedding.
graph_emb (dict) – Edge/mesh embeddings, each entry of shape (B, *, d_h).

embedd_grid_with_target(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor, current_state: torch.Tensor) → torch.Tensor[source]#

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:

Shape (B, num_grid_nodes, d_h). Grid embedding.

Return type:

torch.Tensor

abstract embedd_mesh(batch_size: int) → dict[str, Any][source]#

Embed static mesh node and intra-mesh edge features.

Parameters:: batch_size (int) – Batch size to expand the embeddings to.
Returns:: 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.
Return type:: dict

forward(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) → tuple[torch.Tensor, torch.Tensor | None][source]#

Sample one time step prediction: embed features, sample the latent from the prior, decode, and return the predicted next state. The prediction is stochastic only through the latent sample; no observation noise is added.

Parameters:

prev_state (torch.Tensor) – Shape (B, num_grid_nodes, d_state). X_t.
prev_prev_state (torch.Tensor) – Shape (B, num_grid_nodes, d_state). X_{t-1}.
forcing (torch.Tensor) – Shape (B, num_grid_nodes, d_forcing).

Returns:

new_state (torch.Tensor) – Shape (B, num_grid_nodes, d_state). Predicted X_{t+1} (the decoder mean, given the sampled latent).
pred_std (torch.Tensor or None) – Shape (B, num_grid_nodes, d_state) when output_std is True, otherwise None.

g2m_edge_index: torch.Tensor#

g2m_embedder#

grid_current_embedder#

grid_dim#

grid_prev_embedder#

hierarchical#

latent_dim#

property latent_spatial_dim: int#

Abstractmethod:

Number of mesh nodes the latent variable lives on.

Implemented by the concrete subclass, which knows the mesh graph type: the top mesh level for hierarchical graphs, or every mesh node for flat graphs.

Returns:: Number of latent-carrying mesh nodes.
Return type:: int

m2g_edge_index: torch.Tensor#

m2g_embedder#

mlp_blueprint_end#

num_state_vars#

class neural_lam.models.step_predictors.graph.graph_efm.GraphEFM(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'hierarchical', hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: int | None = None, prior_intra_level_layers: int = 2, encoder_intra_level_layers: int = 2, decoder_intra_level_layers: int = 4, learn_prior: bool = True, prior_dist: str = 'isotropic', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', output_std: bool = False, output_clamping_lower: Dict[str, float] | None = None, output_clamping_upper: Dict[str, float] | None = None)[source]#

Bases: 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.

Build the mesh embedders and the hierarchical encoder/decoder latent modules. The prior is built by the base class.

Parameters:

datastore (BaseDatastore) – Datastore providing static features, standardization statistics and variable counts.
graph_name (str) – Name of the graph directory (under <root>/graph) to load. Must be a hierarchical graph.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
latent_dim (int, optional) – Dimensionality of the latent variable at each top-level mesh node; defaults to hidden_dim when None. Forwarded to the base class, which resolves the default and stores it as self.latent_dim.
prior_intra_level_layers (int) – Number of intra-level GNN layers in the (learned) prior. Forwarded to the base class as prior_layers.
encoder_intra_level_layers (int) – Number of intra-level GNN layers in the variational encoder.
decoder_intra_level_layers (int) – Number of intra-level GNN layers in the latent decoder.
learn_prior (bool) – If True, the prior is a hierarchical graph encoder conditioned on the previous state; if False, a constant Normal(0, 1) prior is used. Forwarded to the base class.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal". Forwarded to the base class.
num_past_forcing_steps (int) – Number of past forcing steps included in the input window.
num_future_forcing_steps (int) – Number of future forcing steps included in the input window.
g2m_gnn_type (str) – GNN type for the grid-to-mesh steps of the prior, encoder and decoder (key in gnn_layers.GNN_TYPES).
m2g_gnn_type (str) – GNN type for the mesh-to-grid step of the decoder (key in gnn_layers.GNN_TYPES).
output_std (bool) – If True, the decoder outputs a per-variable std alongside the mean; if False, forward returns None for the std.
output_clamping_lower (dict of str to float, optional) – Lower clamping limits per output variable.
output_clamping_upper (dict of str to float, optional) – Upper clamping limits per output variable.

build_learnable_prior(latent_dim: int, hidden_dim: int, hidden_layers: int, g2m_gnn_type: str, prior_dist: str, prior_layers: int) → neural_lam.models.latent.HiGraphLatentEncoder[source]#

Build the hierarchical learnable prior encoder.

Parameters:

latent_dim (int) – Dimensionality of the latent variable at each top-level mesh node.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
g2m_gnn_type (str) – GNN type for the grid-to-mesh step of the prior.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal".
prior_layers (int) – Number of intra-level GNN layers in the prior.

Returns:

The learnable prior latent encoder.

Return type:

HiGraphLatentEncoder

check_graph_type(graph_name: str) → None[source]#

Verify the loaded graph is hierarchical.

Parameters:: graph_name (str) – Name of the graph directory that was loaded, for the error message.
Raises:: ValueError – If the loaded graph is flat.

embedd_mesh(batch_size: int) → dict[str, Any][source]#

Embed static mesh node and intra-mesh edge features per level.

Parameters:: batch_size (int) – Batch size to expand the embeddings to.
Returns:: Entries mesh, m2m, mesh_up and mesh_down, each a list with one (B, *, d_h) tensor per mesh level (or inter-level connection).
Return type:: dict

decoder#

embedd_m2m#

encoder#

property latent_spatial_dim: int#

Number of mesh nodes on the top mesh level, where the latent variable lives.

Returns:: Number of top-level mesh nodes.
Return type:: int

m2m_edge_index: neural_lam.utils.BufferList#

m2m_features: list[torch.Tensor]#

mesh_down_edge_index: neural_lam.utils.BufferList#

mesh_down_embedders#

mesh_down_features: list[torch.Tensor]#

mesh_embedders#

mesh_static_features: list[torch.Tensor]#

mesh_up_edge_index: neural_lam.utils.BufferList#

mesh_up_embedders#

mesh_up_features: list[torch.Tensor]#

class neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'multiscale', hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: int | None = None, prior_m2m_layers: int = 2, encoder_m2m_layers: int = 2, decoder_m2m_layers: int = 4, learn_prior: bool = True, prior_dist: str = 'isotropic', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', output_std: bool = False, output_clamping_lower: Dict[str, float] | None = None, output_clamping_upper: Dict[str, float] | None = None)[source]#

Bases: BaseGraphEFM

Graph-based Ensemble Forecasting Model on a flat mesh graph (e.g. a multi-scale graph).

The latent variable lives on the (single) level of the mesh graph. The prior and variational encoder are GraphLatentEncoder``s and the decoder is a ``GraphLatentDecoder.

Build the mesh embedders and the flat-graph encoder/decoder latent modules. The prior is built by the base class.

Parameters:

datastore (BaseDatastore) – Datastore providing static features, standardization statistics and variable counts.
graph_name (str) – Name of the graph directory (under <root>/graph) to load. Must be a flat graph.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
latent_dim (int, optional) – Dimensionality of the latent variable at each mesh node; defaults to hidden_dim when None. Forwarded to the base class, which resolves the default and stores it as self.latent_dim.
prior_m2m_layers (int) – Number of on-mesh (m2m) GNN layers in the (learned) prior. Forwarded to the base class as prior_layers.
encoder_m2m_layers (int) – Number of on-mesh (m2m) GNN layers in the variational encoder.
decoder_m2m_layers (int) – Number of on-mesh (m2m) GNN layers in the latent decoder.
learn_prior (bool) – If True, the prior is a graph encoder conditioned on the previous state; if False, a constant Normal(0, 1) prior is used. Forwarded to the base class.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal". Forwarded to the base class.
num_past_forcing_steps (int) – Number of past forcing steps included in the input window.
num_future_forcing_steps (int) – Number of future forcing steps included in the input window.
g2m_gnn_type (str) – GNN type for the grid-to-mesh steps of the prior, encoder and decoder (key in gnn_layers.GNN_TYPES).
m2g_gnn_type (str) – GNN type for the mesh-to-grid step of the decoder (key in gnn_layers.GNN_TYPES).
output_std (bool) – If True, the decoder outputs a per-variable std alongside the mean; if False, forward returns None for the std.
output_clamping_lower (dict of str to float, optional) – Lower clamping limits per output variable.
output_clamping_upper (dict of str to float, optional) – Upper clamping limits per output variable.

Build the flat-graph learnable prior encoder.

Parameters:

latent_dim (int) – Dimensionality of the latent variable at each mesh node.
hidden_dim (int) – Dimensionality of internal node and edge representations.
hidden_layers (int) – Number of hidden layers in internal MLPs.
g2m_gnn_type (str) – GNN type for the grid-to-mesh step of the prior.
prior_dist (str) – Output distribution of the prior: "isotropic" or "diagonal".
prior_layers (int) – Number of on-mesh (m2m) GNN layers in the prior.

Returns:

The learnable prior latent encoder.

Return type:

GraphLatentEncoder

check_graph_type(graph_name: str) → None[source]#

Verify the loaded graph is flat.

Parameters:: graph_name (str) – Name of the graph directory that was loaded, for the error message.
Raises:: ValueError – If the loaded graph is hierarchical.

embedd_mesh(batch_size: int) → dict[str, Any][source]#

Embed static mesh node and intra-mesh edge features.

Parameters:: batch_size (int) – Batch size to expand the embeddings to.
Returns:: Entries mesh: (B, num_mesh_nodes, d_h) and m2m: (B, M_m2m, d_h).
Return type:: dict

decoder#

encoder#

property latent_spatial_dim: int#

Number of mesh nodes, where the latent variable lives.

Returns:: Number of mesh nodes.
Return type:: int

m2m_edge_index: torch.Tensor#

m2m_embedder#

m2m_features: torch.Tensor#

mesh_embedder#

mesh_static_features: torch.Tensor#