Consolidate and unify cached spatial indexing (KDTree) by bknight1 · Pull Request #182 · underworldcode/underworld3

bknight1 · 2026-05-11T06:59:34Z

This PR refactors how spatial indices (KDTree) are managed across Underworld3 to eliminate redundant construction overhead in solver loops.

Key changes:

Unified Caching: Introduced a consistent _get_kdtree() pattern across Swarm, PICSwarm, and MeshVariable. This ensures that a single C++ index is built per step and reused for all property updates.
Automatic Invalidation: Spatial indices are now automatically invalidated and rebuilt only when necessary (e.g., after particle migration or mesh deformation).
Parallel Routing Fixes:
- Fixed an IndexError in Swarm.migrate() where a 1D rank array was incorrectly indexed as 2D during centroid-based routing.
- Restored default max_its=10 in Swarm.advection migration to ensure particles are correctly routed across multiple MPI partitions in high-velocity flows.

Copilot

Pull request overview

This PR refactors spatial indexing (KDTree) usage across swarms/mesh variables to reduce redundant KDTree construction in hot solver loops, while also adjusting swarm migration routing behavior.

Changes:

Introduces cached _get_kdtree() helpers for Swarm, PICSwarm, and MeshVariable, and updates call sites to reuse cached indices.
Adds a cached mesh-domain centroid KDTree (Mesh._get_domain_kdtree()) and updates Swarm.migrate() to use it.
Adjusts swarm migration/advection behavior (rank assignment indexing change; advection() no longer forces max_its=1).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
`src/underworld3/swarms/pic_swarm.py`	Adds cached swarm KDTree accessor and updates nearest-neighbour map construction to use it.
`src/underworld3/swarm.py`	Adds cached swarm KDTree accessor, switches several interpolations/routing paths to cached trees, and updates migration rank assignment logic.
`src/underworld3/discretisation/discretisation_mesh.py`	Reuses cached variable KDTrees for vertex/DOF mapping and adds a cached centroid KDTree for migration routing.
`src/underworld3/discretisation/discretisation_mesh_variables.py`	Adds cached KDTree for mesh-variable DOF coordinates keyed by mesh version and uses it for RBF interpolation.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

                swarm_rank_array = self.dm.getField("DMSwarm_rank")
                swarm_coord_array = self.dm.getField("DMSwarmPIC_coor").reshape(-1, self.dim)

                if not_my_points.shape[0] > 0:
                    dist, rank = mesh_domain_kdtree.query(
                        swarm_coord_array[not_my_points], k=it + 1, sqr_dists=False
                    )
-                    swarm_rank_array[not_my_points, 0] = rank.reshape(-1, it + 1)[:, it]
+                    swarm_rank_array[not_my_points] = rank.reshape(-1, it + 1)[:, it]



-        tree = uw.kdtree.KDTree(self.X.coords)
+        # Use cached KDTree from coordinate variable
+        tree = self.X._get_kdtree()
        dists, indices = tree.query(self.parent.X.coords, sqr_dists=False)



-        tree = uw.kdtree.KDTree(sub_var.coords)
+        tree = sub_var._get_kdtree()
        dists, indices = tree.query(parent_var.coords, sqr_dists=False)


lmoresi · 2026-05-11T07:08:41Z

Thanks @bknight1 — flagging things to address before this lands, since two related PRs merged just after you opened this one.

What's now on `development`

