Consolidate and unify cached spatial indexing (KDTree)

src/underworld3/discretisation/discretisation_mesh.py

@@ -1239,11 +1239,12 @@ def _build_vertex_map(self):
         Uses coordinate matching at extraction time (before any
         deformation). Cached permanently since topology doesn't change.
         """
-        if hasattr(self, '_vertex_map') and self._vertex_map is not None:
+        if hasattr(self, "_vertex_map") and self._vertex_map is not None:
             return self._vertex_map
 
-        tree = uw.kdtree.KDTree(self.X.coords)
-        dists, indices = tree.query(self.parent.X.coords, sqr_dists=False)
+        # Use cached KDTree from coordinate variable
+        tree = self.X._get_kdtree()
+        dists, indices = tree.query(self.parent.X.coords_nd, sqr_dists=False)
         matched = dists < 1.0e-10
 
         # parent_rows[i] -> sub_rows[i]: matched vertex pairs
@@ -1433,8 +1434,8 @@ def _build_dof_map(self, parent_var, sub_var):
         if key in self._dof_maps:
             return self._dof_maps[key]
 
-        tree = uw.kdtree.KDTree(sub_var.coords)
-        dists, indices = tree.query(parent_var.coords, sqr_dists=False)
+        tree = sub_var._get_kdtree()
+        dists, indices = tree.query(parent_var.coords_nd, sqr_dists=False)
         matched = dists < 1.0e-10
 
         # indices[matched] maps parent row → sub row
@@ -3455,11 +3456,24 @@ def _get_mesh_centroids(self):
     def _get_domain_centroids(self):
 
         import numpy as np
+        from underworld3.utilities import gather_data
 
         domain_centroid = self._centroids.mean(axis=0)
         all_centroids = gather_data(domain_centroid, bcast=True).reshape(-1, self.dim)
         return all_centroids
 
+    def _get_domain_kdtree(self):
+        import underworld3 as uw
+        if (
+            not hasattr(self, "_domain_kdtree")
+            or self._domain_kdtree is None
+            or getattr(self, "_domain_kdtree_version", -1) != self._mesh_version
+        ):
+            centroids = self._get_domain_centroids()
+            self._domain_kdtree = uw.kdtree.KDTree(centroids)
+            self._domain_kdtree_version = self._mesh_version
+        return self._domain_kdtree
+
     def get_min_radius_old(self) -> float:
         """
         This method returns the global minimum distance from any cell centroid to a face.

src/underworld3/discretisation/discretisation_mesh_variables.py

@@ -914,6 +914,22 @@ def unpack_uw_data_from_petsc(self, squeeze=True, sync=True):
         else:
             return data_array_3d
 
