neural_lam.create_graph ======================= .. py:module:: neural_lam.create_graph .. autoapi-nested-parse:: Graph construction utilities for Neural-LAM meshes and grids. Attributes ---------- .. autoapisummary:: neural_lam.create_graph.CURRENT_GRAPH_SPEC_VERSION neural_lam.create_graph.METAINFO_FILENAME Functions --------- .. autoapisummary:: neural_lam.create_graph.assert_edge_index_in_range neural_lam.create_graph.cli neural_lam.create_graph.create_graph neural_lam.create_graph.create_graph_from_datastore neural_lam.create_graph.from_networkx_with_start_index neural_lam.create_graph.mk_2d_graph neural_lam.create_graph.plot_graph neural_lam.create_graph.prepend_node_index neural_lam.create_graph.save_edges neural_lam.create_graph.save_edges_list neural_lam.create_graph.sort_nodes_internally neural_lam.create_graph.zero_offset_edge_index Module Contents --------------- .. py:function:: assert_edge_index_in_range(edge_index: torch.Tensor, sender_range: tuple[int, int], receiver_range: tuple[int, int], name: str) -> None Assert that edge-index rows lie in expected half-open ranges. :param edge_index: Edge index tensor of shape ``[2, E]``. :type edge_index: torch.Tensor :param sender_range: Expected half-open range ``[start, stop)`` for sender indices in row 0. :type sender_range: tuple[int, int] :param receiver_range: Expected half-open range ``[start, stop)`` for receiver indices in row 1. :type receiver_range: tuple[int, int] :param name: Name used in assertion messages. :type name: str :rtype: None .. py:function:: cli(input_args: list[str] | None = None) -> None Parse CLI arguments and call :func:`create_graph_from_datastore`. :param input_args: Argument list forwarded to :class:`argparse.ArgumentParser`. When ``None``, ``sys.argv`` is used. :type input_args: list[str] or None, optional .. py:function:: create_graph(graph_dir_path: str, xy: numpy.ndarray, n_max_levels: int | None = None, hierarchical: bool = False, create_plot: bool = False) -> None Create graph components from `xy` grid coordinates and store in `graph_dir_path`. Creates the following files for all graphs: - g2m_edge_index.pt [2, N_g2m_edges] - g2m_features.pt [N_g2m_edges, d_features] - m2g_edge_index.pt [2, N_m2m_edges] - m2g_features.pt [N_m2m_edges, d_features] - m2m_edge_index.pt list of [2, N_m2m_edges_level], length==n_levels - m2m_features.pt list of [N_m2m_edges_level, d_features], length==n_levels - mesh_features.pt list of [N_mesh_nodes_level, num_mesh_static_vars], length==n_levels where d_features: number of features per edge (currently d_features==3, for edge-length, x and y) N_g2m_edges: number of edges in the graph from grid-to-mesh N_m2g_edges: number of edges in the graph from mesh-to-grid N_m2m_edges_level: number of edges in the graph from mesh-to-mesh at a given level (list index corresponds to the level) num_mesh_static_vars: number of static features per mesh node (currently num_mesh_static_vars==2, for x and y) N_mesh_nodes_level: number of nodes in the mesh at a given level And in addition for hierarchical graphs: - mesh_up_edge_index.pt: list of [2, N_mesh_updown_edges_level], length==n_levels-1 - mesh_up_features.pt: list of [N_mesh_updown_edges_level, d_features], length==n_levels-1 - mesh_down_edge_index.pt: list of [2, N_mesh_updown_edges_level], length==n_levels-1 - mesh_down_features.pt: list of [N_mesh_updown_edges_level, d_features], length==n_levels-1 where N_mesh_updown_edges_level is the number of edges in the graph from mesh-to-mesh between two consecutive levels (list index corresponds index of lower level) :param graph_dir_path: Path to store the graph components. :type graph_dir_path: str :param xy: Grid coordinates, expected to be of shape (Nx, Ny, 2). :type xy: np.ndarray :param n_max_levels: Limit multi-scale mesh to given number of levels, from bottom up (default: None (no limit)). :type n_max_levels: int :param hierarchical: Generate hierarchical mesh graph (default: False). :type hierarchical: bool :param create_plot: If graphs should be plotted during generation (default: False). :type create_plot: bool :rtype: None .. py:function:: create_graph_from_datastore(datastore: neural_lam.datastore.base.BaseDatastore, output_root_path: str, n_max_levels: int | None = None, hierarchical: bool = False, create_plot: bool = False) -> None Generate graph components for ``datastore`` and persist them on disk. :param datastore: Datastore providing ``get_xy`` for state nodes. :type datastore: BaseDatastore :param output_root_path: Directory where the resulting ``*.pt`` graph files are stored. :type output_root_path: str :param n_max_levels: Optional limit of hierarchical mesh levels to build. :type n_max_levels: int or None, optional :param hierarchical: If ``True``, create multi-level hierarchical graphs. Default ``False``. :type hierarchical: bool, optional :param create_plot: If ``True``, display matplotlib previews of the generated graphs. :type create_plot: bool, optional .. py:function:: from_networkx_with_start_index(nx_graph: networkx.Graph, start_index: int) -> torch_geometric.data.Data Convert a NetworkX graph to PyG and shift its edge indices. ``from_networkx`` maps NetworkX node labels to contiguous PyG node indices starting at 0. This helper shifts those indices into the combined node-index space used internally while constructing graph components. The known offsets are removed again before current-format graph tensors are saved. :param nx_graph: Graph to convert. :type nx_graph: networkx.Graph :param start_index: Offset added to both rows of the PyG edge-index tensor. :type start_index: int :returns: Converted PyG graph with shifted edge indices. :rtype: pyg.data.Data .. py:function:: mk_2d_graph(xy: numpy.ndarray, nx: int, ny: int) -> networkx.DiGraph Create a diagonal 2-D grid graph over the ``xy`` positions. :param xy: The grid coordinates. :type xy: np.ndarray :param nx: Number of nodes in the x-dimension. :type nx: int :param ny: Number of nodes in the y-dimension. :type ny: int :returns: The constructed directed 2-D grid graph. :rtype: networkx.DiGraph .. py:function:: plot_graph(graph: torch_geometric.data.Data, title: str | None = None) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes] Render a PyTorch Geometric graph using stored node coordinates. :param graph: Graph containing ``edge_index`` and ``pos`` attributes. :type graph: torch_geometric.data.Data :param title: Optional subplot title. :type title: str or None, optional :returns: Figure and axis handles for further customization. :rtype: tuple[matplotlib.figure.Figure, matplotlib.axes.Axes] .. py:function:: prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph Relabel each node by prepending ``new_index`` to its tuple identifier. :param graph: The graph to relabel. :type graph: networkx.Graph :param new_index: The value to prepend to each node identifier. :type new_index: int :returns: The relabeled graph. :rtype: networkx.Graph .. py:function:: save_edges(graph: torch_geometric.data.Data, name: str, base_path: str) -> None Persist edge indices/features for a PyG graph under ``base_path``. :param graph: The graph containing edge data. :type graph: torch_geometric.data.Data :param name: The name prefix for the saved files. :type name: str :param base_path: The directory path where files should be saved. :type base_path: str .. py:function:: save_edges_list(graphs: list[torch_geometric.data.Data], name: str, base_path: str) -> None Persist edge indices/features for a list of graphs. :param graphs: The list of graphs containing edge data. :type graphs: list of torch_geometric.data.Data :param name: The name prefix for the saved files. :type name: str :param base_path: The directory path where files should be saved. :type base_path: str .. py:function:: sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph Return a copy of ``nx_graph`` with deterministically ordered nodes. :param nx_graph: The input graph to sort nodes for. :type nx_graph: networkx.Graph :returns: A directed graph with nodes sorted alphabetically by their labels. :rtype: networkx.DiGraph .. py:function:: zero_offset_edge_index(edge_index: torch.Tensor, sender_offset: int, receiver_offset: int) -> torch.Tensor Subtract known sender and receiver offsets from an edge index tensor. :param edge_index: Edge index tensor of shape ``[2, E]`` in the combined node-index space. :type edge_index: torch.Tensor :param sender_offset: Offset to subtract from sender indices in row 0. :type sender_offset: int :param receiver_offset: Offset to subtract from receiver indices in row 1. :type receiver_offset: int :returns: Edge index tensor with row-wise offsets removed. :rtype: torch.Tensor .. py:data:: CURRENT_GRAPH_SPEC_VERSION :value: '0.1.0' .. py:data:: METAINFO_FILENAME :value: 'metainfo.yaml'