class TestKDTree2d(unittest.TestCase): def setUp(self): points = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)] self.tree = KDTree(points) def test_traversal(self): logger.info("PreOrder:") self.tree.traversal(lambda n: logger.info(n), "preorder") logger.info("InOrder:") self.tree.traversal(lambda n: logger.info(n), "inorder") logger.info("PostOrder:") self.tree.traversal(lambda n: logger.info(n), "postorder") def test_closest(self): dist, node, count, nodes = self.tree.closest((2.3, 3.3)) logger.info("dist = {0}, node = {1}, visit_count = {2}".format( dist, node, count)) logger.info("visited nodes = {0}".format(nodes)) self.assertTrue(np.all(node.point == (2, 3)), "actual is {0}".format(node.point)) dist, node, count, nodes = self.tree.closest((2.0, 4.5)) logger.info("dist = {0}, node = {1}, visit_count = {2}".format( dist, node, count)) logger.info("visited nodes = {0}".format(nodes)) self.assertTrue(np.all(node.point == (2, 3)), "actual is {0}".format(node.point)) dist, node, count, nodes = self.tree.closest((6.9, 4.2)) logger.info("dist = {0}, node = {1}, visit_count = {2}".format( dist, node, count)) logger.info("visited nodes = {0}".format(nodes)) self.assertTrue(np.all(node.point == (5, 4))) def test_visualization(self): self.tree.traversal(draw_point, 'preorder') show_closest(self.tree, (2.3, 3.3), 'r') show_closest(self.tree, (2.0, 4.5), 'g') show_closest(self.tree, (6.8, 4.2), 'm') plt.show()
class TestKDTree1d(unittest.TestCase): def setUp(self): np.random.seed(0) points = np.random.randint(0, 50, 10) self.tree = KDTree(np.transpose([points])) def test_traversal(self): logger.info("PreOrder:") self.tree.traversal(lambda n: logger.info(n), "preorder") logger.info("InOrder:") self.tree.traversal(lambda n: logger.info(n), "inorder") logger.info("PostOrder:") self.tree.traversal(lambda n: logger.info(n), "postorder") def test_closest(self): dist, node, count, nodes = self.tree.closest([15]) logger.info("dist = {0}, node = {1}, visit_count = {2}".format( dist, node, count)) logger.info("visited nodes = {0}".format(nodes))