Commit 0b103f50 authored by inhuszar's avatar inhuszar
Browse files

Switched to cKDTree for higher performance.

parent 069e9426
......@@ -21,7 +21,7 @@
import ctypes
import numpy as np
cimport numpy as np
from sklearn.neighbors import KDTree
from scipy.spatial import cKDTree
from libc.stdlib cimport malloc, free
# 2D functions
......@@ -137,7 +137,7 @@ def invert2d(field, output=None, vectorder=(0, 1), pretx=None, posttx=None,
# Use the k-d tree algorithm to find the nearest forward-mapped point for
# each mapped grid point of the backward transformation (invwarp).
tree = KDTree(fwd_mapped, leaf_size=2)
tree = cKDTree(fwd_mapped, leafsize=10000000)
bwd_locations = np.stack(
np.mgrid[tuple(slice(ax) for ax in backward_shape)],
axis=-1).reshape((-1, 2))
......@@ -219,7 +219,7 @@ def jacobian2d(field, locations, pretx=None, vectorder=(0, 1), maxiter=1000):
ic[point] = &c_input_coordinates[point, 0]
# Find nearest grid point to each location
tree = KDTree(fwd_grid_coordinates, leaf_size=2)
tree = cKDTree(fwd_grid_coordinates, leafsize=10000000)
cdef np.ndarray[int, ndim=1, mode="c"] nearest_neighbour = \
np.ascontiguousarray(
tree.query(locations, k=1)[1].ravel(), dtype=ctypes.c_int)
......@@ -305,7 +305,7 @@ def invert3d(field, output=None, vectorder=(0, 1, 2), pretx=None, posttx=None,
# Use the k-d tree algorithm to find the nearest forward-mapped point for
# each mapped grid point of the backward transformation (invwarp).
tree = KDTree(fwd_mapped, leaf_size=2)
tree = cKDTree(fwd_mapped, leafsize=10000000)
bwd_locations = np.stack(
np.mgrid[tuple(slice(ax) for ax in backward_shape)],
axis=-1).reshape((-1, 3))
......@@ -405,7 +405,7 @@ def local_affine(field, newlocations, gridlocations, output=None,
p_gridlocations[point] = &c_gridlocations[point, 0]
# Find the nearest gridpoint for every new location using a k-D tree
tree = KDTree(gridlocations, leaf_size=2)
tree = cKDTree(gridlocations, leafsize=10000000)
nearest = tree.query(newlocations, k=1)[1].ravel()
cdef np.ndarray[int, ndim=1, mode="c"] c_nearest = \
np.ascontiguousarray(nearest, dtype=ctypes.c_int)
......@@ -424,7 +424,7 @@ def local_affine(field, newlocations, gridlocations, output=None,
for point in range(n_gridpoints):
p_mapped_gridlocations[point] = &c_mapped_gridlocations[point, 0]
# Call C function to calculate the local affine for each location
# Choose the appropriate version of the C function
if spacedim == 2:
cfunc = CalculateLocalAffine2D
elif spacedim == 3:
......@@ -492,7 +492,7 @@ def jacobian3d(field, locations, pretx=None, vectorder=(0, 1, 2), maxiter=1000):
ic[point] = &c_input_coordinates[point, 0]
# Find nearest grid point to each location
tree = KDTree(fwd_grid_coordinates, leaf_size=2)
tree = cKDTree(fwd_grid_coordinates, leafsize=10000000)
cdef np.ndarray[int, ndim=1, mode="c"] nearest_neighbour = \
np.ascontiguousarray(
tree.query(locations, k=1)[1].ravel(), dtype=ctypes.c_int)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment