Added tests. Updated docstrings/typehints

Author: Ramy-Badr-Ahmed

data_structures/kd_tree/build_kdtree.py

@@ -1,17 +1,16 @@
-from typing import List, Optional
+from typing import Optional
 from .kd_node import KDNode
 
-
-def build_kdtree(points: List[List[float]], depth: int = 0) -> Optional[KDNode]:
+def build_kdtree(points: list[list[float]], depth: int = 0) -> Optional[KDNode]:
     """
-    Builds a KD-Tree from a set of k-dimensional points.
+    Builds a KD-Tree from a list of points.
 
     Args:
-        points (List[List[float]]): A list of k-dimensional points (each point is a list of floats).
-        depth (int): The current depth in the tree. Used to determine the splitting axis. Defaults to 0.
+        points (list[list[float]]): The list of points to build the KD-Tree from.
+        depth (int): The current depth in the tree (used to determine axis for splitting).
 
     Returns:
-        Optional[KDNode]: The root of the KD-Tree or None if the input list is empty.
+        Optional[KDNode]: The root node of the KD-Tree.
     """
     if not points:
         return None
@@ -27,5 +26,5 @@ def build_kdtree(points: List[List[float]], depth: int = 0) -> Optional[KDNode]:
     return KDNode(
         point=points[median_idx],
         left=build_kdtree(points[:median_idx], depth + 1),
-        right=build_kdtree(points[median_idx + 1 :], depth + 1),
-    )
+        right=build_kdtree(points[median_idx + 1:], depth + 1),
+    )

data_structures/kd_tree/example/example_usage.py

@@ -1,5 +1,4 @@
 import numpy as np
-from typing import List
 from hypercube_points import hypercube_points
 from data_structures.kd_tree.build_kdtree import build_kdtree
 from data_structures.kd_tree.nearest_neighbour_search import nearest_neighbour_search
@@ -11,7 +10,7 @@ def main() -> None:
     in a 10-dimensional hypercube and performing a nearest neighbor search.
     """
     num_points: int = 5000
-    cube_size: int = 10
+    cube_size: float = 10.0  # Size of the hypercube (edge length)
     num_dimensions: int = 10
 
     # Generate random points within the hypercube
@@ -20,7 +19,7 @@ def main() -> None:
 
     # Generate a random query point within the same space
     rng = np.random.default_rng()
-    query_point: List[float] = rng.random(num_dimensions).tolist()
+    query_point: list[float] = rng.random(num_dimensions).tolist()
 
     # Perform nearest neighbor search
     nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
@@ -33,6 +32,5 @@ def main() -> None:
     print(f"Distance: {nearest_dist:.4f}")
     print(f"Nodes visited: {nodes_visited}")
 
-
 if __name__ == "__main__":
     main()

data_structures/kd_tree/example/hypercube_points.py

@@ -1,20 +1,16 @@
 import numpy as np
-from typing import Union
 
-
-def hypercube_points(
-    num_points: int, hypercube_size: Union[int, float], num_dimensions: int
-) -> np.ndarray:
+def hypercube_points(num_points: int, hypercube_size: float, num_dimensions: int) -> np.ndarray:
     """
     Generates random points uniformly distributed within an n-dimensional hypercube.
 
     Args:
-        num_points (int): The number of random points to generate.
-        hypercube_size (Union[int, float]): The size of the hypercube (side length).
-        num_dimensions (int): The number of dimensions of the hypercube.
+        num_points (int): Number of points to generate.
+        hypercube_size (float): Size of the hypercube.
+        num_dimensions (int): Number of dimensions of the hypercube.
 
     Returns:
-        np.ndarray: An array of shape (num_points, num_dimensions) with the generated points.
+        np.ndarray: An array of shape (num_points, num_dimensions) with generated points.
     """
     rng = np.random.default_rng()
     return hypercube_size * rng.random((num_points, num_dimensions))

data_structures/kd_tree/kd_node.py

@@ -1,29 +1,23 @@
-from typing import List, Optional
-
+from typing import Optional
 
 class KDNode:
     """
     Represents a node in a KD-Tree.
 
     Attributes:
-        point (List[float]): The k-dimensional point stored in this node.
-        left (Optional[KDNode]): The left subtree of this node.
-        right (Optional[KDNode]): The right subtree of this node.
+        point (list[float]): The point stored in this node.
+        left (Optional[KDNode]): The left child node.
+        right (Optional[KDNode]): The right child node.
     """
 
-    def __init__(
-        self,
-        point: List[float],
-        left: Optional["KDNode"] = None,
-        right: Optional["KDNode"] = None,
-    ) -> None:
+    def __init__(self, point: list[float], left: Optional['KDNode'] = None, right: Optional['KDNode'] = None) -> None:
         """
-        Initializes a KDNode with a point and optional left and right children.
+        Initializes a KDNode with the given point and child nodes.
 
         Args:
-            point (List[float]): The k-dimensional point to be stored in this node.
-            left (Optional[KDNode]): The left subtree of this node. Defaults to None.
-            right (Optional[KDNode]): The right subtree of this node. Defaults to None.
+            point (list[float]): The point stored in this node.
+            left (Optional[KDNode]): The left child node.
+            right (Optional[KDNode]): The right child node.
         """
         self.point = point
         self.left = left

data_structures/kd_tree/nearest_neighbour_search.py

@@ -1,39 +1,31 @@
-from typing import Optional, List, Tuple
-from .kd_node import KDNode
+from typing import Optional
+from data_structures.kd_tree.kd_node import KDNode
 
-
-def nearest_neighbour_search(
-    root: Optional[KDNode], query_point: List[float]
-) -> Tuple[Optional[List[float]], float, int]:
+def nearest_neighbour_search(root: Optional[KDNode], query_point: list[float]) -> tuple[Optional[list[float]], float, int]:
     """
     Performs a nearest neighbor search in a KD-Tree for a given query point.
 
     Args:
         root (Optional[KDNode]): The root node of the KD-Tree.
