neural_lam.models.graph_lam

neural_lam.models.graph_lam#

GraphLAM: the non-hierarchical Neural-LAM architecture.

Module Contents#

class neural_lam.models.graph_lam.GraphLAM(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: neural_lam.models.base_graph_model.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 non-hierarchical GraphLAM variant.

embedd_mesh_nodes()#

Embed static mesh features.

Returns:

Embedded mesh node representations.

  • Shape: (N_mesh, d_h)

Return type:

torch.Tensor

get_num_mesh()#

Compute mesh node counts used for encoding and decoding.

Returns:

Total number of mesh nodes and the number to ignore during encoding/decoding.

Return type:

tuple[int, int]

process_step(mesh_rep)#

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

Parameters:

mesh_rep (torch.Tensor) –

Mesh node representations before processing.

  • Shape: (B, N_mesh, d_h)

Returns:

Updated mesh representations.

  • Shape: (B, N_mesh, d_h)

Return type:

torch.Tensor

m2m_embedder#
mesh_embedder#
processor#
Contents