+    def _get_kdtree(self):
+        """
+        Return a cached KDTree for this variable's DOF locations.
+        Rebuilds automatically if the parent mesh has deformed.
+        """
+        # Use non-dimensional coordinates for internal caching (avoids UnitAwareArray overhead)
+        if (
+            not hasattr(self, "_kdtree")
+            or self._kdtree is None
+            or getattr(self, "_kdtree_mesh_version", -1) != self.mesh._mesh_version
+        ):
+            self._kdtree = uw.kdtree.KDTree(self.coords_nd)
+            self._kdtree_mesh_version = self.mesh._mesh_version
+
+        return self._kdtree
+
     def rbf_interpolate(self, new_coords, meth=0, p=2, verbose=False, nnn=None, rubbish=None):
         """Interpolate variable data to new coordinates using RBF.
 
@@ -954,10 +970,9 @@ def rbf_interpolate(self, new_coords, meth=0, p=2, verbose=False, nnn=None, rubb
         if verbose and uw.mpi.rank == 0:
             print("Building K-D tree", flush=True)
 
-        # Use non-dimensional coordinates for internal RBF interpolation KDTree
-        mesh_kdt = uw.kdtree.KDTree(self.coords_nd)
-        values = mesh_kdt.rbf_interpolator_local(new_coords, D, nnn, p=p, verbose=verbose)
-        del mesh_kdt
+        # Use cached KDTree for interpolation
+        kdt = self._get_kdtree()
+        values = kdt.rbf_interpolator_local(new_coords, D, nnn, p=p, verbose=verbose)
 
         return values
 

src/underworld3/discretisation/enhanced_variables.py

@@ -554,6 +554,10 @@ def rbf_interpolate(self, *args, **kwargs):
         """RBF interpolation."""
         return self._base_var.rbf_interpolate(*args, **kwargs)
 
+    def _get_kdtree(self):
+        """KDTree access."""
+        return self._base_var._get_kdtree()
+
     def pack_raw_data_to_petsc(self, *args, **kwargs):
         """Pack raw data to PETSc format."""
         return self._base_var.pack_raw_data_to_petsc(*args, **kwargs)

src/underworld3/swarm.py

@@ -1100,8 +1100,8 @@ def _rbf_reduce_to_meshVar(self, meshVar, verbose=False):
 
         # 1 - Average particles to nodes with distance weighted average
 
-        # Use non-dimensional coordinates for internal KDTree (matches swarm.data coordinate system)
-        kd = uw.kdtree.KDTree(meshVar.coords_nd)
+        # Use cached KDTree for interpolation (avoids redundant index construction)
+        kd = meshVar._get_kdtree()
 
         with self.swarm.access():
             d, n = kd.query(self.swarm.data, k=1, sqr_dists=False)  # need actual distances
@@ -1414,16 +1414,13 @@ def rbf_interpolate(self, new_coords, verbose=False, nnn=None):
             D = raw_data[not_remeshed].copy()
 
             kdt = uw.kdtree.KDTree(self.swarm._particle_coordinates.data[not_remeshed, :])
+            values = kdt.rbf_interpolator_local(new_coords, D, nnn, 2, verbose)
         else:
             D = raw_data.copy()
-            kdt = uw.kdtree.KDTree(self.swarm._particle_coordinates.data[:, :])
-
-            # kdt.build_index()
-
+            # Use cached KDTree for standard swarms
+            kdt = self.swarm._get_kdtree()
             values = kdt.rbf_interpolator_local(new_coords, D, nnn, 2, verbose)
 
-            del kdt
-
         return values
 
     @property
@@ -2314,12 +2311,12 @@ def _update_proxy_variables(self):
         """
         if self.update_type == 0:
             # Use non-dimensional coordinates for internal level set KDTree
-            kd = uw.kdtree.KDTree(self._meshLevelSetVars[0].coords_nd)
+            kd = self._meshLevelSetVars[0]._get_kdtree()
 
             n_distance, n_indices = kd.query(
                 self.swarm._particle_coordinates.data, k=self.nnn, sqr_dists=False
             )
-            kd_swarm = uw.kdtree.KDTree(self.swarm._particle_coordinates.data)
+            kd_swarm = self.swarm._get_kdtree()
             # n, d, b = kd_swarm.find_closest_point(self._meshLevelSetVars[0].coords)
             d, n = kd_swarm.query(self._meshLevelSetVars[0].coords, k=1, sqr_dists=False)
 
@@ -2630,6 +2627,22 @@ def _invalidate_canonical_data(self):
             if hasattr(var, "_canonical_data"):
                 var._canonical_data = None
 
