Vectorize KDTree queries and edge insertion in connect_nodes_across_graphs

[Raj-Taware]

Description

---

Building on the benchmarks and flamegraph profiling introduced in #117, this issue targets the a bottleneck identified in connect_nodes_across_graphs :- the repeated per-node KDTree queries and edge insertions that could be replaced with vectorized, batched operations.
The current implementation in base.py has three compounding inefficiencies inside a Python loop over all N_target nodes:

• Serial KDTree queries : _find_neighbour_node_idxs_in_source_mesh(xy_target) is called once per target node, passing a single 2D point. All three underlying methods (kdt_s.query() for nearest_neighbour/nearest_neighbours, kdt_s.query_ball_point() for within_radius) natively accept an array of points so N_target individual Python-level calls could become one.
• Per-edge add_edge calls : each edge is inserted individually into the DiGraph. NetworkX supports add_edges_from() which takes a list, cutting down on per-call overhead.
• Per-edge attribute computation : len (Euclidean distance) and vdiff (position difference) are computed as scalar numpy ops inside the inner loop, when they could be derived in a single vectorized pass over all edges once neighbour indices are known.

Also worth noting: containing_rectangle calls connect_nodes_across_graphs recursively with within_radius any changes here need to keep that path intact.

The flamegraph from #117 makes this visible. We can see connect_nodes_across_graphs dominates the call stack:

Proposed Enhancements

---

High-level direction:

• Stack all target node positions into a single array, run one batched KDTree query, build the full edge list from the result
• nearest_neighbour / nearest_neighbours are the most straightforward, batch query returns a regular (N_target, k) array
• within_radius needs a bit more thought since query_ball_point on an array returns a ragged list (variable neighbours per node), so edge-list construction is less clean. Open to ideas on the best approach here.

Happy to discuss before going to implementation.

---

Context & Links

---

• Benchmarks and flamegraph: #117
• Implementation: src/weather_model_graphs/create/base.py
• Dev meeting 2026-04-13: discussed the need for better hardware utilization as grid sizes scale

---

Task List

---

• Batch KDTree queries across all target nodes for nearest_neighbour and nearest_neighbours
• Handle within_radius & query_ball_point ragged output
• Replace per-edge add_edge loop with a single add_edges_from call
• Vectorize len and vdiff edge attribute computation
• Verify containing_rectangle (recursive call) still works correctly
• Run scaling benchmark from [Feat] runtime scaling and flamegraph benchmark CLIs #117 before/after to measure improvement

[joeloskarsson]

This sounds great I think, this is indeed a part of the codebase that is unnecessarily serial and can be improved a lot.

With regards to handling the ragged lists from query_ball_point: I was first thinking that there might be a way to instead use just query, to get a fixed number of neighbors, and then filter these based on their distances. The problem with that is that we don't know the maximum number of neighbors to get a priori. So I think that you might be better off just sticking to query_ball_point and just handling the ragged list. I would try to flatten this ragged list as quickly as possible. Should be able to make an identical list (same ragged shape) with the query nodes for each row, and then you have both the sender and receiver nodes of each edge in the same position. If you flatten these in the same way you just get two 1d lists of nodes that describe the edge connections.

[yuvraajnarula]

@Raj-Taware
The serial KDTree queries and per‑edge overhead have been on my radar too, especially after looking at the flamegraph in #117.

I like the direction you've outlined. A couple of thoughts:

• Flattening the ragged list: Joel's suggestion to flatten sender/receiver arrays in parallel makes sense. We can also use np.repeat for the source indices and np.concatenate for the neighbours to keep it vectorised.
• within_radius: If the ragged output becomes too messy, we could cap the number of neighbours (e.g., k=50) and filter by distance later. That would give a fixed‑shape array and simplify batching, at the cost of some extra distance checks. Might be worth benchmarking both approaches.

I'm currently working on a memory‑profiling extension for the scaling benchmarks (#137). Once that's ready, we could use it to measure the memory impact of your changes as well.

Happy to help review or test the vectorised version on the icosahedral graphs from #76. Let me know if you'd like me to run any preliminary benchmarks on my end.

[Raj-Taware]

@joeloskarsson Agreed, flattening the ragged list seems to be the best option right now.
We will query_ball_point on the full array, then build sender/receiver in parallel with np.repeat and np.concatenate:

counts = [len(row) for row in neigh_idxs]
target_indices = np.repeat(np.arange(N_target), counts)
source_indices = np.concatenate(neigh_idxs)

Since we have 2 flat 1D arrays with the same shape and one edge per entry, it makes the query + distance-filter fallback unnecessary, but I will still benchmark both just to double check.

---

@yuvraajnarula agreed on np.repeat / np.concatenate usage. Once the index arrays are built, lengths and vdiff both come out of a single pass:

vdiff = xy_source[source_indices] - xy_target_all[target_indices]
lengths = np.linalg.norm(vdiff, axis=1)

Would love to pull in the memory profiling from #137 when it's ready, having memory impact alongside runtime will make the before/after comparison a lot more complete. And definitely keen to test on the icosahedral graphs from #76, I'll ping you when the PR is up.

Will include scaling benchmarks from #117 before/after in the PR.