"""Graph construction utilities for Neural-LAM meshes and grids."""
# Standard library
import os
import warnings
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
# Third-party
import matplotlib
import matplotlib.pyplot as plt
import networkx
import numpy as np
import scipy.spatial
import torch
import torch_geometric as pyg
import yaml
from loguru import logger
from torch_geometric.utils.convert import from_networkx
# Local
from .config import load_config_and_datastore
from .datastore.base import BaseDatastore
# Stores the graph storage spec version the graph conforms to.
METAINFO_FILENAME = "metainfo.yaml"
CURRENT_GRAPH_SPEC_VERSION = "0.1.0"
[docs]
def plot_graph(
graph: pyg.data.Data, title: str | None = None
) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]:
"""
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
-------
tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]
Figure and axis handles for further customization.
"""
fig, axis = plt.subplots(figsize=(8, 8), dpi=200) # W,H
edge_index = graph.edge_index
pos = graph.pos
# Fix for re-indexed edge indices only containing mesh nodes at
# higher levels in hierarchy
edge_index = edge_index - edge_index.min()
if pyg.utils.is_undirected(edge_index):
# Keep only 1 direction of edge_index
edge_index = edge_index[:, edge_index[0] < edge_index[1]] # (2, M/2)
# TODO: indicate direction of directed edges
# Move all to cpu and numpy, compute (in)-degrees
degrees = (
pyg.utils.degree(edge_index[1], num_nodes=pos.shape[0]).cpu().numpy()
)
edge_index = edge_index.cpu().numpy()
pos = pos.cpu().numpy()
# Plot edges
from_pos = pos[edge_index[0]] # (M/2, 2)
to_pos = pos[edge_index[1]] # (M/2, 2)
edge_lines = np.stack((from_pos, to_pos), axis=1)
axis.add_collection(
matplotlib.collections.LineCollection(
edge_lines, lw=0.4, colors="black", zorder=1
)
)
# Plot nodes
node_scatter = axis.scatter(
pos[:, 0],
pos[:, 1],
c=degrees,
s=3,
marker="o",
zorder=2,
cmap="viridis",
clim=None,
)
plt.colorbar(node_scatter, aspect=50)
if title is not None:
axis.set_title(title)
return fig, axis
[docs]
def sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph:
"""
Return a copy of ``nx_graph`` with deterministically ordered nodes.
Parameters
----------
nx_graph : networkx.Graph
The input graph to sort nodes for.
Returns
-------
networkx.DiGraph
A directed graph with nodes sorted alphabetically by their labels.
"""
# For some reason the networkx .nodes() return list can not be sorted,
# but this is the ordering used by pyg when converting.
# This function fixes this.
H = networkx.DiGraph()
H.add_nodes_from(sorted(nx_graph.nodes(data=True)))
H.add_edges_from(nx_graph.edges(data=True))
return H
[docs]
def save_edges(graph: pyg.data.Data, name: str, base_path: str) -> None:
"""
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.
"""
torch.save(
graph.edge_index, os.path.join(base_path, f"{name}_edge_index.pt")
)
edge_features = torch.cat((graph.len.unsqueeze(1), graph.vdiff), dim=1).to(
torch.float32
) # Save as float32
torch.save(edge_features, os.path.join(base_path, f"{name}_features.pt"))
[docs]
def save_edges_list(
graphs: list[pyg.data.Data], name: str, base_path: str
) -> None:
"""
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.
"""
torch.save(
[graph.edge_index for graph in graphs],
os.path.join(base_path, f"{name}_edge_index.pt"),
)
edge_features = [
torch.cat((graph.len.unsqueeze(1), graph.vdiff), dim=1).to(
torch.float32
)
for graph in graphs
] # Save as float32
torch.save(edge_features, os.path.join(base_path, f"{name}_features.pt"))
[docs]
def from_networkx_with_start_index(
nx_graph: networkx.Graph, start_index: int
) -> pyg.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.
Parameters
----------
nx_graph : networkx.Graph
Graph to convert.
start_index : int
Offset added to both rows of the PyG edge-index tensor.
Returns
-------
pyg.data.Data
Converted PyG graph with shifted edge indices.
"""
pyg_graph = from_networkx(nx_graph)
pyg_graph.edge_index += start_index
return pyg_graph
[docs]
def 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.
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.
Returns
-------
None
"""
assert edge_index.ndim == 2, (
f"{name}: expected edge_index shape [2, E], "
f"got {tuple(edge_index.shape)}"
)
assert edge_index.shape[0] == 2, (
f"{name}: expected edge_index shape [2, E], "
f"got {tuple(edge_index.shape)}"
)
assert edge_index.shape[1] > 0, f"{name}: expected at least one edge"
sender_start, sender_stop = sender_range
receiver_start, receiver_stop = receiver_range
sender_min = int(edge_index[0].min().item())
sender_max = int(edge_index[0].max().item())
receiver_min = int(edge_index[1].min().item())
receiver_max = int(edge_index[1].max().item())
assert sender_min >= sender_start, (
f"{name}: sender min {sender_min} is below expected range "
f"[{sender_start}, {sender_stop})"
)
assert sender_max < sender_stop, (
f"{name}: sender max {sender_max} is above expected range "
f"[{sender_start}, {sender_stop})"
)
assert receiver_min >= receiver_start, (
f"{name}: receiver min {receiver_min} is below expected range "
f"[{receiver_start}, {receiver_stop})"
)
assert receiver_max < receiver_stop, (
f"{name}: receiver max {receiver_max} is above expected range "
f"[{receiver_start}, {receiver_stop})"
)
[docs]
def 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.
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
-------
torch.Tensor
Edge index tensor with row-wise offsets removed.
"""
return torch.stack(
(edge_index[0] - sender_offset, edge_index[1] - receiver_offset),
dim=0,
)
[docs]
def mk_2d_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph:
"""
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
-------
networkx.DiGraph
The constructed directed 2-D grid graph.
"""
xm, xM = np.amin(xy[:, :, 0][:, 0]), np.amax(xy[:, :, 0][:, 0])
ym, yM = np.amin(xy[:, :, 1][0, :]), np.amax(xy[:, :, 1][0, :])
# avoid nodes on border
dx = (xM - xm) / nx
dy = (yM - ym) / ny
lx = np.linspace(xm + dx / 2, xM - dx / 2, nx)
ly = np.linspace(ym + dy / 2, yM - dy / 2, ny)
mg = np.meshgrid(lx, ly, indexing="ij") # Use 'ij' indexing for (Nx,Ny)
g = networkx.grid_2d_graph(len(lx), len(ly))
for node in g.nodes:
g.nodes[node]["pos"] = np.array([mg[0][node], mg[1][node]])
# add diagonal edges
g.add_edges_from(
[((x, y), (x + 1, y + 1)) for y in range(ny - 1) for x in range(nx - 1)]
+ [
((x + 1, y), (x, y + 1))
for y in range(ny - 1)
for x in range(nx - 1)
]
)
# turn into directed graph
dg = networkx.DiGraph(g)
for u, v in g.edges():
d = np.sqrt(np.sum((g.nodes[u]["pos"] - g.nodes[v]["pos"]) ** 2))
dg.edges[u, v]["len"] = d
dg.edges[u, v]["vdiff"] = g.nodes[u]["pos"] - g.nodes[v]["pos"]
dg.add_edge(v, u)
dg.edges[v, u]["len"] = d
dg.edges[v, u]["vdiff"] = g.nodes[v]["pos"] - g.nodes[u]["pos"]
return dg
[docs]
def prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph:
"""
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
-------
networkx.Graph
The relabeled graph.
"""
# Relabel node indices in graph, insert (graph_level, i, j)
ijk = [tuple((new_index,) + x) for x in graph.nodes]
to_mapping = dict(zip(graph.nodes, ijk))
return networkx.relabel_nodes(graph, to_mapping, copy=True)
[docs]
def create_graph(
graph_dir_path: str,
xy: np.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)
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).
Returns
-------
None
"""
os.makedirs(graph_dir_path, exist_ok=True)
logger.info(f"Writing graph components to {graph_dir_path}")
#
# Mesh
#
# graph geometry
nx = 3 # number of children =nx**2
nlev = int(np.log(max(xy.shape[:2])) / np.log(nx))
nleaf = nx**nlev # leaves at the bottom = nleaf**2
mesh_levels = nlev - 1
if n_max_levels:
# Limit the levels in mesh graph
mesh_levels = min(mesh_levels, n_max_levels)
# print(f"nlev: {nlev}, nleaf: {nleaf}, mesh_levels: {mesh_levels}")
# multi resolution tree levels
G = []
for lev in range(1, mesh_levels + 1):
n = int(nleaf / (nx**lev))
g = mk_2d_graph(xy, n, n)
if create_plot:
plot_graph(from_networkx(g), title=f"Mesh graph, level {lev}")
plt.show()
G.append(g)
if hierarchical:
# Relabel nodes of each level with level index first
G = [
prepend_node_index(graph, level_i)
for level_i, graph in enumerate(G)
]
num_nodes_level = np.array([len(g_level.nodes) for g_level in G])
# First node index in each level in the hierarchical graph
first_index_level = np.concatenate(
(np.zeros(1, dtype=int), np.cumsum(num_nodes_level[:-1]))
)
# Create inter-level mesh edges from lower levels to upper levels.
up_graphs = []
for upper_level_index, lower_level_index, G_upper, G_lower in zip(
range(1, mesh_levels), range(mesh_levels - 1), G[1:], G[:-1]
):
G_up = G_lower.copy()
G_up.clear_edges()
G_up = networkx.DiGraph(G_up)
G_up.add_nodes_from(G_upper.nodes(data=True))
upper_node_list = list(G_upper.nodes)
lower_node_list = list(G_lower.nodes)
upper_node_xy = np.array(
[xy for _, xy in G_upper.nodes.data("pos")]
)
kdt_m = scipy.spatial.KDTree(upper_node_xy)
for lower_node in lower_node_list:
# find 1(?) nearest neighbours (index to vm_xy)
neigh_idx = kdt_m.query(G_up.nodes[lower_node]["pos"], 1)[1]
upper_node = upper_node_list[neigh_idx]
G_up.add_edge(lower_node, upper_node)
d = np.sqrt(
np.sum(
(
G_up.nodes[lower_node]["pos"]
- G_up.nodes[upper_node]["pos"]
)
** 2
)
)
G_up.edges[lower_node, upper_node]["len"] = d
G_up.edges[lower_node, upper_node]["vdiff"] = (
G_up.nodes[lower_node]["pos"]
- G_up.nodes[upper_node]["pos"]
)
lower_level_node_index_offset = int(
first_index_level[lower_level_index]
)
upper_level_node_index_offset = int(
first_index_level[upper_level_index]
)
lower_level_num_nodes = int(num_nodes_level[lower_level_index])
upper_level_num_nodes = int(num_nodes_level[upper_level_index])
lower_level_node_index_range = (
lower_level_node_index_offset,
lower_level_node_index_offset + lower_level_num_nodes,
)
upper_level_node_index_range = (
upper_level_node_index_offset,
upper_level_node_index_offset + upper_level_num_nodes,
)
G_up_int = networkx.convert_node_labels_to_integers(
G_up,
first_label=lower_level_node_index_offset,
ordering="sorted",
)
G_up_int = sort_nodes_internally(G_up_int)
pyg_up = from_networkx_with_start_index(
G_up_int, lower_level_node_index_offset
)
assert_edge_index_in_range(
pyg_up.edge_index,
sender_range=lower_level_node_index_range,
receiver_range=upper_level_node_index_range,
name=f"mesh_up_edge_index[{lower_level_index}] before offset",
)
pyg_up.edge_index = zero_offset_edge_index(
pyg_up.edge_index,
sender_offset=lower_level_node_index_offset,
receiver_offset=upper_level_node_index_offset,
)
assert_edge_index_in_range(
pyg_up.edge_index,
sender_range=(0, lower_level_num_nodes),
receiver_range=(0, upper_level_num_nodes),
name=f"mesh_up_edge_index[{lower_level_index}] after offset",
)
up_graphs.append(pyg_up)
if create_plot:
plot_graph(
pyg_up,
title=(
f"Up graph, {lower_level_index} -> {upper_level_index}"
),
)
plt.show()
down_graphs = []
for lower_level_index, pyg_up in enumerate(up_graphs):
upper_level_index = lower_level_index + 1
lower_level_num_nodes = int(num_nodes_level[lower_level_index])
upper_level_num_nodes = int(num_nodes_level[upper_level_index])
pyg_down = pyg_up.clone()
pyg_down.edge_index = torch.stack(
(pyg_up.edge_index[1], pyg_up.edge_index[0]), dim=0
)
pyg_down.vdiff = -pyg_up.vdiff
assert_edge_index_in_range(
pyg_down.edge_index,
sender_range=(0, upper_level_num_nodes),
receiver_range=(0, lower_level_num_nodes),
name=f"mesh_down_edge_index[{lower_level_index}] after offset",
)
down_graphs.append(pyg_down)
if create_plot:
plot_graph(
pyg_down,
title=(
f"Down graph, {upper_level_index} -> {lower_level_index}" # noqa: E501
),
)
plt.show()
# Save up and down edges
save_edges_list(up_graphs, "mesh_up", graph_dir_path)
save_edges_list(down_graphs, "mesh_down", graph_dir_path)
# Extract intra-level edges for m2m
m2m_graphs = [
from_networkx_with_start_index(
networkx.convert_node_labels_to_integers(
level_graph, first_label=start_index, ordering="sorted"
),
start_index,
)
for level_graph, start_index in zip(G, first_index_level)
]
for level_index, (graph, start_index, num_nodes) in enumerate(
zip(m2m_graphs, first_index_level, num_nodes_level)
):
start_index = int(start_index)
num_nodes = int(num_nodes)
combined_range = (start_index, start_index + num_nodes)
assert_edge_index_in_range(
graph.edge_index,
sender_range=combined_range,
receiver_range=combined_range,
name=f"m2m_edge_index[{level_index}] before offset",
)
graph.edge_index = zero_offset_edge_index(
graph.edge_index,
sender_offset=start_index,
receiver_offset=start_index,
)
assert_edge_index_in_range(
graph.edge_index,
sender_range=(0, num_nodes),
receiver_range=(0, num_nodes),
name=f"m2m_edge_index[{level_index}] after offset",
)
mesh_pos = [graph.pos.to(torch.float32) for graph in m2m_graphs]
# For use in g2m and m2g
G_bottom_mesh = G[0]
joint_mesh_graph = networkx.union_all([graph for graph in G])
all_mesh_nodes = joint_mesh_graph.nodes(data=True)
else:
# combine all levels to one graph
G_tot = G[0]
for lev in range(1, len(G)):
nodes = list(G[lev - 1].nodes)
n = int(np.sqrt(len(nodes)))
ij = (
np.array(nodes)
.reshape((n, n, 2))[1::nx, 1::nx, :]
.reshape(int(n / nx) ** 2, 2)
)
ij_tuples = [tuple(x) for x in ij]
G[lev] = networkx.relabel_nodes(
G[lev], dict(zip(G[lev].nodes, ij_tuples))
)
G_tot = networkx.compose(G_tot, G[lev])
# Relabel mesh nodes to start with 0
G_tot = prepend_node_index(G_tot, 0)
num_nodes_level = np.array([len(G_tot.nodes)])
first_index_level = np.zeros(1, dtype=int)
# relabel nodes to integers (sorted)
G_int = networkx.convert_node_labels_to_integers(
G_tot, first_label=0, ordering="sorted"
)
# Graph to use in g2m and m2g
G_bottom_mesh = G_tot
all_mesh_nodes = G_tot.nodes(data=True)
# export the nx graph to PyTorch geometric
pyg_m2m = from_networkx(G_int)
m2m_graphs = [pyg_m2m]
mesh_pos = [pyg_m2m.pos.to(torch.float32)]
if create_plot:
plot_graph(pyg_m2m, title="Mesh-to-mesh")
plt.show()
# Save m2m edges
save_edges_list(m2m_graphs, "m2m", graph_dir_path)
# Save mesh positions
torch.save(
mesh_pos, os.path.join(graph_dir_path, "mesh_features.pt")
) # mesh pos, in float32
#
# Grid2Mesh
#
# radius within which grid nodes are associated with a mesh node
# (in terms of mesh distance)
DM_SCALE = 0.67
# mesh nodes on lowest level
vm = G_bottom_mesh.nodes
vm_xy = np.array([xy for _, xy in vm.data("pos")])
# distance between mesh nodes
dm = np.sqrt(
np.sum((vm.data("pos")[(0, 1, 0)] - vm.data("pos")[(0, 0, 0)]) ** 2)
)
# grid nodes
Nx, Ny = xy.shape[:2]
G_grid = networkx.grid_2d_graph(Nx, Ny)
G_grid.clear_edges()
# vg features (only pos introduced here)
for node in G_grid.nodes:
# pos is in feature but here explicit for convenience
G_grid.nodes[node]["pos"] = xy[
node[0], node[1]
] # xy is already (Nx,Ny,2)
# add 1000 to node key to separate grid nodes (1000,i,j) from mesh nodes
# (i,j) and impose sorting order such that vm are the first nodes
G_grid = prepend_node_index(G_grid, 1000)
# build kd tree for grid point pos
# order in vg_list should be same as in vg_xy
vg_list = list(G_grid.nodes)
vg_xy = np.array(
[xy[node[1], node[2]] for node in vg_list]
) # xy is already (Nx,Ny,2)
kdt_g = scipy.spatial.KDTree(vg_xy)
# now add (all) mesh nodes, include features (pos)
G_grid.add_nodes_from(all_mesh_nodes)
# Re-create graph with sorted node indices
# Need to do sorting of nodes this way for indices to map correctly to pyg
G_g2m = networkx.Graph()
G_g2m.add_nodes_from(sorted(G_grid.nodes(data=True)))
# turn into directed graph
G_g2m = networkx.DiGraph(G_g2m)
# add edges
for v in vm:
# find neighbours (index to vg_xy)
neigh_idxs = kdt_g.query_ball_point(vm[v]["pos"], dm * DM_SCALE)
for i in neigh_idxs:
u = vg_list[i]
# add edge from grid to mesh
G_g2m.add_edge(u, v)
d = np.sqrt(
np.sum((G_g2m.nodes[u]["pos"] - G_g2m.nodes[v]["pos"]) ** 2)
)
G_g2m.edges[u, v]["len"] = d
G_g2m.edges[u, v]["vdiff"] = (
G_g2m.nodes[u]["pos"] - G_g2m.nodes[v]["pos"]
)
pyg_g2m = from_networkx(G_g2m)
if create_plot:
plot_graph(pyg_g2m, title="Grid-to-mesh")
plt.show()
#
# Mesh2Grid
#
# start out from Grid2Mesh and then replace edges
G_m2g = G_g2m.copy()
G_m2g.clear_edges()
# build kd tree for mesh point pos
# order in vm should be same as in vm_xy
vm_list = list(vm)
kdt_m = scipy.spatial.KDTree(vm_xy)
# add edges from mesh to grid
for v in vg_list:
# find 4 nearest neighbours (index to vm_xy)
neigh_idxs = kdt_m.query(G_m2g.nodes[v]["pos"], 4)[1]
for i in neigh_idxs:
u = vm_list[i]
# add edge from mesh to grid
G_m2g.add_edge(u, v)
d = np.sqrt(
np.sum((G_m2g.nodes[u]["pos"] - G_m2g.nodes[v]["pos"]) ** 2)
)
G_m2g.edges[u, v]["len"] = d
G_m2g.edges[u, v]["vdiff"] = (
G_m2g.nodes[u]["pos"] - G_m2g.nodes[v]["pos"]
)
# relabel nodes to integers (sorted)
G_m2g_int = networkx.convert_node_labels_to_integers(
G_m2g, first_label=0, ordering="sorted"
)
pyg_m2g = from_networkx(G_m2g_int)
if create_plot:
plot_graph(pyg_m2g, title="Mesh-to-grid")
plt.show()
# Save g2m and m2g everything
num_mesh_nodes_total = int(np.sum(num_nodes_level))
num_bottom_mesh_nodes = int(num_nodes_level[0])
num_grid_nodes = Nx * Ny
assert_edge_index_in_range(
pyg_g2m.edge_index,
sender_range=(
num_mesh_nodes_total,
num_mesh_nodes_total + num_grid_nodes,
),
receiver_range=(0, num_bottom_mesh_nodes),
name="g2m_edge_index before offset",
)
pyg_g2m.edge_index = zero_offset_edge_index(
pyg_g2m.edge_index,
sender_offset=num_mesh_nodes_total,
receiver_offset=0,
)
assert_edge_index_in_range(
pyg_g2m.edge_index,
sender_range=(0, num_grid_nodes),
receiver_range=(0, num_bottom_mesh_nodes),
name="g2m_edge_index after offset",
)
assert_edge_index_in_range(
pyg_m2g.edge_index,
sender_range=(0, num_bottom_mesh_nodes),
receiver_range=(
num_mesh_nodes_total,
num_mesh_nodes_total + num_grid_nodes,
),
name="m2g_edge_index before offset",
)
pyg_m2g.edge_index = zero_offset_edge_index(
pyg_m2g.edge_index,
sender_offset=0,
receiver_offset=num_mesh_nodes_total,
)
assert_edge_index_in_range(
pyg_m2g.edge_index,
sender_range=(0, num_bottom_mesh_nodes),
receiver_range=(0, num_grid_nodes),
name="m2g_edge_index after offset",
)
# g2m
save_edges(pyg_g2m, "g2m", graph_dir_path)
# m2g
save_edges(pyg_m2g, "m2g", graph_dir_path)
with open(
os.path.join(graph_dir_path, METAINFO_FILENAME),
"w",
encoding="utf-8",
) as fp:
yaml.dump({"spec_version": CURRENT_GRAPH_SPEC_VERSION}, fp)
[docs]
def create_graph_from_datastore(
datastore: 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.
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.
"""
if not datastore.is_on_rectilinear_grid:
raise NotImplementedError(
"Graph creation requires a datastore on a rectilinear grid."
)
xy = datastore.get_xy(category="state", stacked=False)
create_graph(
graph_dir_path=output_root_path,
xy=xy,
n_max_levels=n_max_levels,
hierarchical=hierarchical,
create_plot=create_plot,
)
[docs]
def cli(input_args: list[str] | None = None) -> None:
"""
Parse CLI arguments and call :func:`create_graph_from_datastore`.
Parameters
----------
input_args : list[str] or None, optional
Argument list forwarded to :class:`argparse.ArgumentParser`. When
``None``, ``sys.argv`` is used.
"""
warnings.warn(
"create_graph.py is deprecated and will be removed in a future "
"version. Use create_graph_with_wmg.py instead, which delegates "
"graph creation to weather-model-graphs (wmg). See "
"https://github.com/mllam/neural-lam/issues/384 for details.",
DeprecationWarning,
stacklevel=2,
)
parser = ArgumentParser(
description="Graph generation for neural-lam",
formatter_class=ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--config_path",
type=str,
help="Path to neural-lam configuration file",
)
parser.add_argument(
"--name",
type=str,
default="multiscale",
help="Name to save graph as",
)
parser.add_argument(
"--plot",
action="store_true",
help="If graphs should be plotted during generation",
)
parser.add_argument(
"--levels",
type=int,
help="Limit multi-scale mesh to given number of levels, "
"from bottom up. If None, no limit.",
)
parser.add_argument(
"--hierarchical",
action="store_true",
help="Generate hierarchical mesh graph. Otherwise multi-scale.",
)
args = parser.parse_args(input_args)
if args.config_path is None:
raise ValueError("Specify your config with --config_path")
# Load neural-lam configuration and datastore to use
_, datastore, _ = load_config_and_datastore(config_path=args.config_path)
create_graph_from_datastore(
datastore=datastore,
output_root_path=os.path.join(datastore.root_path, "graph", args.name),
n_max_levels=args.levels,
hierarchical=args.hierarchical,
create_plot=args.plot,
)
if __name__ == "__main__":
cli()
