Using cykdtree¶
From Python¶
The cykdtree
package provides a python interface to a C++ kdtree
implementation. The primary interface to cykdtree
is the PyKDTree
class:
from cykdtree import PyKDTree
import numpy as np
# Randomly generate a million particle positions
np.random.seed(0xdeadbeef)
positions = np.random.random((int(1e6), 3))
# Initialize PyKDTree object
tree = PyKDTree(
positions,
left_edge=np.array([0., 0., 0.]),
right_edge=np.array([1., 1., 1.]),
periodic=np.array([True, True, True]),
leafsize=30
)
# Working with the tree
print("The number of tree leaves is %i" % len(tree.leaves))
print("")
node = tree.get(np.array([0.5, 0.5, 0.5]))
print("The KDTree node containing the point [0.5, 0.5, 0.5] is:\n%s" % node)
print("")
print("The KDTree IDs of that node's neighbors:\n%s" % node.neighbors)
This should print the following output:
The number of tree leaves is 65536
The KDTree node containing the point [0.5, 0.5, 0.5] is:
PyNode(id=22235, npts=15, start_idx=339280, stop_idx=339295,
left_edge=[ 0.49900133 0.47859409 0.466449 ],
right_edge=[ 0.52928342 0.50126165 0.50050787])
The KDTree IDs of that node's neighbors:
[8190, 8191, 13778, 13783, 22232, 22233, 22234, 22235, 22236, 22237, 22238, 22239, 26816, 26824, 26825, 26826, 26827, 39773, 43589, 53828, 53829]
One can think of the KDTree as a spatially sorted list of particles. These
sorted indices are stored in the idx
attribute of the PyKDTree
. One can
use idx
and the slice
attribute of a PyNode
instance to find the
positions of the particles contained in that node:
node_particle_positions = positions[tree.idx[node.slice]]
print(node_particle_positions)
This should print the positions of all of the particles contained in the KDTree node we selected earlier.
From Cython¶
This is a bit more cumbersome, as your C++ compiler will need to be able to
find the cykdtree
headers when it builds your code. The most straightforward
way to get this working is to integrate cykdtree
into your setup.py
script:
ex = Extension("my_cython_cykdtree_wrapper",
sources=["my_cython_cykdtree_wrapper.pyx"]
include_dirs=[numpy.get_include, cykdtree.get_include()])
This will allow you to write Cython code that cimports code from cykdtree.