def test_KDTree_exceptions(self):
bucket_size = self.bucket_size
nr_points = self.nr_points
radius = self.radius
coords = random((nr_points, 3)) * 100000000000000
with self.assertRaises(Exception) as context:
kdt = kdtrees.KDTree(coords, bucket_size)
self.assertTrue("coordinate values should lie between -1e6 and 1e6" in
str(context.exception))
with self.assertRaises(Exception) as context:
kdt = kdtrees.KDTree(random((nr_points, 3 - 2)), bucket_size)
self.assertTrue("expected a Nx3 numpy array" in str(context.exception))
def test_KDTree_point_search(self):
"""Test searching all points within a certain radius of center.
Using the KDTree C module, search all point pairs that are
within radius, and compare the results to a manual search.
"""
bucket_size = self.bucket_size
nr_points = self.nr_points
for radius in (self.radius, 100 * self.radius):
for i in range(0, 10):
# kd tree search
coords = random((nr_points, 3))
center = random(3)
kdt = kdtrees.KDTree(coords, bucket_size)
points1 = kdt.search(center, radius)
points1.sort(key=lambda point: point.index) # noqa: E731
# manual search
points2 = []
for i in range(0, nr_points):
p = coords[i]
v = p - center
r = sqrt(dot(v, v))
if r <= radius:
point2 = kdtrees.Point(i, r)
points2.append(point2)
# compare results
self.assertEqual(len(points1), len(points2))
for point1, point2 in zip(points1, points2):
self.assertEqual(point1.index, point2.index)
self.assertAlmostEqual(point1.radius, point2.radius)
def test_KDTree_neighbor_search_simple(self):
"""Test all fixed radius neighbor search.
Test all fixed radius neighbor search using the KD tree C
module, and compare the results to those of a simple but
slow algorithm.
"""
bucket_size = self.bucket_size
nr_points = self.nr_points
radius = self.radius
for i in range(0, 10):
# KD tree search
coords = random((nr_points, 3))
kdt = kdtrees.KDTree(coords, bucket_size)
neighbors1 = kdt.neighbor_search(radius)
# same search, using a simple but slow algorithm
neighbors2 = kdt.neighbor_simple_search(radius)
# compare results
self.assertEqual(len(neighbors1), len(neighbors2))
key = lambda neighbor: (neighbor.index1, neighbor.index2
) # noqa: E731
neighbors1.sort(key=key)
neighbors2.sort(key=key)
for neighbor1, neighbor2 in zip(neighbors1, neighbors2):
self.assertEqual(neighbor1.index1, neighbor2.index1)
self.assertEqual(neighbor1.index2, neighbor2.index2)
self.assertAlmostEqual(neighbor1.radius, neighbor2.radius)
def test_KDTree_neighbor_search_manual(self):
"""Test all fixed radius neighbor search.
Test all fixed radius neighbor search using the KD tree C
module, and compare the results to those of a manual search.
"""
bucket_size = self.bucket_size
nr_points = self.nr_points // 10 # fewer points to speed up the test
for radius in (self.radius, 3 * self.radius):
for i in range(0, 5):
# KD tree search
coords = random((nr_points, 3))
kdt = kdtrees.KDTree(coords, bucket_size)
neighbors1 = kdt.neighbor_search(radius)
# manual search
neighbors2 = []
indices = argsort(coords[:, 0])
for j1 in range(nr_points):
index1 = indices[j1]
p1 = coords[index1]
for j2 in range(j1 + 1, nr_points):
index2 = indices[j2]
p2 = coords[index2]
if p2[0] - p1[0] > radius:
break
v = p1 - p2
r = sqrt(dot(v, v))
if r <= radius:
if index1 < index2:
i1, i2 = index1, index2
else:
i1, i2 = index2, index1
neighbor = kdtrees.Neighbor(i1, i2, r)
neighbors2.append(neighbor)
self.assertEqual(len(neighbors1), len(neighbors2))
key = lambda neighbor: (neighbor.index1, neighbor.index2
) # noqa: E731
neighbors1.sort(key=key)
neighbors2.sort(key=key)
for neighbor1, neighbor2 in zip(neighbors1, neighbors2):
self.assertEqual(neighbor1.index1, neighbor2.index1)
self.assertEqual(neighbor1.index2, neighbor2.index2)
self.assertAlmostEqual(neighbor1.radius, neighbor2.radius)