-        query_point (List[float]): The point for which the nearest neighbor is being searched.
+        query_point (list[float]): The point for which the nearest neighbor is being searched.
 
     Returns:
-        Tuple[Optional[List[float]], float, int]:
+        tuple[Optional[list[float]], float, int]:
             - The nearest point found in the KD-Tree to the query point.
             - The squared distance to the nearest point.
             - The number of nodes visited during the search.
     """
-    nearest_point: Optional[List[float]] = None
+    nearest_point: Optional[list[float]] = None
     nearest_dist: float = float("inf")
     nodes_visited: int = 0
 
     def search(node: Optional[KDNode], depth: int = 0) -> None:
         """
-        Recursively searches the KD-Tree to find the nearest point to the query point.
+        Recursively searches the KD-Tree for the nearest neighbor.
 
         Args:
-            node (Optional[KDNode]): The current node being examined.
-            depth (int): The current depth of the tree, which determines the axis to split on.
-
-        Updates:
-            nearest_point: The closest point found so far in the KD-Tree.
-            nearest_dist: The squared distance from the query point to the nearest point found.
-            nodes_visited: The number of nodes visited during the search.
+            node (Optional[KDNode]): The current node in the KD-Tree.
+            depth (int): The current depth in the tree.
         """
         nonlocal nearest_point, nearest_dist, nodes_visited
         if node is None:
@@ -43,18 +35,15 @@ def search(node: Optional[KDNode], depth: int = 0) -> None:
 
         # Calculate the current distance (squared distance)
         current_point = node.point
-        current_dist = sum(
-            (query_coord - point_coord) ** 2
-            for query_coord, point_coord in zip(query_point, current_point)
-        )
+        current_dist = sum((query_coord - point_coord) ** 2 for query_coord, point_coord in zip(query_point, current_point))
 
         # Update nearest point if the current node is closer
         if nearest_point is None or current_dist < nearest_dist:
             nearest_point = current_point
             nearest_dist = current_dist
 
         # Determine which subtree to search first (based on axis and query point)
-        k = len(query_point)  # dimensionality of points
+        k = len(query_point)  # Dimensionality of points
         axis = depth % k
 
         if query_point[axis] <= current_point[axis]:
@@ -67,7 +56,7 @@ def search(node: Optional[KDNode], depth: int = 0) -> None:
         # Search the nearer subtree first
         search(nearer_subtree, depth + 1)
 
-        # If the further subtree has a closer point, search it
+        # If the further subtree has a closer point
         if (query_point[axis] - current_point[axis]) ** 2 < nearest_dist:
             search(further_subtree, depth + 1)
 

data_structures/kd_tree/tests/__init__.py

@@ -0,0 +1,70 @@
+import unittest
+import numpy as np
+from data_structures.kd_tree.build_kdtree import build_kdtree
+from data_structures.kd_tree.nearest_neighbour_search import nearest_neighbour_search
+from data_structures.kd_tree.kd_node import KDNode
+from data_structures.kd_tree.example.hypercube_points import hypercube_points
+
+class TestKDTree(unittest.TestCase):
+
+    def setUp(self):
+        """
+        Set up test data.
+        """
+        self.num_points = 10
+        self.cube_size = 10.0
+        self.num_dimensions = 2
+        self.points = hypercube_points(self.num_points, self.cube_size, self.num_dimensions)
+        self.kdtree = build_kdtree(self.points.tolist())
+
+    def test_build_kdtree(self):
+        """
+        Test that KD-Tree is built correctly.
+        """
+        # Check if root is not None
+        self.assertIsNotNone(self.kdtree)
+
+        # Check if root has correct dimensions
+        self.assertEqual(len(self.kdtree.point), self.num_dimensions)
+
+        # Check that the tree is balanced to some extent (simplistic check)
+        self.assertIsInstance(self.kdtree, KDNode)
+
+    def test_nearest_neighbour_search(self):
+        """
+        Test the nearest neighbor search function.
+        """
+        rng = np.random.default_rng()
+        query_point = rng.random(self.num_dimensions).tolist()
+
+        nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
+            self.kdtree, query_point
+        )
+
+        # Check that nearest point is not None
+        self.assertIsNotNone(nearest_point)
+
+        # Check that distance is a non-negative number
+        self.assertGreaterEqual(nearest_dist, 0)
+
+        # Check that nodes visited is a non-negative integer
+        self.assertGreaterEqual(nodes_visited, 0)
+
+    def test_edge_cases(self):
+        """
+        Test edge cases such as an empty KD-Tree.
+        """
+        empty_kdtree = build_kdtree([])
+        query_point = [0.0] * self.num_dimensions
+
+        nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
+            empty_kdtree, query_point
+        )
+
+        # With an empty KD-Tree, nearest_point should be None
+        self.assertIsNone(nearest_point)
+        self.assertEqual(nearest_dist, float("inf"))
+        self.assertEqual(nodes_visited, 0)
+
+if __name__ == '__main__':
+    unittest.main()