test_kdtree.py

import numpy as np
import pytest

from data_structures.kd_tree.build_kdtree import build_kdtree
from data_structures.kd_tree.example.hypercube_points import hypercube_points
from data_structures.kd_tree.kd_node import KDNode
from data_structures.kd_tree.nearest_neighbour_search import nearest_neighbour_search


@pytest.mark.parametrize(
    ("num_points", "cube_size", "num_dimensions", "expected_result"),
    [
        (0, 10.0, 2, None),  # Empty points list
        (10, 10.0, 2, KDNode),  # 2D points
        (10, 10.0, 3, KDNode),  # 3D points
    ],
)
def test_build_kdtree(num_points, cube_size, num_dimensions, expected_result):
    """
    Test that KD-Tree is built correctly.

    Cases:
        - Empty points list.
        - Valid points list with correct dimensions.
    """
    points = hypercube_points(num_points, cube_size, num_dimensions).tolist() if num_points > 0 else []

    kdtree = build_kdtree(points)

    if expected_result is None:
        # Empty points list case
        assert kdtree is None, f"Expected None for empty points list, got {kdtree}"
    else:
        # Check if KD-Tree is built correctly
        assert kdtree is not None, "Expected a KDNode, got None"
        assert isinstance(kdtree, KDNode), f"Expected KDNode instance, got {type(kdtree)}"
        assert len(kdtree.point) == num_dimensions, f"Expected point dimension {num_dimensions}, got {len(kdtree.point)}"


def test_nearest_neighbour_search():
    """
    Test the nearest neighbor search function.
    """
    num_points = 10
    cube_size = 10.0
    num_dimensions = 2
    points = hypercube_points(num_points, cube_size, num_dimensions)
    kdtree = build_kdtree(points.tolist())

    rng = np.random.default_rng()
    query_point = rng.random(num_dimensions).tolist()

    nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
        kdtree, query_point
    )

    # Check that nearest point is not None
    assert nearest_point is not None

    # Check that distance is a non-negative number
    assert nearest_dist >= 0

    # Check that nodes visited is a non-negative integer
    assert nodes_visited >= 0


def test_edge_cases():
    """
    Test edge cases such as an empty KD-Tree.
    """
    empty_kdtree = build_kdtree([])
    query_point = [0.0] * 2  # Using a default 2D query point

    nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
        empty_kdtree, query_point
    )

    # With an empty KD-Tree, nearest_point should be None
    assert nearest_point is None
    assert nearest_dist == float("inf")
    assert nodes_visited == 0


if __name__ == "__main__":
    pytest.main()