From 9bca7f5268e66d3250e99ddfa7ee8533db0624d4 Mon Sep 17 00:00:00 2001 From: Raj_taware Date: Thu, 23 Apr 2026 02:30:32 +0530 Subject: [PATCH 1/5] test: add property-based correctness tests for connect_nodes_across_graphs --- tests/test_connect_nodes_vectorized.py | 97 ++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 tests/test_connect_nodes_vectorized.py diff --git a/tests/test_connect_nodes_vectorized.py b/tests/test_connect_nodes_vectorized.py new file mode 100644 index 0000000..05d05db --- /dev/null +++ b/tests/test_connect_nodes_vectorized.py @@ -0,0 +1,97 @@ +# tests/test_connect_nodes_vectorized.py +import numpy as np +import pytest +import weather_model_graphs as wmg +from weather_model_graphs.create.base import connect_nodes_across_graphs +import tests.utils as test_utils + + +def _make_source_target(): + """Small deterministic mesh source + grid target for fast tests.""" + xy = test_utils.create_fake_xy(N=10) + source = wmg.create.mesh.create_single_level_2d_mesh_graph(xy=xy, nx=5, ny=5) + target = wmg.create.grid.create_grid_graph_nodes(xy) + return source, target + + +def _assert_edge_attrs_valid(G, G_source, G_target): + """Common invariants for all connection methods.""" + assert set(G.nodes) == set(G_source.nodes) | set(G_target.nodes) + + for u, v, data in G.edges(data=True): + assert u in G_source.nodes, f"Edge source {u} not in G_source" + assert v in G_target.nodes, f"Edge target {v} not in G_target" + assert "len" in data, "Missing 'len' attribute" + assert "vdiff" in data, "Missing 'vdiff' attribute" + + pos_u = np.array(G.nodes[u]["pos"]) + pos_v = np.array(G.nodes[v]["pos"]) + expected_vdiff = pos_u - pos_v + expected_len = np.linalg.norm(expected_vdiff) + + assert np.isclose(data["len"], expected_len, atol=1e-10), ( + f"Edge ({u},{v}): len={data['len']}, expected={expected_len}" + ) + assert np.allclose(data["vdiff"], expected_vdiff, atol=1e-10), ( + f"Edge ({u},{v}): vdiff={data['vdiff']}, expected={expected_vdiff}" + ) + + +def test_nearest_neighbour_one_edge_per_target(): + source, target = _make_source_target() + G = connect_nodes_across_graphs(source, target, method="nearest_neighbour") + + for node in target.nodes: + preds = list(G.predecessors(node)) + assert len(preds) == 1, f"Target {node} has {len(preds)} predecessors, expected 1" + assert preds[0] in source.nodes + + _assert_edge_attrs_valid(G, source, target) + + +@pytest.mark.parametrize("k", [3, 4, 8]) +def test_nearest_neighbours_at_most_k_per_target(k): + source, target = _make_source_target() + G = connect_nodes_across_graphs( + source, target, method="nearest_neighbours", max_num_neighbours=k + ) + + for node in target.nodes: + preds = list(G.predecessors(node)) + assert len(preds) <= k, ( + f"Target {node} has {len(preds)} predecessors, expected <= {k}" + ) + for p in preds: + assert p in source.nodes + + _assert_edge_attrs_valid(G, source, target) + + +@pytest.mark.parametrize("max_dist", [1.5, 3.0, 6.0]) +def test_within_radius_all_edges_within_dist(max_dist): + source, target = _make_source_target() + G = connect_nodes_across_graphs( + source, target, method="within_radius", max_dist=max_dist + ) + + for u, v, data in G.edges(data=True): + assert data["len"] <= max_dist + 1e-10, ( + f"Edge ({u},{v}) len={data['len']} exceeds max_dist={max_dist}" + ) + + _assert_edge_attrs_valid(G, source, target) + + +def test_within_radius_rel_max_dist(): + source, target = _make_source_target() + G = connect_nodes_across_graphs( + source, target, method="within_radius", rel_max_dist=1.0 + ) + _assert_edge_attrs_valid(G, source, target) + + +def test_containing_rectangle_nodes_preserved(): + source, target = _make_source_target() + G = connect_nodes_across_graphs(source, target, method="containing_rectangle") + assert set(G.nodes) == set(source.nodes) | set(target.nodes) + _assert_edge_attrs_valid(G, source, target) From 203ed3b984a6469097736bcbf423683bb1beed2e Mon Sep 17 00:00:00 2001 From: Raj_taware Date: Thu, 23 Apr 2026 02:40:20 +0530 Subject: [PATCH 2/5] Fix code quality issues in test_connect_nodes_vectorized.py - Fix 1: Add xfail marker for k=1 in nearest_neighbours parametrize (line 53-58) k=1 triggers scalar-vs-array bug in serial loop; fixed by vectorization - Fix 2: Add non-empty guard in within_radius test (line 83) Assert that edges exist before testing their properties - Fix 3: Change '<= k' to '== k' in nearest_neighbours test (line 68) Source has 25 nodes >= max k=8, so every target always gets exactly k neighbours - Fix 4: Remove duplicate node-set check in containing_rectangle test (line 101-104) _assert_edge_attrs_valid already performs this check - Fix 5: Fix import ordering (line 1-7) Add blank line between third-party and first-party imports Sort imports correctly per isort convention --- tests/test_connect_nodes_vectorized.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/test_connect_nodes_vectorized.py b/tests/test_connect_nodes_vectorized.py index 05d05db..9e2263a 100644 --- a/tests/test_connect_nodes_vectorized.py +++ b/tests/test_connect_nodes_vectorized.py @@ -1,9 +1,10 @@ # tests/test_connect_nodes_vectorized.py import numpy as np import pytest + +import tests.utils as test_utils import weather_model_graphs as wmg from weather_model_graphs.create.base import connect_nodes_across_graphs -import tests.utils as test_utils def _make_source_target(): @@ -49,7 +50,12 @@ def test_nearest_neighbour_one_edge_per_target(): _assert_edge_attrs_valid(G, source, target) -@pytest.mark.parametrize("k", [3, 4, 8]) +@pytest.mark.parametrize("k", [ + pytest.param(1, marks=pytest.mark.xfail( + strict=True, reason="k=1 triggers scalar-vs-array bug in serial loop; fixed by vectorization" + )), + 3, 4, 8, +]) def test_nearest_neighbours_at_most_k_per_target(k): source, target = _make_source_target() G = connect_nodes_across_graphs( @@ -58,8 +64,9 @@ def test_nearest_neighbours_at_most_k_per_target(k): for node in target.nodes: preds = list(G.predecessors(node)) - assert len(preds) <= k, ( - f"Target {node} has {len(preds)} predecessors, expected <= {k}" + # source has 25 nodes >= max k=8, so every target always gets exactly k neighbours + assert len(preds) == k, ( + f"Target {node} has {len(preds)} predecessors, expected exactly {k}" ) for p in preds: assert p in source.nodes @@ -73,6 +80,7 @@ def test_within_radius_all_edges_within_dist(max_dist): G = connect_nodes_across_graphs( source, target, method="within_radius", max_dist=max_dist ) + assert G.number_of_edges() > 0, f"Expected edges with max_dist={max_dist}, got 0" for u, v, data in G.edges(data=True): assert data["len"] <= max_dist + 1e-10, ( @@ -93,5 +101,4 @@ def test_within_radius_rel_max_dist(): def test_containing_rectangle_nodes_preserved(): source, target = _make_source_target() G = connect_nodes_across_graphs(source, target, method="containing_rectangle") - assert set(G.nodes) == set(source.nodes) | set(target.nodes) _assert_edge_attrs_valid(G, source, target) From 46da5a9baf493dd8584bb9465cc4c90905373bee Mon Sep 17 00:00:00 2001 From: Raj_taware Date: Thu, 23 Apr 2026 02:42:59 +0530 Subject: [PATCH 3/5] refactor: align xy_source ordering with source_nodes_list in connect_nodes_across_graphs --- src/weather_model_graphs/create/base.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index fbfa93b..fc0b6ea 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -255,13 +255,16 @@ def connect_nodes_across_graphs( Graph containing the nodes in `G_source` and `G_target` and directed edges from nodes in `G_source` to nodes in `G_target` """ - source_nodes_list = list(G_source.nodes) + source_nodes_list = sorted(G_source.nodes) target_nodes_list = list(G_target.nodes) - # build kd tree for source nodes (e.g. the mesh nodes when constructing m2g) - xy_source = np.array([G_source.nodes[node]["pos"] for node in G_source.nodes]) + # build kd tree for source nodes; order must match source_nodes_list + xy_source = np.array([G_source.nodes[n]["pos"] for n in source_nodes_list]) kdt_s = scipy.spatial.KDTree(xy_source) + xy_target_all = np.array([G_target.nodes[n]["pos"] for n in target_nodes_list]) + N_target = len(target_nodes_list) + # Determine method and perform checks once # Conditionally define _find_neighbour_node_idxs_in_source_mesh for use in # loop later @@ -400,9 +403,6 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target): G_connect.add_nodes_from(sorted(G_source.nodes(data=True))) G_connect.add_nodes_from(sorted(G_target.nodes(data=True))) - # sort nodes by index - source_nodes_list = sorted(G_source.nodes) - # add edges for target_node in target_nodes_list: xy_target = G_target.nodes[target_node]["pos"] From 2b8653ff788062dc124e09557a1ded551a39bf77 Mon Sep 17 00:00:00 2001 From: Raj_taware Date: Thu, 23 Apr 2026 02:51:51 +0530 Subject: [PATCH 4/5] perf: vectorize KDTree queries and edge insertion in connect_nodes_across_graphs - Batch all target positions and run one KDTree query per method - Flatten nearest_neighbours result with ravel+repeat - Flatten within_radius ragged output with np.repeat/np.concatenate - Compute vdiff and len in one numpy pass over all edges - Replace per-edge add_edge loop with add_edges_from - Remove now-resolved xfail on nearest_neighbours k=1 (scalar-vs-array bug no longer exists in the vectorized path) --- src/weather_model_graphs/create/base.py | 64 ++++++++++++------------- tests/test_connect_nodes_vectorized.py | 7 +-- 2 files changed, 32 insertions(+), 39 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index fc0b6ea..d0d7191 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -326,10 +326,9 @@ def _edge_filter(edge_prop): raise Exception( "to use `nearest_neighbour` you should not set `max_dist`, `rel_max_dist`or `max_num_neighbours`" ) - - def _find_neighbour_node_idxs_in_source_mesh(xy_target): - neigh_idx = kdt_s.query(xy_target, 1)[1] - return [neigh_idx] + # Batch query: shape (N_target,) + source_indices = kdt_s.query(xy_target_all, 1)[1] + target_indices = np.arange(N_target) elif method == "nearest_neighbours": if max_num_neighbours is None: @@ -340,10 +339,10 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target): raise Exception( "to use `nearest_neighbours` you should not set `max_dist` or `rel_max_dist`" ) - - def _find_neighbour_node_idxs_in_source_mesh(xy_target): - neigh_idxs = kdt_s.query(xy_target, max_num_neighbours)[1] - return neigh_idxs + # Batch query: shape (N_target, max_num_neighbours) + neigh_idxs = kdt_s.query(xy_target_all, max_num_neighbours)[1] + source_indices = neigh_idxs.ravel() + target_indices = np.repeat(np.arange(N_target), max_num_neighbours) elif method == "within_radius": if max_num_neighbours is not None: @@ -392,38 +391,37 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target): "to use `witin_radius` method you shold set `max_dist` or `rel_max_dist" ) - def _find_neighbour_node_idxs_in_source_mesh(xy_target): - neigh_idxs = kdt_s.query_ball_point(xy_target, query_dist) - return neigh_idxs + # Batch query: ragged because each target may have different neighbour count + neigh_idxs_ragged = kdt_s.query_ball_point(xy_target_all, query_dist) + counts = [len(row) for row in neigh_idxs_ragged] + target_indices = np.repeat(np.arange(N_target), counts) + nonempty_rows = [row for row in neigh_idxs_ragged if len(row) > 0] + source_indices = ( + np.concatenate(nonempty_rows).astype(int) + if nonempty_rows + else np.array([], dtype=int) + ) else: raise NotImplementedError(method) + # Vectorized edge attribute computation + xy_src = xy_source[source_indices] # (N_edges, 2) + xy_tgt = xy_target_all[target_indices] # (N_edges, 2) + vdiffs = xy_src - xy_tgt # (N_edges, 2) + lengths = np.linalg.norm(vdiffs, axis=1) # (N_edges,) + G_connect = networkx.DiGraph() G_connect.add_nodes_from(sorted(G_source.nodes(data=True))) G_connect.add_nodes_from(sorted(G_target.nodes(data=True))) - # add edges - for target_node in target_nodes_list: - xy_target = G_target.nodes[target_node]["pos"] - neigh_idxs = _find_neighbour_node_idxs_in_source_mesh(xy_target) - for i in neigh_idxs: - source_node = source_nodes_list[i] - # add edge from source to target - G_connect.add_edge(source_node, target_node) - d = np.sqrt( - np.sum( - ( - G_connect.nodes[source_node]["pos"] - - G_connect.nodes[target_node]["pos"] - ) - ** 2 - ) - ) - G_connect.edges[source_node, target_node]["len"] = d - G_connect.edges[source_node, target_node]["vdiff"] = ( - G_connect.nodes[source_node]["pos"] - - G_connect.nodes[target_node]["pos"] - ) + G_connect.add_edges_from( + ( + source_nodes_list[si], + target_nodes_list[ti], + {"len": float(lengths[e]), "vdiff": vdiffs[e]}, + ) + for e, (si, ti) in enumerate(zip(source_indices, target_indices)) + ) return G_connect diff --git a/tests/test_connect_nodes_vectorized.py b/tests/test_connect_nodes_vectorized.py index 9e2263a..a170ef2 100644 --- a/tests/test_connect_nodes_vectorized.py +++ b/tests/test_connect_nodes_vectorized.py @@ -50,12 +50,7 @@ def test_nearest_neighbour_one_edge_per_target(): _assert_edge_attrs_valid(G, source, target) -@pytest.mark.parametrize("k", [ - pytest.param(1, marks=pytest.mark.xfail( - strict=True, reason="k=1 triggers scalar-vs-array bug in serial loop; fixed by vectorization" - )), - 3, 4, 8, -]) +@pytest.mark.parametrize("k", [1, 3, 4, 8]) def test_nearest_neighbours_at_most_k_per_target(k): source, target = _make_source_target() G = connect_nodes_across_graphs( From 5fd1c94c4a4bf0043135573e37fc8cf68cde7390 Mon Sep 17 00:00:00 2001 From: Raj_taware Date: Thu, 23 Apr 2026 02:59:07 +0530 Subject: [PATCH 5/5] fix: update stale comment and add fixture precondition guard in vectorization --- src/weather_model_graphs/create/base.py | 5 ++--- tests/test_connect_nodes_vectorized.py | 3 +++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index d0d7191..f246379 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -265,9 +265,8 @@ def connect_nodes_across_graphs( xy_target_all = np.array([G_target.nodes[n]["pos"] for n in target_nodes_list]) N_target = len(target_nodes_list) - # Determine method and perform checks once - # Conditionally define _find_neighbour_node_idxs_in_source_mesh for use in - # loop later + # Validate arguments and run batch KDTree query for the chosen method. + # Each branch sets source_indices and target_indices (1-D, aligned). if method == "containing_rectangle": if ( max_dist is not None diff --git a/tests/test_connect_nodes_vectorized.py b/tests/test_connect_nodes_vectorized.py index a170ef2..4158b0f 100644 --- a/tests/test_connect_nodes_vectorized.py +++ b/tests/test_connect_nodes_vectorized.py @@ -57,6 +57,9 @@ def test_nearest_neighbours_at_most_k_per_target(k): source, target, method="nearest_neighbours", max_num_neighbours=k ) + # Precondition: fixture must have enough source nodes for all targets to get exactly k + assert source.number_of_nodes() >= k, "fixture too small; test precondition violated" + for node in target.nodes: preds = list(G.predecessors(node)) # source has 25 nodes >= max k=8, so every target always gets exactly k neighbours