neural_lam.models.base_hi_graph_model

neural_lam.models.base_hi_graph_model#

Base implementations for hierarchical (multi-level) graph models.

Module Contents#

class neural_lam.models.base_hi_graph_model.BaseHiGraphModel(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: neural_lam.models.base_graph_model.BaseGraphModel

Base class for hierarchical graph models.

Extend BaseGraphModel with hierarchical mesh structures.

embedd_mesh_nodes()#

Embed static mesh features for the bottom level of the hierarchy.

Returns:

Embedded representations for the base-level mesh nodes.

  • Shape: (num_mesh_nodes[0], 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]

abstract hi_processor_step(mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep)#

Internal processor step executed between mesh init and read-out.

Parameters:

  • mesh_rep_levels (list[torch.Tensor]) –

    Mesh representations for each level.

    • Shape: (B, num_mesh_nodes[l], d_h)

  • mesh_same_rep (list[torch.Tensor]) –

    Same-level edge representations per level.

    • Shape: (B, M_same[l], d_h)

  • mesh_up_rep (list[torch.Tensor]) –

    Edge representations from level l to l+1.

    • Shape: (B, M_up[l -> l+1], d_h)

  • mesh_down_rep (list[torch.Tensor]) –

    Edge representations from level l+1 down to l.

    • Shape: (B, M_down[l <- l+1], d_h)

Returns:

  • tuple[ – list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]

  • ] – Updated representations for (mesh, same-level, up edges, down edges) in that order.

process_step(mesh_rep)#

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

Parameters:

mesh_rep (torch.Tensor) –

Base-level mesh representations prior to the processor.

  • Shape: (B, num_mesh_nodes, d_h)

Returns:

Updated base-level mesh representations.

  • Shape: (B, num_mesh_nodes, d_h)

Return type:

torch.Tensor

level_mesh_sizes#
mesh_down_embedders#
mesh_embedders#
mesh_init_gnns#
mesh_read_gnns#
mesh_same_embedders#
mesh_up_embedders#
num_levels#
Contents