perf: Vectorize KDTree queries and edge insertion in connect_nodes_across_graphs

src/weather_model_graphs/create/base.py

@@ -255,16 +255,18 @@ 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)
 
-    # Determine method and perform checks once
-    # Conditionally define _find_neighbour_node_idxs_in_source_mesh for use in
-    # loop later
+    xy_target_all = np.array([G_target.nodes[n]["pos"] for n in target_nodes_list])
+    N_target = len(target_nodes_list)
+
+    # 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
@@ -323,10 +325,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:
@@ -337,10 +338,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:
@@ -389,41 +390,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)))
 
-    # 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"]
-        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

tests/test_connect_nodes_vectorized.py

@@ -0,0 +1,102 @@
+# 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
+
+
+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", [1, 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
+    )
+
+    # 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
+        assert len(preds) == k, (
+            f"Target {node} has {len(preds)} predecessors, expected exactly {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
+    )
+    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, (
+            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_edge_attrs_valid(G, source, target)