diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index bf3e40752..9488273ed 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/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. diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 510787459..7a1930084 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/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 diff --git a/src/underworld3/discretisation/enhanced_variables.py b/src/underworld3/discretisation/enhanced_variables.py index 6b2aea316..53353d633 100644 --- a/src/underworld3/discretisation/enhanced_variables.py +++ b/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) diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 22662a95f..2683b94ff 100644 --- a/src/underworld3/swarm.py +++ b/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 diff --git a/src/underworld3/swarms/pic_swarm.py b/src/underworld3/swarms/pic_swarm.py index 62d27ccf9..20a0462f2 100644 --- a/src/underworld3/swarms/pic_swarm.py +++ b/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 diff --git a/tests/test_0780_memprobe.py b/tests/test_0780_memprobe.py index d1b4b5414..8cf2d92d9 100644 --- a/tests/test_0780_memprobe.py +++ b/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