neural_lam.models.step_predictors.graph.graph_lam

neural_lam.models.step_predictors.graph.graph_lam#

Graph-based LAM model with a flat mesh.

Classes#

GraphLAM

Full graph-based LAM model that can be used with different

Module Contents#

class neural_lam.models.step_predictors.graph.graph_lam.GraphLAM(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', **_kwargs: object)[source]#

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

Full graph-based LAM model that can be used with different (non-hierarchical )graphs. Mainly based on GraphCast, but the model from Keisler (2022) is almost identical. Used for GC-LAM and L1-LAM in Oskarsson et al. (2023).

Initialize the GraphLAM 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.

embedd_mesh_nodes() → torch.Tensor[source]#

Embed static mesh node features.

Returns:

Shape (num_mesh_nodes, hidden_dim). Embedded mesh node representations. Dims: num_mesh_nodes is the number of mesh nodes and hidden_dim is the hidden dimension.

Return type:

torch.Tensor

get_num_mesh() → tuple[int, int][source]#

Compute number of mesh nodes from loaded features, and number of mesh nodes that should be ignored in encoding/decoding.

Returns:

  • num_mesh_nodes (int) – The number of mesh nodes.

  • num_ignore_mesh_nodes (int) – The number of mesh nodes to ignore.

process_step(mesh_rep: torch.Tensor) → torch.Tensor[source]#

Process the mesh representation through the flat message-passing processor (all nodes at a single level).

Parameters:

mesh_rep (torch.Tensor) – Shape (B, num_mesh_nodes, hidden_dim). Current mesh node representations. Dims: B is batch size, num_mesh_nodes is the number of mesh nodes, and hidden_dim is the hidden dimension.

Returns:

Shape (B, num_mesh_nodes, hidden_dim). Updated mesh node representations. Dims: same as mesh_rep.

Return type:

torch.Tensor

m2m_edge_index: torch.Tensor#
m2m_embedder#
m2m_features: torch.Tensor#
mesh_embedder#
mesh_static_features: torch.Tensor#
processor#