neural_lam.models.step_predictors.graph.base

neural_lam.models.step_predictors.graph.base#

Base class for graph-based step predictors.

Classes#

BaseGraphModel

Base (abstract) class for graph-based models building on

Module Contents#

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

Bases: neural_lam.models.step_predictors.base.StepPredictor

Base (abstract) class for graph-based models building on the encode-process-decode idea.

Initialize the BaseGraphModel.

Parameters:

datastore (BaseDatastore) – Datastore supplying data and information about dataset and forecast region.
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.

abstract 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

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

Advance the state by one step using the encode-process-decode graph pipeline: embed grid + edge features, map grid -> mesh via the g2m GNN, run the (subclass-defined) processor on the mesh, decode mesh -> grid via the m2g GNN, project to a state delta, rescale with difference statistics, and add to prev_state (with optional clamping). Returns (X_{t+1}, optional std).

Parameters:

prev_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The current state X_t. Dims: B is batch size, num_grid_nodes is the number of spatial grid nodes, and num_state_vars is the number of state variables.
prev_prev_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The previous state X_{t-1}, used as additional conditioning. Dims: same as prev_state.
forcing (torch.Tensor) – Shape (B, num_grid_nodes, num_forcing_vars). External forcings for this step (already concatenated past/current/future windows). Dims: B is batch size, num_grid_nodes is the number of spatial grid nodes, and num_forcing_vars is the forcing feature dimension.

Returns:

new_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The predicted next state X_{t+1} after delta-add and clamping. Dims: same as prev_state.
pred_std (torch.Tensor or None) – Shape (B, num_grid_nodes, num_state_vars) when output_std is True, otherwise None. Per-feature predicted standard deviation (raw softplus output, not rescaled by diff statistics). Dims: same as prev_state.

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

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

Process the mesh representation in the encode-process-decode framework, running one or more message-passing steps.

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