Skip to content

Use uw.kdtree.KDTree in mesh.adapt (drop scipy.spatial.cKDTree) #150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lmoresi merged 1 commit into from
Apr 28, 2026
Merged

Use uw.kdtree.KDTree in mesh.adapt (drop scipy.spatial.cKDTree) #150

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions src/underworld3/discretisation/discretisation_mesh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Copilot AI Apr 28, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

uw.kdtree.KDTree.query() returns i reshaped to 1D for k==1, but it leaves d as shape (n, 1). Here that makes matched 2D and then indices[matched] will raise because indices is 1D. Flatten dists (e.g., dists = numpy.asarray(dists).reshape(-1) / dists = dists[:,0]) before building the boolean mask, or adjust the query call/implementation so dists is 1D when k==1 (SciPy-compatible).

Suggested change
dists, indices = tree.query(self.parent.X.coords, sqr_dists=False)
dists, indices = tree.query(self.parent.X.coords, sqr_dists=False)
dists = numpy.asarray(dists).reshape(-1)

Copilot uses AI. Check for mistakes.
matched = dists < 1.0e-10

# parent_rows[i] -> sub_rows[i]: matched vertex pairs
Expand Down Expand Up @@ -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)
Comment on lines +1382 to +1384

Copilot AI Apr 28, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the IDW transfer path, old_coords is stored via numpy.array(var.coords) (strips UnitAwareArray metadata), but new_coords = old_var.coords can be unit-aware when model units are active. uw.kdtree.KDTree.query() will raise when the tree was built with dimensionless points but query coords have units. Consider using .coords_nd for both arrays (preferred for internal algorithms) or otherwise ensure both old_coords and new_coords are consistently unitless/unitful before building/querying the tree.

Copilot uses AI. Check for mistakes.

# Inverse distance weighting
weights = 1.0 / (dists + 1e-30)
Expand Down Expand Up @@ -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
Expand Down
Loading