Added parametric tests

Author: lucasimi

tests/test_unit_cover.py

@@ -2,7 +2,6 @@
 Unit tests for the cover algorithms.
 """
 
-import numpy as np
 import pytest
 
 from tdamapper.core import TrivialCover
@@ -14,51 +13,7 @@
     StandardCubicalCover,
 )
 from tdamapper.utils.unionfind import UnionFind
-
-
-def dataset_simple():
-    """
-    Create a simple dataset of points in a 2D space.
-
-    This dataset consists of four points forming the corners of a rectangle
-    such that two sides are longer than the other two.
-    """
-    return [
-        np.array([0.0, 1.0]),
-        np.array([1.1, 0.0]),
-        np.array([0.0, 0.0]),
-        np.array([1.1, 1.0]),
-    ]
-
-
-def dataset_random(dim=1, num=1000):
-    """
-    Create a random dataset of points in the unit square.
-    """
-    return [np.random.rand(dim) for _ in range(num)]
-
-
-def dataset_two_lines(num=1000):
-    """
-    Create a dataset consisting of two lines in the unit square.
-    One line is horizontal at y=0, the other is vertical at x=1.
-    """
-    t = np.linspace(0.0, 1.0, num)
-    line1 = np.array([[x, 0.0] for x in t])
-    line2 = np.array([[x, 1.0] for x in t])
-    return np.concatenate((line1, line2), axis=0)
-
-
-def dataset_grid(num=1000):
-    """
-    Create a grid dataset in the unit square.
-    The grid consists of points evenly spaced in both dimensions.
-    """
-    t = np.linspace(0.0, 1.0, num)
-    s = np.linspace(0.0, 1.0, num)
-    grid = np.array([[x, y] for x in t for y in s])
-    return grid
-
+from tests.test_utils import dataset_grid, dataset_simple, dataset_two_lines
 
 SIMPLE = dataset_simple()
 

tests/test_unit_vptree.py

@@ -1,25 +1,58 @@
+"""
+Unit tests for the vp-tree implementations.
+"""
+
 import random
 
 import numpy as np
+import pytest
 
 from tdamapper.utils.metrics import get_metric
 from tdamapper.utils.vptree_flat.vptree import VPTree as FVPT
 from tdamapper.utils.vptree_hier.vptree import VPTree as HVPT
 from tests.ball_tree import SkBallTree
+from tests.test_utils import (
+    dataset_grid,
+    dataset_random,
+    dataset_simple,
+    dataset_two_lines,
+)
+
+
+def distance(metric):
+    """
+    Get the distance function for the specified metric.
+    """
+    return get_metric(metric)
+
 
-distance = get_metric("euclidean")
+def distance_refs(metric, data):
+    """
+    Get the distance function for the specified metric, using data references.
+    This is useful for testing with datasets that are not numpy arrays.
+    """
+    d = get_metric(metric)
 
+    def dist_refs(i, j):
+        return d(data[i, :], data[j, :])
+
+    return dist_refs
 
-def dataset(dim=10, num=1000):
-    return [np.random.rand(dim) for _ in range(num)]
 
+SIMPLE = dataset_simple()
+SIMPLE_REFS = np.array(list(range(len(SIMPLE))))
 
-eps = 0.25
+TWO_LINES = dataset_two_lines(100)
+TWO_LINES_REFS = np.array(list(range(len(TWO_LINES))))
 
-neighbors = 5
+GRID = dataset_grid(10)
+GRID_REFS = np.array(list(range(len(GRID))))
 
+RANDOM = dataset_random(2, 100)
+RANDOM_REFS = np.array(list(range(len(RANDOM))))
 
-def _test_ball_search(data, dist, vpt):
+
+def _test_ball_search(data, dist, vpt, eps):
     for _ in range(len(data) // 10):
         point = random.choice(data)
         ball = vpt.ball_search(point, eps)
@@ -31,7 +64,7 @@ def _test_ball_search(data, dist, vpt):
             assert any(d(x, y) == 0.0 for y in ball)
 
 
-def _test_knn_search(data, dist, vpt):
+def _test_knn_search(data, dist, vpt, neighbors):
     for _ in range(len(data) // 10):
         point = random.choice(data)
         neigh = vpt.knn_search(point, neighbors)
@@ -54,65 +87,28 @@ def _test_nn_search(data, dist, vpt):
         assert 0.0 == d(val, neigh[0])
 
 
-def _test_vptree(builder, data, dist):
-    vpt = builder(data, metric=dist, leaf_radius=eps, leaf_capacity=neighbors)
-    _test_ball_search(data, dist, vpt)
-    _test_knn_search(data, dist, vpt)
-    _test_nn_search(data, dist, vpt)
-    vpt = builder(
-        data,
-        metric=dist,
-        leaf_radius=eps,
-        leaf_capacity=neighbors,
-        pivoting="random",
-    )
-    _test_ball_search(data, dist, vpt)
-    _test_knn_search(data, dist, vpt)
-    _test_nn_search(data, dist, vpt)
+def _test_vptree(builder, data, dist, eps, neighbors, pivoting):
     vpt = builder(
         data,
         metric=dist,
         leaf_radius=eps,
         leaf_capacity=neighbors,
-        pivoting="furthest",
+        pivoting=pivoting,
     )
-    _test_ball_search(data, dist, vpt)
-    _test_knn_search(data, dist, vpt)
+    _test_ball_search(data, dist, vpt, eps)
+    _test_knn_search(data, dist, vpt, neighbors)
     _test_nn_search(data, dist, vpt)
 
 
-def test_vptree_hier_refs():
-    data = dataset()
-    data_refs = list(range(len(data)))
-    d = get_metric(distance)
-
-    def dist_refs(i, j):
-        return d(data[i], data[j])
-
-    _test_vptree(HVPT, data_refs, dist_refs)
-
-
-def test_vptree_hier_data():
-    data = dataset()
-    _test_vptree(HVPT, data, distance)
-
-
-def test_vptree_flat_refs():
-    data = dataset()
-    data_refs = list(range(len(data)))
-    d = get_metric(distance)
-
-    def dist_refs(i, j):
-        return d(data[i], data[j])
-
-    _test_vptree(FVPT, data_refs, dist_refs)
-
-
-def test_vptree_flat_data():
-    data = dataset()
-    _test_vptree(FVPT, data, distance)
-
-
-def test_ball_tree_data():
-    data = dataset()
-    _test_vptree(SkBallTree, data, distance)
+@pytest.mark.parametrize("pivoting", ["disabled", "random", "furthest"])
+@pytest.mark.parametrize("eps", [0.1, 0.25, 0.5])
+@pytest.mark.parametrize("neighbors", [1, 5, 10])
+@pytest.mark.parametrize("builder", [HVPT, FVPT])
+@pytest.mark.parametrize("metric", ["euclidean"])
+@pytest.mark.parametrize("dataset", [SIMPLE, TWO_LINES, GRID, RANDOM])
+def test_vptree(builder, dataset, metric, eps, neighbors, pivoting):
+    """
+    Test the vp-tree implementations with various datasets and metrics.
+    """
+    metric = get_metric(metric)
+    _test_vptree(builder, dataset, metric, eps, neighbors, pivoting)

tests/test_utils.py

@@ -0,0 +1,51 @@
+"""
+Test utilities.
+"""
+
+import numpy as np
+
+
+def dataset_simple():
+    """
+    Create a simple dataset of points in a 2D space.
+
+    This dataset consists of four points forming the corners of a rectangle
+    such that two sides are longer than the other two.
+    """
+    return np.array(
+        [
+            [0.0, 1.0],
+            [1.1, 0.0],
+            [0.0, 0.0],
+            [1.1, 1.0],
+        ]
+    )
+
+
+def dataset_random(dim=1, num=1000):
+    """
+    Create a random dataset of points in the unit square.
+    """
+    return np.array([np.random.rand(dim) for _ in range(num)])
+
+
+def dataset_two_lines(num=1000):
+    """
+    Create a dataset consisting of two lines in the unit square.
+    One line is horizontal at y=0, the other is vertical at x=1.
+    """
+    t = np.linspace(0.0, 1.0, num)
+    line1 = np.array([[x, 0.0] for x in t])
+    line2 = np.array([[x, 1.0] for x in t])
+    return np.concatenate((line1, line2), axis=0)
+
+
+def dataset_grid(num=1000):
+    """
+    Create a grid dataset in the unit square.
+    The grid consists of points evenly spaced in both dimensions.
+    """
+    t = np.linspace(0.0, 1.0, num)
+    s = np.linspace(0.0, 1.0, num)
+    grid = np.array([[x, y] for x in t for y in s])
+    return grid