neural_lam.models.step_predictors.graph.hi_lam

neural_lam.models.step_predictors.graph.hi_lam#

Hierarchical graph-based LAM model.

Classes#

HiLAM

Hierarchical graph model with message passing that goes sequentially down

Module Contents#

class neural_lam.models.step_predictors.graph.hi_lam.HiLAM(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')[source]#

Bases: neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel

Hierarchical graph model with message passing that goes sequentially down and up the hierarchy during processing. The Hi-LAM model from Oskarsson et al. (2023)

Initialize the HiLAM model.

Parameters:

  • datastore (BaseDatastore) – The datastore providing grid metadata and data access.

  • graph_name (str, default "multiscale") – The name of the graph to load.

  • hidden_dim (int, default 64) – The dimension of the hidden representations.

  • hidden_layers (int, default 1) – The number of hidden layers in the MLPs.

  • processor_layers (int, default 4) – The number of processor layers in the GNN.

  • mesh_aggr (str, default "sum") – The aggregation method for mesh nodes.

  • num_past_forcing_steps (int, default 1) – The number of past forcing steps to include.

  • num_future_forcing_steps (int, default 1) – The number of future forcing steps to include.

  • output_std (bool, default False) – Whether to output a predicted standard deviation.

  • output_clamping_lower (dict, optional) – Lower clamping limits for state variables.

  • output_clamping_upper (dict, optional) – Upper clamping limits for state variables.

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]][source]#

Run all processor steps (down then up at each layer depth) over the hierarchical mesh.

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.

  • 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.

  • 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.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep).

Return type:

tuple of (list, list, list, list)

make_down_gnns() → torch.nn.ModuleList[source]#

Make GNNs for processing steps down through the hierarchy.

Returns:

List of GNNs for each inter-level gap (downwards).

Return type:

nn.ModuleList

make_same_gnns() → torch.nn.ModuleList[source]#

Make intra-level GNNs.

Returns:

List of GNNs for each level.

Return type:

nn.ModuleList

make_up_gnns() → torch.nn.ModuleList[source]#

Make GNNs for processing steps up through the hierarchy.

Returns:

List of GNNs for each inter-level gap (upwards).

Return type:

nn.ModuleList

mesh_down_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_down_rep: list[torch.Tensor], down_gnns: torch.nn.ModuleList, same_gnns: torch.nn.ModuleList) → tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]][source]#

Run the downward pass of vertical processing, alternating between down-edges and same-level edges from the top to the bottom level.

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.

  • 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.

  • down_gnns (nn.ModuleList) – GNNs for downward edges, one per inter-level gap.

  • same_gnns (nn.ModuleList) – GNNs for same-level edges, one per level.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_down_rep).

Return type:

tuple of (list, list, list)

mesh_up_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_up_rep: list[torch.Tensor], up_gnns: torch.nn.ModuleList, same_gnns: torch.nn.ModuleList) → tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]][source]#

Run the upward pass of vertical processing, alternating between up-edges and same-level edges from the bottom to the top level.

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.

  • 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.

  • up_gnns (nn.ModuleList) – GNNs for upward edges, one per inter-level gap.

  • same_gnns (nn.ModuleList) – GNNs for same-level edges, one per level.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_up_rep).

Return type:

tuple of (list, list, list)

mesh_down_gnns#
mesh_down_same_gnns#
mesh_up_gnns#
mesh_up_same_gnns#