neural_lam.utils.graph ====================== .. py:module:: neural_lam.utils.graph .. autoapi-nested-parse:: Loading and zero-indexing of mesh/grid graph tensors. Attributes ---------- .. autoapisummary:: neural_lam.utils.graph.LEGACY_GRAPH_SPEC_VERSION Functions --------- .. autoapisummary:: neural_lam.utils.graph.compute_grid_input_dim neural_lam.utils.graph.load_and_register_graph neural_lam.utils.graph.load_graph neural_lam.utils.graph.zero_index_edge_index neural_lam.utils.graph.zero_index_g2m neural_lam.utils.graph.zero_index_m2g Module Contents --------------- .. py:function:: compute_grid_input_dim(datastore: neural_lam.datastore.BaseDatastore, num_past_forcing_steps: int, num_future_forcing_steps: int) -> int Compute the total grid input dimensionality of a graph step predictor. The grid input concatenates the two previous states, the grid static features and the windowed forcing (past + current + future forcing steps). :param datastore: Datastore providing the number of state, static and forcing variables. :type datastore: BaseDatastore :param num_past_forcing_steps: Number of past forcing steps included in the input window. :type num_past_forcing_steps: int :param num_future_forcing_steps: Number of future forcing steps included in the input window. :type num_future_forcing_steps: int :returns: Total grid input dimensionality. :rtype: int .. py:function:: load_and_register_graph(module: torch.nn.Module, datastore: neural_lam.datastore.BaseDatastore, graph_name: str, mesh_node_features_scaling: float) -> bool Load a graph and register its tensors on ``module``. Loads the graph ``graph_name`` from the datastore's graph directory via :func:`load_graph`, then registers each tensor as a non-persistent buffer and each non-tensor (e.g. ``BufferList``) as a plain attribute on ``module``. :param module: Module to register the graph tensors and attributes on, in place. :type module: torch.nn.Module :param datastore: Datastore whose ``root_path`` holds the ``graph`` directory. :type datastore: BaseDatastore :param graph_name: Name of the graph directory (under ``/graph``) to load. :type graph_name: str :param mesh_node_features_scaling: Scalar used to normalize mesh node coordinate features for graphs in the current on-disk format; forwarded to :func:`load_graph`. :type mesh_node_features_scaling: float :returns: Whether the loaded graph is hierarchical. :rtype: bool .. py:function:: load_graph(graph_dir_path: str | pathlib.Path, mesh_node_features_scaling: float, device: str = 'cpu') -> tuple[bool, dict[str, Any]] Load all tensors representing the graph from `graph_dir_path`. Needs the following files for all graphs: - m2m_edge_index.pt - g2m_edge_index.pt - m2g_edge_index.pt - m2m_features.pt - g2m_features.pt - m2g_features.pt - mesh_features.pt And in addition for hierarchical graphs: - mesh_up_edge_index.pt - mesh_down_edge_index.pt - mesh_up_features.pt - mesh_down_features.pt :param graph_dir_path: Path to directory containing the graph files. :type graph_dir_path: str :param mesh_node_features_scaling: Scalar used to normalize mesh node coordinate features for graphs in the current on-disk format. :type mesh_node_features_scaling: float :param device: Device to load tensors to. :type device: str :returns: * **hierarchical** (*bool*) -- Whether the graph is hierarchical. * **graph** (*dict*) -- Dictionary containing the graph tensors, with keys as follows: - g2m_edge_index - m2g_edge_index - m2m_edge_index - mesh_up_edge_index - mesh_down_edge_index - g2m_features - m2g_features - m2m_features - mesh_up_features - mesh_down_features - mesh_static_features .. py:function:: zero_index_edge_index(edge_index: torch.Tensor) -> torch.Tensor Make both sender and receiver indices of edge_index start at 0. :param edge_index: Edge index tensor of shape (2, num_edges). :type edge_index: torch.Tensor :returns: Edge index tensor with indices starting at 0. :rtype: torch.Tensor .. py:function:: zero_index_g2m(g2m_edge_index: torch.Tensor, mesh_static_features: list[torch.Tensor], mesh_first: bool, restore: bool = False) -> torch.Tensor Zero-index the g2m (grid-to-mesh) edge index, or undo this operation. Special handling is needed since not all mesh nodes may be present. :param g2m_edge_index: Edge index tensor of shape (2, num_edges). :type g2m_edge_index: torch.Tensor :param mesh_static_features: Mesh node feature tensors. :type mesh_static_features: list of torch.Tensor :param mesh_first: If True, mesh nodes are indexed before grid nodes. :type mesh_first: bool :param restore: If True, undo zero-indexing (restore original indices). :type restore: bool :returns: Edge index tensor with zero-based or restored indices. :rtype: torch.Tensor .. py:function:: zero_index_m2g(m2g_edge_index: torch.Tensor, mesh_static_features: list[torch.Tensor], mesh_first: bool, restore: bool = False) -> torch.Tensor Zero-index the m2g (mesh-to-grid) edge index, or undo this operation. Special handling is needed since not all mesh nodes may be present. :param m2g_edge_index: Edge index tensor of shape (2, num_edges). :type m2g_edge_index: torch.Tensor :param mesh_static_features: Mesh node feature tensors. :type mesh_static_features: list of torch.Tensor :param mesh_first: If True, mesh nodes are indexed before grid nodes. :type mesh_first: bool :param restore: If True, undo zero-indexing (restore original indices). :type restore: bool :returns: Edge index tensor with zero-based or restored indices. :rtype: torch.Tensor .. py:data:: LEGACY_GRAPH_SPEC_VERSION :value: 'legacy'