neural_lam.models.base_graph_model

neural_lam.models.base_graph_model#

Base classes for Neural-LAM graph models.

Module Contents#

class neural_lam.models.base_graph_model.BaseGraphModel(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: neural_lam.models.ar_model.ARModel

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

Initialize the graph-model scaffolding shared by concrete variants.

Parameters:

  • args (argparse.Namespace) – Training/runtime arguments describing graph paths and widths.

  • config (NeuralLAMConfig) – Experiment configuration for clamping and weighting.

  • datastore (BaseDatastore) – Datastore providing static features and metadata (e.g. graph path).

abstract embedd_mesh_nodes()#

Embed static mesh features for downstream processing.

Returns:

Embedded mesh node representations.

  • Shape: (num_mesh_nodes, d_h)

Return type:

torch.Tensor

get_clamped_new_state(state_delta, prev_state)#

Clamp predicted deltas and add them to the previous state.

The clamped values follow f(f^{-1}(X_t) + model({X_t, X_{t-1}, ...}, forcing)) so that the model learns to emit outputs in the range of the inverse clamping function.

Parameters:

  • state_delta (torch.Tensor) –

    Predicted change to apply to the previous state.

    • Shape: (B, num_grid_nodes, feature_dim)

  • prev_state (torch.Tensor) –

    Previous state X_t.

    • Shape: (B, num_grid_nodes, feature_dim)

Returns:

Clamped next state X_{t+1}.

  • Shape: (B, num_grid_nodes, feature_dim)

Return type:

torch.Tensor

abstract get_num_mesh()#

Compute mesh node counts used for encoding and decoding.

Returns:

Total number of mesh nodes and the number that should be ignored during encoding/decoding.

Return type:

tuple[int, int]

predict_step(prev_state, prev_prev_state, forcing)#

Advance the state by one step using the prediction model.

Parameters:

  • prev_state (torch.Tensor) –

    Current state X_t.

    • Shape: (B, num_grid_nodes, feature_dim)

  • prev_prev_state (torch.Tensor) –

    Previous state X_{t-1}.

    • Shape: (B, num_grid_nodes, feature_dim)

  • forcing (torch.Tensor) –

    Forcing inputs applied at the prediction step.

    • Shape: (B, num_grid_nodes, forcing_dim)

Returns:

Tuple (new_state, pred_std) where pred_std is None when the model does not emit uncertainty estimates.

  • Shape: (B, num_grid_nodes, feature_dim) for new_state and (B, num_grid_nodes, d_f) for pred_std when present.

Return type:

tuple[torch.Tensor, torch.Tensor | None]

prepare_clamping_params(config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Prepare per-feature parameters for clamping model outputs.

Parameters:

  • config (NeuralLAMConfig) – Model and training configuration containing clamping settings.

  • datastore (BaseDatastore) – Datastore that provides the ordering of state variables.

abstract process_step(mesh_rep)#

Run the processor portion of the encode-process-decode framework.

Parameters:

mesh_rep (torch.Tensor) –

Mesh node representations prior to the processor.

  • Shape: (B, num_mesh_nodes, d_h)

Returns:

Updated mesh representations after processing.

  • Shape: (B, num_mesh_nodes, d_h)

Return type:

torch.Tensor

encoding_grid_mlp#
g2m_embedder#
g2m_gnn#
grid_embedder#
m2g_embedder#
m2g_gnn#
mlp_blueprint_end#
output_map#
Contents