+        # Invalidate cached spatial index
+        self._kdtree = None
+
+    def _get_kdtree(self):
+        """
+        Return a cached KDTree for the swarm particle coordinates.
+        Invalidated automatically whenever particles migrate or positions change.
+        """
+        # Note: self.data returns unit-aware array if units are active,
+        # but kdtree construction expects non-dimensional values.
+        # Use _particle_coordinates.data directly.
+        if not hasattr(self, "_kdtree") or self._kdtree is None:
+            self._kdtree = uw.kdtree.KDTree(self._particle_coordinates.data)
+
+        return self._kdtree
+
     def _route_by_nearest_centroid(self):
         """Migrate every particle to the rank whose domain-centroid is closest.
 
@@ -3320,8 +3333,7 @@ def migrate(
         if delete_lost_points is None:
             delete_lost_points = self.clip_to_mesh
 
-        centroids = self.mesh._get_domain_centroids()
-        mesh_domain_kdtree = uw.kdtree.KDTree(centroids)
+        mesh_domain_kdtree = self.mesh._get_domain_kdtree()
 
         # This will only worry about particles that are not already claimed !
         #
@@ -3366,7 +3378,10 @@ def migrate(
                     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.reshape(-1)[not_my_points] = rank.reshape(
+                        -1, it + 1
+                    )[:, it].flatten()
 
                 self.dm.restoreField("DMSwarm_rank")
                 self.dm.restoreField("DMSwarmPIC_coor")
@@ -4280,8 +4295,7 @@ def _data_layout(self, i, j=None):
     @timing.routine_timer_decorator
     def _get_map(self, var):
         # generate tree if not avaiable
-        if not self._index:
-            self._index = uw.kdtree.KDTree(self.data)
+        kd = self._get_kdtree()
 
         # get or generate map
         meshvar_coords = var._meshVar.coords
@@ -4295,7 +4309,7 @@ def _get_map(self, var):
         h.update(meshvar_coords)
         digest = h.intdigest()
         if digest not in self._nnmapdict:
-            self._nnmapdict[digest] = self._index.query(meshvar_coords, k=1, sqr_dists=False)[1]
+            self._nnmapdict[digest] = kd.query(meshvar_coords, k=1, sqr_dists=False)[1]
         return self._nnmapdict[digest]
 
     @timing.routine_timer_decorator

src/underworld3/swarms/pic_swarm.py

@@ -223,7 +223,7 @@ def __init__(self, mesh, recycle_rate=0, verbose=False):
             )
 
         self._X0_uninitialised = True
-        self._index = None
+        self._kdtree = None
         self._nnmapdict = {}
 
         super().__init__()
@@ -987,7 +987,7 @@ def __exit__(self, *args):
                     self.em_swarm.dm.migrate(remove_sent_points=True)
 
                     # void these things too
-                    self.em_swarm._index = None
+                    self.em_swarm._kdtree = None
                     self.em_swarm._nnmapdict = {}
 
                 # do var updates
@@ -1042,12 +1042,21 @@ def _data_layout(self, i, j=None):
         if self.vtype == uw.VarType.MATRIX:
             return i + j * self.shape[0]
 
+    def _get_kdtree(self):
+        """
+        Return a cached KDTree for the swarm particle coordinates.
+        Invalidated automatically whenever particles migrate or positions change.
+        """
+        if not hasattr(self, "_kdtree") or self._kdtree is None:
+            with self.access():
+                self._kdtree = uw.kdtree.KDTree(self._coord_var.data)
+
+        return self._kdtree
+
     @timing.routine_timer_decorator
     def _get_map(self, var):
         # generate tree if not avaiable
-        if not self._index:
-            with self.access():
-                self._index = uw.kdtree.KDTree(self.data)
+        kd = self._get_kdtree()
 
         # get or generate map
         meshvar_coords = var._meshVar.coords
@@ -1061,8 +1070,8 @@ def _get_map(self, var):
         h.update(meshvar_coords)
         digest = h.intdigest()
         if digest not in self._nnmapdict:
-            # self._nnmapdict[digest] = self._index.find_closest_point(meshvar_coords)[0]
-            self._nnmapdict[digest] = self._index.query(meshvar_coords, k=1, sqr_dists=False)[0]
+            # self._nnmapdict[digest] = self._kdtree.find_closest_point(meshvar_coords)[0]
+            self._nnmapdict[digest] = kd.query(meshvar_coords, k=1, sqr_dists=False)[1]
         return self._nnmapdict[digest]
 
     @timing.routine_timer_decorator

tests/test_0780_memprobe.py

@@ -151,3 +151,56 @@ def f(x):
     assert f(3) == 6
     captured = capsys.readouterr()
     assert "[memprobe] test-fn" in captured.out
+
+
+@pytest.mark.level_1
+@pytest.mark.tier_a
+def test_mesh_variable_kdtree_caching():
+    """MeshVariable._get_kdtree must cache the tree and reuse it."""
+    mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4))
+    v = uw.discretisation.MeshVariable("v", mesh, 1)
+
+    gc.collect()
+    before_total = uw.kdtree.total_constructed()
+
+    # First access builds
+    kd1 = v._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 1
+
+    # Second access reuses
+    kd2 = v._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 1
+    assert kd1 is kd2
+
+    # Mesh deformation/version change invalidates
+    mesh._mesh_version += 1
+    kd3 = v._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 2
+    assert kd3 is not kd1
+
+
+@pytest.mark.level_1
+@pytest.mark.tier_a
+def test_swarm_kdtree_caching():
+    """Swarm._get_kdtree must cache the tree and reuse it."""
+    mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4))
+    swarm = uw.swarm.Swarm(mesh)
+    swarm.populate(fill_param=1)
+
+    gc.collect()
+    before_total = uw.kdtree.total_constructed()
+
+    # First access builds
+    kd1 = swarm._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 1
+
+    # Second access reuses
+    kd2 = swarm._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 1
+    assert kd1 is kd2
+
+    # Migration/Invalidation should drop the cache
+    swarm._invalidate_canonical_data()
+    kd3 = swarm._get_kdtree()
+    assert uw.kdtree.total_constructed() - before_total == 2
+    assert kd3 is not kd1