neural_lam.models.step_predictors.graph.hi_lam ============================================== .. py:module:: neural_lam.models.step_predictors.graph.hi_lam .. autoapi-nested-parse:: Hierarchical graph-based LAM model. Classes ------- .. autoapisummary:: neural_lam.models.step_predictors.graph.hi_lam.HiLAM Module Contents --------------- .. py:class:: HiLAM(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', mesh_up_gnn_type: str = 'InteractionNet', mesh_down_gnn_type: str = 'InteractionNet') Bases: :py:obj:`neural_lam.models.step_predictors.graph.hierarchical.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 HiLAM model. :param datastore: The datastore providing grid metadata and data access. :type datastore: BaseDatastore :param graph_name: The name of the graph to load. :type graph_name: str, default "multiscale" :param hidden_dim: The dimension of the hidden representations. :type hidden_dim: int, default 64 :param hidden_layers: The number of hidden layers in the MLPs. :type hidden_layers: int, default 1 :param processor_layers: The number of processor layers in the GNN. :type processor_layers: int, default 4 :param mesh_aggr: The aggregation method for mesh nodes. :type mesh_aggr: str, default "sum" :param num_past_forcing_steps: The number of past forcing steps to include. :type num_past_forcing_steps: int, default 1 :param num_future_forcing_steps: The number of future forcing steps to include. :type num_future_forcing_steps: int, default 1 :param output_std: Whether to output a predicted standard deviation. :type output_std: bool, default False :param output_clamping_lower: Lower clamping limits for state variables. :type output_clamping_lower: dict, optional :param output_clamping_upper: Upper clamping limits for state variables. :type output_clamping_upper: dict, optional .. py:method:: hi_processor_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_up_rep: list[torch.Tensor], mesh_down_rep: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]] Run all processor steps (down then up at each layer depth) over the hierarchical mesh. :param mesh_rep_levels: One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at each hierarchy level. Dims: ``B`` is batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, and ``hidden_dim`` is the hidden dimension. :type mesh_rep_levels: list of torch.Tensor :param mesh_same_rep: One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``. Same-level edge representations. :type mesh_same_rep: list of torch.Tensor :param mesh_up_rep: One tensor per inter-level gap, each of shape ``(B, num_edges[l], hidden_dim)``. Upward edge representations. :type mesh_up_rep: list of torch.Tensor :param mesh_down_rep: One tensor per inter-level gap, each of shape ``(B, num_edges[l], hidden_dim)``. Downward edge representations. :type mesh_down_rep: list of torch.Tensor :returns: Updated ``(mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep)``. :rtype: tuple of (list, list, list, list) .. py:method:: make_down_gnns() -> torch.nn.ModuleList Make GNNs for processing steps down through the hierarchy. :returns: List of GNNs for each inter-level gap (downwards). :rtype: nn.ModuleList .. py:method:: make_same_gnns() -> torch.nn.ModuleList Make intra-level GNNs. :returns: List of GNNs for each level. :rtype: nn.ModuleList .. py:method:: make_up_gnns() -> torch.nn.ModuleList Make GNNs for processing steps up through the hierarchy. :returns: List of GNNs for each inter-level gap (upwards). :rtype: nn.ModuleList .. py:method:: mesh_down_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_down_rep: list[torch.Tensor], down_gnns: torch.nn.ModuleList, same_gnns: torch.nn.ModuleList) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]] Run the downward pass of vertical processing, alternating between down-edges and same-level edges from the top to the bottom level. :param mesh_rep_levels: One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at each hierarchy level. Dims: ``B`` is batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, and ``hidden_dim`` is the hidden dimension. :type mesh_rep_levels: list of torch.Tensor :param mesh_same_rep: One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``. Same-level edge representations. :type mesh_same_rep: list of torch.Tensor :param mesh_down_rep: One tensor per inter-level gap, each of shape ``(B, num_edges[l], hidden_dim)``. Downward edge representations from level ``l+1`` to ``l``. :type mesh_down_rep: list of torch.Tensor :param down_gnns: GNNs for downward edges, one per inter-level gap. :type down_gnns: nn.ModuleList :param same_gnns: GNNs for same-level edges, one per level. :type same_gnns: nn.ModuleList :returns: Updated ``(mesh_rep_levels, mesh_same_rep, mesh_down_rep)``. :rtype: tuple of (list, list, list) .. py:method:: mesh_up_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_up_rep: list[torch.Tensor], up_gnns: torch.nn.ModuleList, same_gnns: torch.nn.ModuleList) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]] Run the upward pass of vertical processing, alternating between up-edges and same-level edges from the bottom to the top level. :param mesh_rep_levels: One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at each hierarchy level. Dims: ``B`` is batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, and ``hidden_dim`` is the hidden dimension. :type mesh_rep_levels: list of torch.Tensor :param mesh_same_rep: One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``. Same-level edge representations. :type mesh_same_rep: list of torch.Tensor :param mesh_up_rep: One tensor per inter-level gap, each of shape ``(B, num_edges[l], hidden_dim)``. Upward edge representations from level ``l`` to ``l+1``. :type mesh_up_rep: list of torch.Tensor :param up_gnns: GNNs for upward edges, one per inter-level gap. :type up_gnns: nn.ModuleList :param same_gnns: GNNs for same-level edges, one per level. :type same_gnns: nn.ModuleList :returns: Updated ``(mesh_rep_levels, mesh_same_rep, mesh_up_rep)``. :rtype: tuple of (list, list, list) .. py:attribute:: mesh_down_gnns .. py:attribute:: mesh_down_same_gnns .. py:attribute:: mesh_up_gnns .. py:attribute:: mesh_up_same_gnns