neural_lam.models.step_predictors.graph.hierarchical

Base implementations for hierarchical (multi-level) graph models.

Classes

BaseHiGraphModel

Base class for hierarchical graph models.

Module Contents

class neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'multiscale', hidden_dim: int = 64, hidden_layers: int = 1, processor_layers: int = 4, mesh_aggr: str = 'sum', 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, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', mesh_up_gnn_type: str = 'InteractionNet', mesh_down_gnn_type: str = 'InteractionNet')

Bases: neural_lam.models.step_predictors.graph.base.BaseGraphModel

Base class for hierarchical graph models.

Extend BaseGraphModel with hierarchical mesh structures.

embedd_mesh_nodes() → torch.Tensor

Embed static mesh node features for the bottom level only; remaining levels are embedded at the start of process_step.

Returns:

Shape (num_mesh_nodes[0], hidden_dim). Embedded bottom-level mesh node representations. Dims: num_mesh_nodes[0] is the number of nodes at level 0 and hidden_dim is the hidden dimension.

Return type:

torch.Tensor

get_num_mesh() → tuple[int, int]

Compute mesh node counts used for encoding and decoding.

Returns:

Total number of mesh nodes and the number to ignore during encoding/decoding.

Return type:

tuple[int, int]

abstract hi_processor_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_up_rep: list[torch.Tensor], mesh_down_rep: list[torch.Tensor]) → tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]

Internal processor step between mesh init and read-out.

Parameters:

• mesh_rep_levels (list of torch.Tensor) – One tensor per level, each of shape (B, num_mesh_nodes[l], hidden_dim). Node representations at each hierarchy level. Dims: B is batch size, num_mesh_nodes[l] is the node count at level l, and hidden_dim is the hidden dimension.

• mesh_same_rep (list of torch.Tensor) – One tensor per level, each of shape (B, num_edges[l], hidden_dim). Same-level edge representations. num_edges[l] is the edge count at level l.

• mesh_up_rep (list of torch.Tensor) – One tensor per inter-level gap, each of shape (B, num_edges[l], hidden_dim). Upward edge representations from level l to l+1.

• mesh_down_rep (list of torch.Tensor) – One tensor per inter-level gap, each of shape (B, num_edges[l], hidden_dim). Downward edge representations from level l+1 to l.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep) in the same order as the inputs.

Return type:

tuple of (list, list, list, list)

process_step(mesh_rep: torch.Tensor) → torch.Tensor

Process the mesh representation across all hierarchy levels, implementing the full init-process-readout cycle.

Parameters:

mesh_rep (torch.Tensor) – Shape (B, num_mesh_nodes[0], hidden_dim). Bottom-level mesh node representations from the encoder. Dims: B is batch size, num_mesh_nodes[0] is the number of nodes at level 0, and hidden_dim is the hidden dimension.

Returns:

Shape (B, num_mesh_nodes[0], hidden_dim). Updated bottom-level mesh node representations. Dims: same as mesh_rep.

Return type:

torch.Tensor

level_mesh_sizes

m2m_edge_index: list[torch.Tensor]

m2m_features: list[torch.Tensor]

mesh_down_edge_index: list[torch.Tensor]

mesh_down_embedders

mesh_down_features: list[torch.Tensor]

mesh_down_gnn_type = 'InteractionNet'

mesh_embedders

mesh_init_gnns

mesh_read_gnns

mesh_same_embedders

mesh_static_features: list[torch.Tensor]

mesh_up_edge_index: list[torch.Tensor]

mesh_up_embedders

mesh_up_features: list[torch.Tensor]

mesh_up_gnn_type = 'InteractionNet'

num_levels