neural_lam.models.step_predictors.graph.hi_lam_parallel

neural_lam.models.step_predictors.graph.hi_lam_parallel#

Parallel hierarchical graph-based LAM model.

Classes#

HiLAMParallel

Version of HiLAM where all message passing in the hierarchical mesh (up,

Module Contents#

class neural_lam.models.step_predictors.graph.hi_lam_parallel.HiLAMParallel(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

Version of HiLAM where all message passing in the hierarchical mesh (up, down, inter-level) is ran in parallel.

This is a somewhat simpler alternative to the sequential message passing of Hi-LAM.

Initialize the HiLAMParallel 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 in parallel across all edge types and hierarchy levels.

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)

edge_split_sections#