neural_lam.models.hi_lam

neural_lam.models.hi_lam#

Sequential up/down hierarchical Neural-LAM model (Hi-LAM).

Module Contents#

class neural_lam.models.hi_lam.HiLAM(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: neural_lam.models.base_hi_graph_model.BaseHiGraphModel

Hierarchical graph model with message passing that goes sequentially down and up the hierarchy during processing. The Hi-LAM model from Oskarsson et al. (2023)

Initialize the sequential up/down hierarchical processor.

hi_processor_step(mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep)#

Execute one full processor iteration (down + up sweeps).

Parameters:

  • mesh_rep_levels (list[torch.Tensor]) –

    Mesh representations for each level.

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

  • mesh_same_rep (list[torch.Tensor]) –

    Same-level edge representations.

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

  • mesh_up_rep (list[torch.Tensor]) –

    Upward edge representations.

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

  • mesh_down_rep (list[torch.Tensor]) –

    Downward edge representations.

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

Returns:

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

  • ] – Updated representations (mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep) after both sweeps.

make_down_gnns(args)#

Make GNNs for processing steps down through the hierarchy.

make_same_gnns(args)#

Make intra-level GNNs.

make_up_gnns(args)#

Make GNNs for processing steps up through the hierarchy.

mesh_down_step(mesh_rep_levels, mesh_same_rep, mesh_down_rep, down_gnns, same_gnns)#

Run the downward half of the hierarchical processing sweep.

Parameters:

  • mesh_rep_levels (list[torch.Tensor]) –

    Mesh representations for each level.

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

  • mesh_same_rep (list[torch.Tensor]) –

    Same-level edge representations.

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

  • mesh_down_rep (list[torch.Tensor]) –

    Downward edge representations.

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

  • down_gnns (Sequence[InteractionNet]) – Message-passing networks applied to downward edges.

  • same_gnns (Sequence[InteractionNet]) – Message-passing networks for same-level processing.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_down_rep).

Return type:

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

mesh_up_step(mesh_rep_levels, mesh_same_rep, mesh_up_rep, up_gnns, same_gnns)#

Run the upward half of the hierarchical processing sweep.

Parameters:

  • mesh_rep_levels (list[torch.Tensor]) –

    Mesh representations for each level.

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

  • mesh_same_rep (list[torch.Tensor]) –

    Same-level edge representations.

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

  • mesh_up_rep (list[torch.Tensor]) –

    Upward edge representations.

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

  • up_gnns (Sequence[InteractionNet]) – Message-passing networks applied to upward edges.

  • same_gnns (Sequence[InteractionNet]) – Message-passing networks for same-level processing.

Returns:

Updated (mesh_rep_levels, mesh_same_rep, mesh_up_rep).

Return type:

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

mesh_down_gnns#
mesh_down_same_gnns#
mesh_up_gnns#
mesh_up_same_gnns#
Contents