Skip to content
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

precompute finite value mask for kdtree query #2201

Merged
merged 9 commits into from
Jul 14, 2023
Merged

precompute finite value mask for kdtree query #2201

Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
0236e6b
precompute finite value mask for kdtree query
chrishavlin Jun 27, 2023
04b63f6
add a test
chrishavlin Jun 27, 2023
e8919e2
linting fixes
chrishavlin Jun 27, 2023
7d0eacc
add the non-finite points to the mask
chrishavlin Jun 29, 2023
fc0f1f3
Merge remote-tracking branch 'upstream/main' into scipy_1pt11_kdtree_…
chrishavlin Jul 3, 2023
175a27d
add _kdtree_handle
chrishavlin Jul 7, 2023
fd08137
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 7, 2023
68b438c
shorten test name
chrishavlin Jul 7, 2023
58c5fe5
review suggestions
chrishavlin Jul 13, 2023
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
12 changes: 10 additions & 2 deletions lib/cartopy/img_transform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,16 +268,24 @@ def regrid(array, source_x_coords, source_y_coords, source_proj, target_proj,
target_x_points.flatten(),
target_y_points.flatten())

# Find mask of valid points before querying kdtree: scipy >= 1.11 errors
# when querying nan points, might as well use for pykdtree too.
indices = np.zeros(target_xyz.shape[0], dtype=int)
Copy link
Contributor Author

Choose a reason for hiding this comment

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

was I right in having the indices array default to 0? I based this on the indices[mask]=0 line after this, wasn't entirely certain.

Copy link
Contributor

Choose a reason for hiding this comment

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

I think this is fine to use 0 here as you noted because I think it is just settting it to some valid index. But, I think we still need to add these non-finite locations you found here into the mask below as well? Something like mask = ~finite_xyz | (indices >= len(xyz))

finite_xyz = np.all(np.isfinite(target_xyz), axis=-1)

if _is_pykdtree:
kdtree = pykdtree.kdtree.KDTree(xyz)
# Use sqr_dists=True because we don't care about distances,
# and it saves a sqrt.
_, indices = kdtree.query(target_xyz, k=1, sqr_dists=True)
_, indices[finite_xyz] = kdtree.query(target_xyz[finite_xyz, :],
k=1,
sqr_dists=True)
greglucas marked this conversation as resolved.
Show resolved Hide resolved
else:
# Versions of scipy >= v0.16 added the balanced_tree argument,
# which caused the KDTree to hang with this input.
kdtree = scipy.spatial.cKDTree(xyz, balanced_tree=False)
_, indices = kdtree.query(target_xyz, k=1)
_, indices[finite_xyz] = kdtree.query(target_xyz[finite_xyz, :], k=1)

mask = indices >= len(xyz)
chrishavlin marked this conversation as resolved.
Show resolved Hide resolved
indices[mask] = 0

Expand Down