def time_tree_construction(Ntime, LStime, ndim=2):
pts, le, re, ls = make_points(Ntime, ndim, leafsize=LStime)
t0 = time.time()
cykdtree.PyKDTree(pts, le, re, leafsize=LStime)
t1 = time.time()
print("{} {}D points, leafsize {}: took {} s".format(
Ntime, ndim, LStime, t1 - t0))
def time_neighbor_search(Ntime, LStime, ndim=2):
pts, le, re, ls = make_points(Ntime, ndim, leafsize=LStime)
tree = cykdtree.PyParallelKDTree(pts, le, re, leafsize=LStime)
t0 = time.time()
tree.get_neighbor_ids(0.5*np.ones(tree.ndim, 'double'))
t1 = time.time()
print("{} {}D points, leafsize {}: took {} s".format(Ntime, ndim, LStime, t1-t0))
def test_search_errors(periodic=False, ndim=2):
pts, le, re, ls = make_points(100, ndim)
tree = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
if not periodic:
assert_raises(ValueError, tree.get, np.ones(ndim, 'double'))
assert_raises(AssertionError, tree.get, np.zeros(ndim+1, 'double'))
def test_spawn_parallel(nproc=1, npts=20, ndim=2, periodic=False,
leafsize=3):
if MPI is None:
return
pts, le, re, ls = make_points(npts, ndim, leafsize=leafsize)
Tseri = cykdtree.spawn_parallel(pts, nproc, leafsize=leafsize,
left_edge=le, right_edge=re,
periodic=periodic)
def test_get_neighbor_ids(npts=100, ndim=2, periodic=False):
pts, le, re, ls = make_points(npts, ndim)
tree = cykdtree.PyKDTree(pts, le, re, leafsize=ls, periodic=periodic)
pos_list = [le, (le + re) / 2.]
if periodic:
pos_list.append(re)
for pos in pos_list:
tree.get_neighbor_ids(pos)
def test_save_load():
for periodic in (True, False):
for ndim in range(1, 5):
pts, le, re, ls = make_points(100, ndim)
tree = cykdtree.PyKDTree(pts, le, re, leafsize=ls,
periodic=periodic, data_version=ndim+12)
with tempfile.NamedTemporaryFile() as tf:
tree.save(tf.name)
restore_tree = cykdtree.PyKDTree.from_file(tf.name)
tree.assert_equal(restore_tree)
def test_PyKDTree(npts=100,
ndim=2,
periodic=False,
use_sliding_midpoint=False):
pts, le, re, ls = make_points(npts, ndim)
cykdtree.PyKDTree(pts,
le,
re,
leafsize=ls,
periodic=periodic,
use_sliding_midpoint=use_sliding_midpoint)
def test_plot2D_serial():
fname_test = "test_plot2D_serial.png"
pts, le, re, ls = make_points(100, 2)
tree = PyKDTree(pts, le, re, leafsize=ls)
axs = plot2D_serial(tree,
pts,
title="Serial Test",
plotfile=fname_test,
label_boxes=True)
os.remove(fname_test)
# plot2D_serial(tree, pts, axs=axs)
del axs
def test_consolidate_process_bounds(periodic=False, ndim=2):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
pts, le, re, ls = make_points(20, ndim, leafsize=3)
Tpara = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
LEpara, REpara = Tpara.consolidate_process_bounds()
assert_equal(LEpara.shape, (size, ndim))
assert_equal(REpara.shape, (size, ndim))
np.testing.assert_allclose(LEpara[rank,:], Tpara.left_edge)
np.testing.assert_allclose(REpara[rank,:], Tpara.right_edge)
def test_search(periodic=False, ndim=2):
pts, le, re, ls = make_points(100, ndim)
tree = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
if periodic:
vals = [0, 0.5, 1.0]
else:
vals = [0, 0.5, 0.9]
for v in vals:
pos = v*np.ones(ndim, 'double')
out = tree.get(pos)
if out is not None:
out.neighbors
def test_get_neighbor_ids(periodic=False, ndim=2):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
pts, le, re, ls = make_points(100, ndim)
tree = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
if periodic:
vals = [0, 0.5, 1.0]
else:
vals = [0, 0.5, 0.9]
for v in vals:
pos = v*np.ones(ndim, 'float')
tree.get_neighbor_ids(pos)
def test_plot2D_parallel():
if PyParallelKDTree is None:
return
fname_test = "test_plot2D_parallel.png"
pts, le, re, ls = make_points(100, 2)
tree = PyParallelKDTree(pts, le, re, leafsize=ls)
axs = plot2D_parallel(tree,
pts,
title="Parallel Test",
plotfile=fname_test,
label_boxes=True,
label_procs=True)
os.remove(fname_test)
# plot2D_parallel(tree, pts, axs=axs)
del axs
def test_consolidate_edges(periodic=False, ndim=2):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
pts, le, re, ls = make_points(20, ndim, leafsize=3)
Tpara = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
LEpara, REpara = Tpara.consolidate_edges()
if rank == 0:
Tseri = cykdtree.PyKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
LEseri, REseri = Tseri.consolidate_edges()
else:
LEseri, REseri = None, None
LEseri, REseri = comm.bcast((LEseri, REseri), root=0)
np.testing.assert_allclose(LEpara, LEseri)
np.testing.assert_allclose(REpara, REseri)
def test_consolidate(periodic=False, ndim=2):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
pts, le, re, ls = make_points(20, ndim, leafsize=3)
Tpara0 = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
Tpara = Tpara0.consolidate()
if rank == 0:
if False:
from cykdtree.plot import plot2D_serial
plot2D_serial(Tpara, label_boxes=True,
plotfile='test_consolidate.png')
Tseri = cykdtree.PyKDTree(pts, le, re, leafsize=ls,
periodic=periodic)
Tpara.assert_equal(Tseri, strict_idx=False)
else:
assert(Tpara is None)
def test_PyParallelKDTree_errors(periodic=False, ndim=2):
pts, le, re, ls = make_points(20, ndim, leafsize=3)
assert_raises(ValueError, cykdtree.PyParallelKDTree, pts,
le, re, leafsize=1)
def test_PyParallelKDTree(periodic=False, ndim=2, use_sliding_midpoint=False):
pts, le, re, ls = make_points(20, ndim, leafsize=3)
Tpara = cykdtree.PyParallelKDTree(pts, le, re, leafsize=ls,
periodic=periodic,
use_sliding_midpoint=use_sliding_midpoint)
def test_search_errors(npts=100, ndim=2):
pts, le, re, ls = make_points(npts, ndim)
tree = cykdtree.PyKDTree(pts, le, re, leafsize=ls)
assert_raises(ValueError, tree.get, re)
def test_PyKDTree_errors():
pts, le, re, ls = make_points(100, 2)
assert_raises(ValueError, cykdtree.PyKDTree, pts, le, re, leafsize=1)
