neural_lam.create_graph

neural_lam.create_graph#

Graph construction utilities for Neural-LAM meshes and grids.

Module Contents#

neural_lam.create_graph.cli(input_args=None)#

Parse CLI arguments and call create_graph_from_datastore().

Parameters:

input_args (list[str] or None, optional) – Argument list forwarded to argparse.ArgumentParser. When None, sys.argv is used.

neural_lam.create_graph.create_graph(graph_dir_path: str, xy: numpy.ndarray, n_max_levels: int | None = None, hierarchical: bool | None = False, create_plot: bool | None = False)#

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, d_mesh_static],

    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)

d_mesh_static:

number of static features per mesh node (currently d_mesh_static==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)

Parameters:

  • graph_dir_path (str) – Path to store the graph components.

  • xy (np.ndarray) – Grid coordinates, expected to be of shape (Nx, Ny, 2).

  • n_max_levels (int) – Limit multi-scale mesh to given number of levels, from bottom up (default: None (no limit)).

  • hierarchical (bool) – Generate hierarchical mesh graph (default: False).

  • create_plot (bool) – If graphs should be plotted during generation (default: False).

Return type:

None

neural_lam.create_graph.create_graph_from_datastore(datastore: neural_lam.datastore.base.BaseRegularGridDatastore, output_root_path: str, n_max_levels: int | None = None, hierarchical: bool = False, create_plot: bool = False)#

Generate graph components for datastore and persist them on disk.

Parameters:

  • datastore (BaseRegularGridDatastore) – Datastore providing get_xy for state nodes.

  • output_root_path (str) – Directory where the resulting *.pt graph files are stored.

  • n_max_levels (int or None, optional) – Optional limit of hierarchical mesh levels to build.

  • hierarchical (bool, optional) – If True, create multi-level hierarchical graphs. Default False.

  • create_plot (bool, optional) – If True, display matplotlib previews of the generated graphs.

neural_lam.create_graph.from_networkx_with_start_index(nx_graph, start_index)#

Convert a NetworkX graph to PyG and offset node indices.

neural_lam.create_graph.mk_2d_graph(xy, nx, ny)#

Create a diagonal 2-D grid graph over the xy positions.

neural_lam.create_graph.plot_graph(graph, title=None)#

Render a PyTorch Geometric graph using stored node coordinates.

Parameters:

  • graph (torch_geometric.data.Data) – Graph containing edge_index and pos attributes.

  • title (str or None, optional) – Optional subplot title.

Returns:

Figure and axis handles for further customization.

Return type:

tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]

neural_lam.create_graph.prepend_node_index(graph, new_index)#

Relabel each node by prepending new_index to its tuple identifier.

neural_lam.create_graph.save_edges(graph, name, base_path)#

Persist edge indices/features for a PyG graph under base_path.

neural_lam.create_graph.save_edges_list(graphs, name, base_path)#

Persist edge indices/features for a list of graphs.

neural_lam.create_graph.sort_nodes_internally(nx_graph)#

Return a copy of nx_graph with deterministically ordered nodes.

Contents