From 85fa77bd4eea699f8e89fb03fb18851a81c58c12 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 28 Apr 2026 18:50:25 +1000 Subject: [PATCH] Use uw.kdtree.KDTree (nanoflann) instead of scipy.spatial.cKDTree in mesh.adapt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three sites in discretisation_mesh.py used scipy.spatial.cKDTree for the mesh-to-mesh variable transfer that runs during Mesh.adapt: - sub_vertex_map (line 1247) — parent-submesh vertex match - _re_extract_from_parent variable IDW transfer (line 1385) - _build_dof_map (line 1440) — parent-submesh DOF match These predate the nanoflann switch and were the only scipy.spatial.cKDTree usages left in src/. uw.kdtree.KDTree (nanoflann-backed) is already used in read_timestep, swarm migration, and the RBF interpolator path; running mesh.adapt on the same backend brings consistency and the small nanoflann-vs-scipy speedup. The .query API is scipy-compatible — same (d, i) return shape, same k= parameter — but defaults to squared distances. Passing sqr_dists=False preserves the previous linear-distance semantics so the existing threshold (1.0e-10 for exact-match) and IDW weighting (1.0 / (dists + 1e-30)) are unchanged. Verified: pytest -m "level_1 and tier_a" on amr-dev passes 56/3/0. Underworld development team with AI support from Claude Code --- .../discretisation/discretisation_mesh.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 4161675bb..c2d760159 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -1242,10 +1242,8 @@ def _build_vertex_map(self): if hasattr(self, '_vertex_map') and self._vertex_map is not None: return self._vertex_map - from scipy.spatial import cKDTree - - tree = cKDTree(self.X.coords) - dists, indices = tree.query(self.parent.X.coords) + tree = uw.kdtree.KDTree(self.X.coords) + dists, indices = tree.query(self.parent.X.coords, sqr_dists=False) matched = dists < 1.0e-10 # parent_rows[i] -> sub_rows[i]: matched vertex pairs @@ -1378,13 +1376,12 @@ def mesh_update_callback(array, change_context): # Interpolate from backed-up data via kd-tree IDW if var_name in old_var_backups: try: - from scipy.spatial import cKDTree old_coords, old_data = old_var_backups[var_name] new_coords = old_var.coords - tree = cKDTree(old_coords) + tree = uw.kdtree.KDTree(old_coords) nnn = 3 if self.dim == 2 else 4 - dists, indices = tree.query(new_coords, k=nnn) + dists, indices = tree.query(new_coords, k=nnn, sqr_dists=False) # Inverse distance weighting weights = 1.0 / (dists + 1e-30) @@ -1431,14 +1428,13 @@ def _build_dof_map(self, parent_var, sub_var): Returns (sub_rows, parent_rows) — numpy arrays of matching DOF indices. """ import numpy as np - from scipy.spatial import cKDTree key = (id(parent_var), id(sub_var)) if key in self._dof_maps: return self._dof_maps[key] - tree = cKDTree(sub_var.coords) - dists, indices = tree.query(parent_var.coords) + tree = uw.kdtree.KDTree(sub_var.coords) + dists, indices = tree.query(parent_var.coords, sqr_dists=False) matched = dists < 1.0e-10 # indices[matched] maps parent row → sub row