Fix swarm particle loss across rank boundaries during advection (#175) #177 (merged 07:04 UTC) — same Swarm.advection migration fix as you describe here (drops max_its=1 override, falls back to the default max_its=10). Includes a parallel regression test at tests/parallel/test_0765_swarm_advection_no_loss.py (Ben's reproducer at order=1 + a solid-body rotation at order=2).
Add memprobe diagnostic module for memory-growth tracking (#176) #179 (merged 07:03 UTC) — uw.utilities.memprobe diagnostic module + KDTree live-count counters in ckdtree.pyx (__cinit__/__dealloc__ increment/decrement a module-level _live_instances).

To do

Rebase onto current development. Conflict expected in src/underworld3/swarm.py around the end-of-advection self.migrate(...) call. Resolution: keep your refactor, drop your max_its=10 line (it's already there post-Fix swarm particle loss across rank boundaries during advection (#175) #177), keep your other migration changes.
Keep your Swarm.migrate() IndexError fix (1D rank array indexed as 2D). That's a genuine new fix not in Fix swarm particle loss across rank boundaries during advection (#175) #177; the rebase shouldn't touch it.
Verify against Add memprobe diagnostic module for memory-growth tracking (#176) #179's KDTree counter. Your _get_kdtree() cache will hold trees longer than the old pattern did, so uw.kdtree.live_count() will report a higher steady-state. That's correct — they really are alive — but please confirm:
- pytest tests/test_0780_memprobe.py::test_kdtree_live_count_tracks_construction_destruction still passes
- total_constructed() no longer monotonically climbs on each step in a tight solver loop (that's the point of your cache — should be a flat curve after warm-up)
- Consider adding a one-line assertion to your tests that exercises the cache: after N steps, live_count() is bounded and total_constructed() minus initial is bounded too.
Cross-validation with Fix comprehensive memory leaks in solver setup, interpolation caching, and SubDM syncing #178 + Core memory leak fixes: Cython deallocation & callback hardening #181: once rebased, a quick UW_MEMPROBE=1 run on a tight solver loop should show flat per-step RSS and zero kdtree: total_constructed deltas after warm-up. That's the leak-free signature for both the cache refactor and the upstream fixes working together.

Underworld development team with AI support from Claude Code

Memory diagnostics revealed significant RSS growth caused by repeated instantiation of KDTree objects (52 per step in AdvectionDiffusion). Each instantiation builds a new C++ index, creating memory pressure and fragmentation. Fix: - Implement _get_domain_kdtree() on Mesh for cached centroid routing - Cache KDTree in MeshVariable.rbf_interpolate(), keyed by mesh version - Use cached trees in Swarm.migrate() and evaluation hot-paths Reduces KDTree constructions from 52 to 0 per timestep in stable runs. Tests show stable RSS and PASSED test_0006_memory_leak.py.

Refactored Swarm, PICSwarm, and MeshVariable to use a unified _get_kdtree() pattern. This ensures consistent, cached access to spatial indices across the project, eliminating redundant builds. Key improvements: - Swarm coordinates automatically invalidate the cache on migration. - Mesh domain KDTree is now version-aware, rebuilding on deformation. - Mesh vertex and DOF mapping now utilize cached coordinate trees. - Fixed off-by-one indexing bug in PICSwarm._get_map.

Removed incorrect 2nd dimension index [ , 0] when assigning to 'swarm_rank_array'. This array is 1-dimensional (rank identifiers) across all supported PETSc versions in this context. This fix ensures particles that fall outside the local domain can be successfully routed to their nearest rank-centroid during migration.

…e._get_kdtree Addressing review comments for PR #182: - Rebased on latest development (resolved conflict in Swarm.advection) - Fixed IndexError in Swarm.migrate using a flattened one-line assignment - Exposed _get_kdtree in EnhancedMeshVariable for direct access - Added regression tests for KDTree caching in both MeshVariable and Swarm

bknight1 · 2026-05-12T02:54:49Z

Addressed review comments:

Rebased on latest development branch.
Simplified Swarm.migrate IndexError fix using .reshape().flatten() for cleaner 1D/2D handling.
Exposed _get_kdtree() in MeshVariable (EnhancedMeshVariable) for API consistency.
Added regression tests in test_0780_memprobe.py for both MeshVariable and Swarm KDTree caching.
Verified KDTree live-count and total_constructed metrics.

…e._get_kdtree Addressing review comments for PR #182: - Rebased on latest development (resolved conflict in Swarm.advection) - Fixed IndexError in Swarm.migrate using a flattened one-line assignment - Exposed _get_kdtree in EnhancedMeshVariable for direct access - Added regression tests for KDTree caching in both MeshVariable and Swarm

bknight1 · 2026-05-12T02:59:18Z

Further updates addressing Copilot and user suggestions:

Refactored Swarm.migrate() IndexError fix to use idiomatic .reshape().flatten() for both 1D and 2D results.
Fixed coordinate system mismatch in discretisation_mesh.py (_build_vertex_map and _build_dof_map): switched query coordinates to .coords_nd to match the KDTree's non-dimensional basis, as flagged by Copilot review.

…e._get_kdtree Addressing review comments for PR #182: - Rebased on latest development (resolved conflict in Swarm.advection) - Fixed IndexError in Swarm.migrate using a flattened one-line assignment - Exposed _get_kdtree in EnhancedMeshVariable for direct access - Added regression tests for KDTree caching in both MeshVariable and Swarm

bknight1 · 2026-05-12T03:11:01Z

Consistency audit for _get_kdtree():

Verified that all _get_kdtree() implementations (Swarm, PICSwarm, MeshVariable) use non-dimensional coordinates (coords_nd or raw .data).
Standardised PICSwarm._get_kdtree() to use self._coord_var.data directly and renamed its internal cache attribute from _index to _kdtree for consistency across classes.
Confirmed that Mesh._get_domain_kdtree() uses non-dimensional centroids.

Copilot AI review requested due to automatic review settings May 11, 2026 06:59

bknight1 requested a review from lmoresi as a code owner May 11, 2026 06:59

Copilot started reviewing on behalf of bknight1 May 11, 2026 07:00 View session

Copilot AI reviewed May 11, 2026

View reviewed changes

This was referenced May 11, 2026

[BUG] - memory leak #176

Open

Core memory leak fixes: Cython deallocation & callback hardening #181

Open

bknight1 added 3 commits May 12, 2026 10:09

bknight1 force-pushed the feature/kdtree-refactor branch from 9f91b11 to e5c6f96 Compare May 12, 2026 02:53

bknight1 force-pushed the feature/kdtree-refactor branch from e5c6f96 to 1229937 Compare May 12, 2026 02:54

bknight1 force-pushed the feature/kdtree-refactor branch from 1229937 to 60a6e43 Compare May 12, 2026 02:58

bknight1 force-pushed the feature/kdtree-refactor branch from 60a6e43 to d549f13 Compare May 12, 2026 03:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Consolidate and unify cached spatial indexing (KDTree)#182

Consolidate and unify cached spatial indexing (KDTree)#182
bknight1 wants to merge 4 commits into
developmentfrom
feature/kdtree-refactor

bknight1 commented May 11, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

lmoresi commented May 11, 2026

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

bknight1 commented May 11, 2026

Key changes:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

lmoresi commented May 11, 2026

What's now on development

To do

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

bknight1 commented May 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

What's now on `development`