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)