neural_lam.create_graph

neural_lam.create_graph#

Graph construction utilities for Neural-LAM meshes and grids.

Attributes#

CURRENT_GRAPH_SPEC_VERSION

METAINFO_FILENAME

Functions#

assert_edge_index_in_range(→ None)

Assert that edge-index rows lie in expected half-open ranges.

cli(→ None)

Parse CLI arguments and call create_graph_from_datastore().

create_graph(→ None)

Create graph components from xy grid coordinates and store in

create_graph_from_datastore(→ None)

Generate graph components for datastore and persist them on disk.

from_networkx_with_start_index(→ torch_geometric.data.Data)

Convert a NetworkX graph to PyG and shift its edge indices.

mk_2d_graph(→ networkx.DiGraph)

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

plot_graph(→ tuple[matplotlib.figure.Figure, ...)

Render a PyTorch Geometric graph using stored node coordinates.

prepend_node_index(→ networkx.Graph)

Relabel each node by prepending new_index to its tuple identifier.

save_edges(→ None)

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

save_edges_list(→ None)

Persist edge indices/features for a list of graphs.

sort_nodes_internally(→ networkx.DiGraph)

Return a copy of nx_graph with deterministically ordered nodes.

zero_offset_edge_index(→ torch.Tensor)

Subtract known sender and receiver offsets from an edge index tensor.

Module Contents#

neural_lam.create_graph.assert_edge_index_in_range(edge_index: torch.Tensor, sender_range: tuple[int, int], receiver_range: tuple[int, int], name: str) → None[source]#

Assert that edge-index rows lie in expected half-open ranges.

Parameters:

  • edge_index (torch.Tensor) – Edge index tensor of shape [2, E].

  • sender_range (tuple[int, int]) – Expected half-open range [start, stop) for sender indices in row 0.

  • receiver_range (tuple[int, int]) – Expected half-open range [start, stop) for receiver indices in row 1.

  • name (str) – Name used in assertion messages.

Return type:

None

neural_lam.create_graph.cli(input_args: list[str] | None = None) → None[source]#

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 = False, create_plot: bool = False) → None[source]#

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)

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.BaseDatastore, output_root_path: str, n_max_levels: int | None = None, hierarchical: bool = False, create_plot: bool = False) → None[source]#

Generate graph components for datastore and persist them on disk.

Parameters:

  • datastore (BaseDatastore) – 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: networkx.Graph, start_index: int) → torch_geometric.data.Data[source]#

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.

Parameters:

  • nx_graph (networkx.Graph) – Graph to convert.

  • start_index (int) – Offset added to both rows of the PyG edge-index tensor.

Returns:

Converted PyG graph with shifted edge indices.

Return type:

pyg.data.Data

neural_lam.create_graph.mk_2d_graph(xy: numpy.ndarray, nx: int, ny: int) → networkx.DiGraph[source]#

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

Parameters:

  • xy (np.ndarray) – The grid coordinates.

  • nx (int) – Number of nodes in the x-dimension.

  • ny (int) – Number of nodes in the y-dimension.

Returns:

The constructed directed 2-D grid graph.

Return type:

networkx.DiGraph

neural_lam.create_graph.plot_graph(graph: torch_geometric.data.Data, title: str | None = None) → tuple[matplotlib.figure.Figure, matplotlib.axes.Axes][source]#

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: networkx.Graph, new_index: int) → networkx.Graph[source]#

Relabel each node by prepending new_index to its tuple identifier.

Parameters:

  • graph (networkx.Graph) – The graph to relabel.

  • new_index (int) – The value to prepend to each node identifier.

Returns:

The relabeled graph.

Return type:

networkx.Graph

neural_lam.create_graph.save_edges(graph: torch_geometric.data.Data, name: str, base_path: str) → None[source]#

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

Parameters:

  • graph (torch_geometric.data.Data) – The graph containing edge data.

  • name (str) – The name prefix for the saved files.

  • base_path (str) – The directory path where files should be saved.

neural_lam.create_graph.save_edges_list(graphs: list[torch_geometric.data.Data], name: str, base_path: str) → None[source]#

Persist edge indices/features for a list of graphs.

Parameters:

  • graphs (list of torch_geometric.data.Data) – The list of graphs containing edge data.

  • name (str) – The name prefix for the saved files.

  • base_path (str) – The directory path where files should be saved.

neural_lam.create_graph.sort_nodes_internally(nx_graph: networkx.Graph) → networkx.DiGraph[source]#

Return a copy of nx_graph with deterministically ordered nodes.

Parameters:

nx_graph (networkx.Graph) – The input graph to sort nodes for.

Returns:

A directed graph with nodes sorted alphabetically by their labels.

Return type:

networkx.DiGraph

neural_lam.create_graph.zero_offset_edge_index(edge_index: torch.Tensor, sender_offset: int, receiver_offset: int) → torch.Tensor[source]#

Subtract known sender and receiver offsets from an edge index tensor.

Parameters:

  • edge_index (torch.Tensor) – Edge index tensor of shape [2, E] in the combined node-index space.

  • sender_offset (int) – Offset to subtract from sender indices in row 0.

  • receiver_offset (int) – Offset to subtract from receiver indices in row 1.

Returns:

Edge index tensor with row-wise offsets removed.

Return type:

torch.Tensor

neural_lam.create_graph.CURRENT_GRAPH_SPEC_VERSION = '0.1.0'#
neural_lam.create_graph.METAINFO_FILENAME = 'metainfo.yaml'#