| |
| |
|
|
| import os |
| from numpy.testing import (assert_equal, assert_array_equal, assert_, |
| assert_almost_equal, assert_array_almost_equal, |
| assert_allclose) |
| from pytest import raises as assert_raises |
| import pytest |
| from platform import python_implementation |
| import numpy as np |
| from scipy.spatial import KDTree, Rectangle, distance_matrix, cKDTree |
| from scipy.spatial._ckdtree import cKDTreeNode |
| from scipy.spatial import minkowski_distance |
|
|
| import itertools |
|
|
| @pytest.fixture(params=[KDTree, cKDTree]) |
| def kdtree_type(request): |
| return request.param |
|
|
|
|
| def KDTreeTest(kls): |
| """Class decorator to create test cases for KDTree and cKDTree |
| |
| Tests use the class variable ``kdtree_type`` as the tree constructor. |
| """ |
| if not kls.__name__.startswith('_Test'): |
| raise RuntimeError("Expected a class name starting with _Test") |
|
|
| for tree in (KDTree, cKDTree): |
| test_name = kls.__name__[1:] + '_' + tree.__name__ |
|
|
| if test_name in globals(): |
| raise RuntimeError("Duplicated test name: " + test_name) |
|
|
| |
| test_case = type(test_name, (kls,), {'kdtree_type': tree}) |
| globals()[test_name] = test_case |
| return kls |
|
|
|
|
| def distance_box(a, b, p, boxsize): |
| diff = a - b |
| diff[diff > 0.5 * boxsize] -= boxsize |
| diff[diff < -0.5 * boxsize] += boxsize |
| d = minkowski_distance(diff, 0, p) |
| return d |
|
|
| class ConsistencyTests: |
| def distance(self, a, b, p): |
| return minkowski_distance(a, b, p) |
|
|
| def test_nearest(self): |
| x = self.x |
| d, i = self.kdtree.query(x, 1) |
| assert_almost_equal(d**2, np.sum((x-self.data[i])**2)) |
| eps = 1e-8 |
| assert_(np.all(np.sum((self.data-x[np.newaxis, :])**2, axis=1) > d**2-eps)) |
|
|
| def test_m_nearest(self): |
| x = self.x |
| m = self.m |
| dd, ii = self.kdtree.query(x, m) |
| d = np.amax(dd) |
| i = ii[np.argmax(dd)] |
| assert_almost_equal(d**2, np.sum((x-self.data[i])**2)) |
| eps = 1e-8 |
| assert_equal( |
| np.sum(np.sum((self.data-x[np.newaxis, :])**2, axis=1) < d**2+eps), |
| m, |
| ) |
|
|
| def test_points_near(self): |
| x = self.x |
| d = self.d |
| dd, ii = self.kdtree.query(x, k=self.kdtree.n, distance_upper_bound=d) |
| eps = 1e-8 |
| hits = 0 |
| for near_d, near_i in zip(dd, ii): |
| if near_d == np.inf: |
| continue |
| hits += 1 |
| assert_almost_equal(near_d**2, np.sum((x-self.data[near_i])**2)) |
| assert_(near_d < d+eps, f"near_d={near_d:g} should be less than {d:g}") |
| assert_equal(np.sum(self.distance(self.data, x, 2) < d**2+eps), hits) |
|
|
| def test_points_near_l1(self): |
| x = self.x |
| d = self.d |
| dd, ii = self.kdtree.query(x, k=self.kdtree.n, p=1, distance_upper_bound=d) |
| eps = 1e-8 |
| hits = 0 |
| for near_d, near_i in zip(dd, ii): |
| if near_d == np.inf: |
| continue |
| hits += 1 |
| assert_almost_equal(near_d, self.distance(x, self.data[near_i], 1)) |
| assert_(near_d < d+eps, f"near_d={near_d:g} should be less than {d:g}") |
| assert_equal(np.sum(self.distance(self.data, x, 1) < d+eps), hits) |
|
|
| def test_points_near_linf(self): |
| x = self.x |
| d = self.d |
| dd, ii = self.kdtree.query(x, k=self.kdtree.n, p=np.inf, distance_upper_bound=d) |
| eps = 1e-8 |
| hits = 0 |
| for near_d, near_i in zip(dd, ii): |
| if near_d == np.inf: |
| continue |
| hits += 1 |
| assert_almost_equal(near_d, self.distance(x, self.data[near_i], np.inf)) |
| assert_(near_d < d+eps, f"near_d={near_d:g} should be less than {d:g}") |
| assert_equal(np.sum(self.distance(self.data, x, np.inf) < d+eps), hits) |
|
|
| def test_approx(self): |
| x = self.x |
| k = self.k |
| eps = 0.1 |
| d_real, i_real = self.kdtree.query(x, k) |
| d, i = self.kdtree.query(x, k, eps=eps) |
| assert_(np.all(d <= d_real*(1+eps))) |
|
|
|
|
| @KDTreeTest |
| class _Test_random(ConsistencyTests): |
| def setup_method(self): |
| self.n = 100 |
| self.m = 4 |
| np.random.seed(1234) |
| self.data = np.random.randn(self.n, self.m) |
| self.kdtree = self.kdtree_type(self.data, leafsize=2) |
| self.x = np.random.randn(self.m) |
| self.d = 0.2 |
| self.k = 10 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_far(_Test_random): |
| def setup_method(self): |
| super().setup_method() |
| self.x = np.random.randn(self.m)+10 |
|
|
|
|
| @KDTreeTest |
| class _Test_small(ConsistencyTests): |
| def setup_method(self): |
| self.data = np.array([[0, 0, 0], |
| [0, 0, 1], |
| [0, 1, 0], |
| [0, 1, 1], |
| [1, 0, 0], |
| [1, 0, 1], |
| [1, 1, 0], |
| [1, 1, 1]]) |
| self.kdtree = self.kdtree_type(self.data) |
| self.n = self.kdtree.n |
| self.m = self.kdtree.m |
| np.random.seed(1234) |
| self.x = np.random.randn(3) |
| self.d = 0.5 |
| self.k = 4 |
|
|
| def test_nearest(self): |
| assert_array_equal( |
| self.kdtree.query((0, 0, 0.1), 1), |
| (0.1, 0)) |
|
|
| def test_nearest_two(self): |
| assert_array_equal( |
| self.kdtree.query((0, 0, 0.1), 2), |
| ([0.1, 0.9], [0, 1])) |
|
|
|
|
| @KDTreeTest |
| class _Test_small_nonleaf(_Test_small): |
| def setup_method(self): |
| super().setup_method() |
| self.kdtree = self.kdtree_type(self.data, leafsize=1) |
|
|
|
|
| class Test_vectorization_KDTree: |
| def setup_method(self): |
| self.data = np.array([[0, 0, 0], |
| [0, 0, 1], |
| [0, 1, 0], |
| [0, 1, 1], |
| [1, 0, 0], |
| [1, 0, 1], |
| [1, 1, 0], |
| [1, 1, 1]]) |
| self.kdtree = KDTree(self.data) |
|
|
| def test_single_query(self): |
| d, i = self.kdtree.query(np.array([0, 0, 0])) |
| assert_(isinstance(d, float)) |
| assert_(np.issubdtype(i, np.signedinteger)) |
|
|
| def test_vectorized_query(self): |
| d, i = self.kdtree.query(np.zeros((2, 4, 3))) |
| assert_equal(np.shape(d), (2, 4)) |
| assert_equal(np.shape(i), (2, 4)) |
|
|
| def test_single_query_multiple_neighbors(self): |
| s = 23 |
| kk = self.kdtree.n+s |
| d, i = self.kdtree.query(np.array([0, 0, 0]), k=kk) |
| assert_equal(np.shape(d), (kk,)) |
| assert_equal(np.shape(i), (kk,)) |
| assert_(np.all(~np.isfinite(d[-s:]))) |
| assert_(np.all(i[-s:] == self.kdtree.n)) |
|
|
| def test_vectorized_query_multiple_neighbors(self): |
| s = 23 |
| kk = self.kdtree.n+s |
| d, i = self.kdtree.query(np.zeros((2, 4, 3)), k=kk) |
| assert_equal(np.shape(d), (2, 4, kk)) |
| assert_equal(np.shape(i), (2, 4, kk)) |
| assert_(np.all(~np.isfinite(d[:, :, -s:]))) |
| assert_(np.all(i[:, :, -s:] == self.kdtree.n)) |
|
|
| def test_query_raises_for_k_none(self): |
| x = 1.0 |
| with pytest.raises(ValueError, match="k must be an integer or*"): |
| self.kdtree.query(x, k=None) |
|
|
| class Test_vectorization_cKDTree: |
| def setup_method(self): |
| self.data = np.array([[0, 0, 0], |
| [0, 0, 1], |
| [0, 1, 0], |
| [0, 1, 1], |
| [1, 0, 0], |
| [1, 0, 1], |
| [1, 1, 0], |
| [1, 1, 1]]) |
| self.kdtree = cKDTree(self.data) |
|
|
| def test_single_query(self): |
| d, i = self.kdtree.query([0, 0, 0]) |
| assert_(isinstance(d, float)) |
| assert_(isinstance(i, int)) |
|
|
| def test_vectorized_query(self): |
| d, i = self.kdtree.query(np.zeros((2, 4, 3))) |
| assert_equal(np.shape(d), (2, 4)) |
| assert_equal(np.shape(i), (2, 4)) |
|
|
| def test_vectorized_query_noncontiguous_values(self): |
| np.random.seed(1234) |
| qs = np.random.randn(3, 1000).T |
| ds, i_s = self.kdtree.query(qs) |
| for q, d, i in zip(qs, ds, i_s): |
| assert_equal(self.kdtree.query(q), (d, i)) |
|
|
| def test_single_query_multiple_neighbors(self): |
| s = 23 |
| kk = self.kdtree.n+s |
| d, i = self.kdtree.query([0, 0, 0], k=kk) |
| assert_equal(np.shape(d), (kk,)) |
| assert_equal(np.shape(i), (kk,)) |
| assert_(np.all(~np.isfinite(d[-s:]))) |
| assert_(np.all(i[-s:] == self.kdtree.n)) |
|
|
| def test_vectorized_query_multiple_neighbors(self): |
| s = 23 |
| kk = self.kdtree.n+s |
| d, i = self.kdtree.query(np.zeros((2, 4, 3)), k=kk) |
| assert_equal(np.shape(d), (2, 4, kk)) |
| assert_equal(np.shape(i), (2, 4, kk)) |
| assert_(np.all(~np.isfinite(d[:, :, -s:]))) |
| assert_(np.all(i[:, :, -s:] == self.kdtree.n)) |
|
|
| class ball_consistency: |
| tol = 0.0 |
|
|
| def distance(self, a, b, p): |
| return minkowski_distance(a * 1.0, b * 1.0, p) |
|
|
| def test_in_ball(self): |
| x = np.atleast_2d(self.x) |
| d = np.broadcast_to(self.d, x.shape[:-1]) |
| l = self.T.query_ball_point(x, self.d, p=self.p, eps=self.eps) |
| for i, ind in enumerate(l): |
| dist = self.distance(self.data[ind], x[i], self.p) - d[i]*(1.+self.eps) |
| norm = self.distance(self.data[ind], x[i], self.p) + d[i]*(1.+self.eps) |
| assert_array_equal(dist < self.tol * norm, True) |
|
|
| def test_found_all(self): |
| x = np.atleast_2d(self.x) |
| d = np.broadcast_to(self.d, x.shape[:-1]) |
| l = self.T.query_ball_point(x, self.d, p=self.p, eps=self.eps) |
| for i, ind in enumerate(l): |
| c = np.ones(self.T.n, dtype=bool) |
| c[ind] = False |
| dist = self.distance(self.data[c], x[i], self.p) - d[i]/(1.+self.eps) |
| norm = self.distance(self.data[c], x[i], self.p) + d[i]/(1.+self.eps) |
| assert_array_equal(dist > -self.tol * norm, True) |
|
|
| @KDTreeTest |
| class _Test_random_ball(ball_consistency): |
| def setup_method(self): |
| n = 100 |
| m = 4 |
| np.random.seed(1234) |
| self.data = np.random.randn(n, m) |
| self.T = self.kdtree_type(self.data, leafsize=2) |
| self.x = np.random.randn(m) |
| self.p = 2. |
| self.eps = 0 |
| self.d = 0.2 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_periodic(ball_consistency): |
| def distance(self, a, b, p): |
| return distance_box(a, b, p, 1.0) |
|
|
| def setup_method(self): |
| n = 10000 |
| m = 4 |
| np.random.seed(1234) |
| self.data = np.random.uniform(size=(n, m)) |
| self.T = self.kdtree_type(self.data, leafsize=2, boxsize=1) |
| self.x = np.full(m, 0.1) |
| self.p = 2. |
| self.eps = 0 |
| self.d = 0.2 |
|
|
| def test_in_ball_outside(self): |
| l = self.T.query_ball_point(self.x + 1.0, self.d, p=self.p, eps=self.eps) |
| for i in l: |
| assert_(self.distance(self.data[i], self.x, self.p) <= self.d*(1.+self.eps)) |
| l = self.T.query_ball_point(self.x - 1.0, self.d, p=self.p, eps=self.eps) |
| for i in l: |
| assert_(self.distance(self.data[i], self.x, self.p) <= self.d*(1.+self.eps)) |
|
|
| def test_found_all_outside(self): |
| c = np.ones(self.T.n, dtype=bool) |
| l = self.T.query_ball_point(self.x + 1.0, self.d, p=self.p, eps=self.eps) |
| c[l] = False |
| assert np.all( |
| self.distance(self.data[c], self.x, self.p) >= self.d/(1.+self.eps) |
| ) |
|
|
| l = self.T.query_ball_point(self.x - 1.0, self.d, p=self.p, eps=self.eps) |
| c[l] = False |
| assert np.all( |
| self.distance(self.data[c], self.x, self.p) >= self.d/(1.+self.eps) |
| ) |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_largep_issue9890(ball_consistency): |
|
|
| |
| tol = 1e-13 |
|
|
| def setup_method(self): |
| n = 1000 |
| m = 2 |
| np.random.seed(123) |
| self.data = np.random.randint(100, 1000, size=(n, m)) |
| self.T = self.kdtree_type(self.data) |
| self.x = self.data |
| self.p = 100 |
| self.eps = 0 |
| self.d = 10 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_approx(_Test_random_ball): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.eps = 0.1 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_approx_periodic(_Test_random_ball): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.eps = 0.1 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_far(_Test_random_ball): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.d = 2. |
|
|
| @KDTreeTest |
| class _Test_random_ball_far_periodic(_Test_random_ball_periodic): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.d = 2. |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_l1(_Test_random_ball): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.p = 1 |
|
|
|
|
| @KDTreeTest |
| class _Test_random_ball_linf(_Test_random_ball): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.p = np.inf |
|
|
|
|
| def test_random_ball_vectorized(kdtree_type): |
| n = 20 |
| m = 5 |
| np.random.seed(1234) |
| T = kdtree_type(np.random.randn(n, m)) |
|
|
| r = T.query_ball_point(np.random.randn(2, 3, m), 1) |
| assert_equal(r.shape, (2, 3)) |
| assert_(isinstance(r[0, 0], list)) |
|
|
|
|
| def test_query_ball_point_multithreading(kdtree_type): |
| np.random.seed(0) |
| n = 5000 |
| k = 2 |
| points = np.random.randn(n, k) |
| T = kdtree_type(points) |
| l1 = T.query_ball_point(points, 0.003, workers=1) |
| l2 = T.query_ball_point(points, 0.003, workers=64) |
| l3 = T.query_ball_point(points, 0.003, workers=-1) |
|
|
| for i in range(n): |
| if l1[i] or l2[i]: |
| assert_array_equal(l1[i], l2[i]) |
|
|
| for i in range(n): |
| if l1[i] or l3[i]: |
| assert_array_equal(l1[i], l3[i]) |
|
|
|
|
| class two_trees_consistency: |
|
|
| def distance(self, a, b, p): |
| return minkowski_distance(a, b, p) |
|
|
| def test_all_in_ball(self): |
| r = self.T1.query_ball_tree(self.T2, self.d, p=self.p, eps=self.eps) |
| for i, l in enumerate(r): |
| for j in l: |
| assert (self.distance(self.data1[i], self.data2[j], self.p) |
| <= self.d*(1.+self.eps)) |
|
|
| def test_found_all(self): |
| r = self.T1.query_ball_tree(self.T2, self.d, p=self.p, eps=self.eps) |
| for i, l in enumerate(r): |
| c = np.ones(self.T2.n, dtype=bool) |
| c[l] = False |
| assert np.all(self.distance(self.data2[c], self.data1[i], self.p) |
| >= self.d/(1.+self.eps)) |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees(two_trees_consistency): |
|
|
| def setup_method(self): |
| n = 50 |
| m = 4 |
| np.random.seed(1234) |
| self.data1 = np.random.randn(n, m) |
| self.T1 = self.kdtree_type(self.data1, leafsize=2) |
| self.data2 = np.random.randn(n, m) |
| self.T2 = self.kdtree_type(self.data2, leafsize=2) |
| self.p = 2. |
| self.eps = 0 |
| self.d = 0.2 |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees_periodic(two_trees_consistency): |
| def distance(self, a, b, p): |
| return distance_box(a, b, p, 1.0) |
|
|
| def setup_method(self): |
| n = 50 |
| m = 4 |
| np.random.seed(1234) |
| self.data1 = np.random.uniform(size=(n, m)) |
| self.T1 = self.kdtree_type(self.data1, leafsize=2, boxsize=1.0) |
| self.data2 = np.random.uniform(size=(n, m)) |
| self.T2 = self.kdtree_type(self.data2, leafsize=2, boxsize=1.0) |
| self.p = 2. |
| self.eps = 0 |
| self.d = 0.2 |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees_far(_Test_two_random_trees): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.d = 2 |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees_far_periodic(_Test_two_random_trees_periodic): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.d = 2 |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees_linf(_Test_two_random_trees): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.p = np.inf |
|
|
|
|
| @KDTreeTest |
| class _Test_two_random_trees_linf_periodic(_Test_two_random_trees_periodic): |
|
|
| def setup_method(self): |
| super().setup_method() |
| self.p = np.inf |
|
|
|
|
| class Test_rectangle: |
|
|
| def setup_method(self): |
| self.rect = Rectangle([0, 0], [1, 1]) |
|
|
| def test_min_inside(self): |
| assert_almost_equal(self.rect.min_distance_point([0.5, 0.5]), 0) |
|
|
| def test_min_one_side(self): |
| assert_almost_equal(self.rect.min_distance_point([0.5, 1.5]), 0.5) |
|
|
| def test_min_two_sides(self): |
| assert_almost_equal(self.rect.min_distance_point([2, 2]), np.sqrt(2)) |
|
|
| def test_max_inside(self): |
| assert_almost_equal(self.rect.max_distance_point([0.5, 0.5]), 1/np.sqrt(2)) |
|
|
| def test_max_one_side(self): |
| assert_almost_equal(self.rect.max_distance_point([0.5, 1.5]), |
| np.hypot(0.5, 1.5)) |
|
|
| def test_max_two_sides(self): |
| assert_almost_equal(self.rect.max_distance_point([2, 2]), 2*np.sqrt(2)) |
|
|
| def test_split(self): |
| less, greater = self.rect.split(0, 0.1) |
| assert_array_equal(less.maxes, [0.1, 1]) |
| assert_array_equal(less.mins, [0, 0]) |
| assert_array_equal(greater.maxes, [1, 1]) |
| assert_array_equal(greater.mins, [0.1, 0]) |
|
|
|
|
| def test_distance_l2(): |
| assert_almost_equal(minkowski_distance([0, 0], [1, 1], 2), np.sqrt(2)) |
|
|
|
|
| def test_distance_l1(): |
| assert_almost_equal(minkowski_distance([0, 0], [1, 1], 1), 2) |
|
|
|
|
| def test_distance_linf(): |
| assert_almost_equal(minkowski_distance([0, 0], [1, 1], np.inf), 1) |
|
|
|
|
| def test_distance_vectorization(): |
| np.random.seed(1234) |
| x = np.random.randn(10, 1, 3) |
| y = np.random.randn(1, 7, 3) |
| assert_equal(minkowski_distance(x, y).shape, (10, 7)) |
|
|
|
|
| class count_neighbors_consistency: |
| def test_one_radius(self): |
| r = 0.2 |
| assert_equal(self.T1.count_neighbors(self.T2, r), |
| np.sum([len(l) for l in self.T1.query_ball_tree(self.T2, r)])) |
|
|
| def test_large_radius(self): |
| r = 1000 |
| assert_equal(self.T1.count_neighbors(self.T2, r), |
| np.sum([len(l) for l in self.T1.query_ball_tree(self.T2, r)])) |
|
|
| def test_multiple_radius(self): |
| rs = np.exp(np.linspace(np.log(0.01), np.log(10), 3)) |
| results = self.T1.count_neighbors(self.T2, rs) |
| assert_(np.all(np.diff(results) >= 0)) |
| for r, result in zip(rs, results): |
| assert_equal(self.T1.count_neighbors(self.T2, r), result) |
|
|
| @KDTreeTest |
| class _Test_count_neighbors(count_neighbors_consistency): |
| def setup_method(self): |
| n = 50 |
| m = 2 |
| np.random.seed(1234) |
| self.T1 = self.kdtree_type(np.random.randn(n, m), leafsize=2) |
| self.T2 = self.kdtree_type(np.random.randn(n, m), leafsize=2) |
|
|
|
|
| class sparse_distance_matrix_consistency: |
|
|
| def distance(self, a, b, p): |
| return minkowski_distance(a, b, p) |
|
|
| def test_consistency_with_neighbors(self): |
| M = self.T1.sparse_distance_matrix(self.T2, self.r) |
| r = self.T1.query_ball_tree(self.T2, self.r) |
| for i, l in enumerate(r): |
| for j in l: |
| assert_almost_equal( |
| M[i, j], |
| self.distance(self.T1.data[i], self.T2.data[j], self.p), |
| decimal=14 |
| ) |
| for ((i, j), d) in M.items(): |
| assert_(j in r[i]) |
|
|
| def test_zero_distance(self): |
| |
| self.T1.sparse_distance_matrix(self.T1, self.r) |
|
|
| def test_consistency(self): |
| |
| M1 = self.T1.sparse_distance_matrix(self.T2, self.r) |
| expected = distance_matrix(self.T1.data, self.T2.data) |
| expected[expected > self.r] = 0 |
| assert_array_almost_equal(M1.toarray(), expected, decimal=14) |
|
|
| def test_against_logic_error_regression(self): |
| |
| np.random.seed(0) |
| too_many = np.array(np.random.randn(18, 2), dtype=int) |
| tree = self.kdtree_type( |
| too_many, balanced_tree=False, compact_nodes=False) |
| d = tree.sparse_distance_matrix(tree, 3).toarray() |
| assert_array_almost_equal(d, d.T, decimal=14) |
|
|
| def test_ckdtree_return_types(self): |
| |
| ref = np.zeros((self.n, self.n)) |
| for i in range(self.n): |
| for j in range(self.n): |
| v = self.data1[i, :] - self.data2[j, :] |
| ref[i, j] = np.dot(v, v) |
| ref = np.sqrt(ref) |
| ref[ref > self.r] = 0. |
| |
| dist = np.zeros((self.n, self.n)) |
| r = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dict') |
| for i, j in r.keys(): |
| dist[i, j] = r[(i, j)] |
| assert_array_almost_equal(ref, dist, decimal=14) |
| |
| dist = np.zeros((self.n, self.n)) |
| r = self.T1.sparse_distance_matrix(self.T2, self.r, |
| output_type='ndarray') |
| for k in range(r.shape[0]): |
| i = r['i'][k] |
| j = r['j'][k] |
| v = r['v'][k] |
| dist[i, j] = v |
| assert_array_almost_equal(ref, dist, decimal=14) |
| |
| r = self.T1.sparse_distance_matrix(self.T2, self.r, |
| output_type='dok_matrix') |
| assert_array_almost_equal(ref, r.toarray(), decimal=14) |
| |
| r = self.T1.sparse_distance_matrix(self.T2, self.r, |
| output_type='coo_matrix') |
| assert_array_almost_equal(ref, r.toarray(), decimal=14) |
|
|
|
|
| @KDTreeTest |
| class _Test_sparse_distance_matrix(sparse_distance_matrix_consistency): |
| def setup_method(self): |
| n = 50 |
| m = 4 |
| np.random.seed(1234) |
| data1 = np.random.randn(n, m) |
| data2 = np.random.randn(n, m) |
| self.T1 = self.kdtree_type(data1, leafsize=2) |
| self.T2 = self.kdtree_type(data2, leafsize=2) |
| self.r = 0.5 |
| self.p = 2 |
| self.data1 = data1 |
| self.data2 = data2 |
| self.n = n |
| self.m = m |
|
|
|
|
| def test_distance_matrix(): |
| m = 10 |
| n = 11 |
| k = 4 |
| np.random.seed(1234) |
| xs = np.random.randn(m, k) |
| ys = np.random.randn(n, k) |
| ds = distance_matrix(xs, ys) |
| assert_equal(ds.shape, (m, n)) |
| for i in range(m): |
| for j in range(n): |
| assert_almost_equal(minkowski_distance(xs[i], ys[j]), ds[i, j]) |
|
|
|
|
| def test_distance_matrix_looping(): |
| m = 10 |
| n = 11 |
| k = 4 |
| np.random.seed(1234) |
| xs = np.random.randn(m, k) |
| ys = np.random.randn(n, k) |
| ds = distance_matrix(xs, ys) |
| dsl = distance_matrix(xs, ys, threshold=1) |
| assert_equal(ds, dsl) |
|
|
|
|
| def check_onetree_query(T, d): |
| r = T.query_ball_tree(T, d) |
| s = set() |
| for i, l in enumerate(r): |
| for j in l: |
| if i < j: |
| s.add((i, j)) |
|
|
| assert_(s == T.query_pairs(d)) |
|
|
| def test_onetree_query(kdtree_type): |
| np.random.seed(0) |
| n = 50 |
| k = 4 |
| points = np.random.randn(n, k) |
| T = kdtree_type(points) |
| check_onetree_query(T, 0.1) |
|
|
| points = np.random.randn(3*n, k) |
| points[:n] *= 0.001 |
| points[n:2*n] += 2 |
| T = kdtree_type(points) |
| check_onetree_query(T, 0.1) |
| check_onetree_query(T, 0.001) |
| check_onetree_query(T, 0.00001) |
| check_onetree_query(T, 1e-6) |
|
|
|
|
| def test_query_pairs_single_node(kdtree_type): |
| tree = kdtree_type([[0, 1]]) |
| assert_equal(tree.query_pairs(0.5), set()) |
|
|
|
|
| def test_kdtree_query_pairs(kdtree_type): |
| np.random.seed(0) |
| n = 50 |
| k = 2 |
| r = 0.1 |
| r2 = r**2 |
| points = np.random.randn(n, k) |
| T = kdtree_type(points) |
| |
| brute = set() |
| for i in range(n): |
| for j in range(i+1, n): |
| v = points[i, :] - points[j, :] |
| if np.dot(v, v) <= r2: |
| brute.add((i, j)) |
| l0 = sorted(brute) |
| |
| s = T.query_pairs(r) |
| l1 = sorted(s) |
| assert_array_equal(l0, l1) |
| |
| s = T.query_pairs(r, output_type='set') |
| l1 = sorted(s) |
| assert_array_equal(l0, l1) |
| |
| s = set() |
| arr = T.query_pairs(r, output_type='ndarray') |
| for i in range(arr.shape[0]): |
| s.add((int(arr[i, 0]), int(arr[i, 1]))) |
| l2 = sorted(s) |
| assert_array_equal(l0, l2) |
|
|
|
|
| def test_query_pairs_eps(kdtree_type): |
| spacing = np.sqrt(2) |
| |
| x_range = np.linspace(0, 3 * spacing, 4) |
| y_range = np.linspace(0, 3 * spacing, 4) |
| xy_array = [(xi, yi) for xi in x_range for yi in y_range] |
| tree = kdtree_type(xy_array) |
| pairs_eps = tree.query_pairs(r=spacing, eps=.1) |
| |
| pairs = tree.query_pairs(r=spacing * 1.01) |
| |
| assert_equal(pairs, pairs_eps) |
|
|
|
|
| def test_ball_point_ints(kdtree_type): |
| |
| x, y = np.mgrid[0:4, 0:4] |
| points = list(zip(x.ravel(), y.ravel())) |
| tree = kdtree_type(points) |
| assert_equal(sorted([4, 8, 9, 12]), |
| sorted(tree.query_ball_point((2, 0), 1))) |
| points = np.asarray(points, dtype=float) |
| tree = kdtree_type(points) |
| assert_equal(sorted([4, 8, 9, 12]), |
| sorted(tree.query_ball_point((2, 0), 1))) |
|
|
|
|
| def test_kdtree_comparisons(): |
| |
| nodes = [KDTree.node() for _ in range(3)] |
| assert_equal(sorted(nodes), sorted(nodes[::-1])) |
|
|
|
|
| def test_kdtree_build_modes(kdtree_type): |
| |
| np.random.seed(0) |
| n = 5000 |
| k = 4 |
| points = np.random.randn(n, k) |
| T1 = kdtree_type(points).query(points, k=5)[-1] |
| T2 = kdtree_type(points, compact_nodes=False).query(points, k=5)[-1] |
| T3 = kdtree_type(points, balanced_tree=False).query(points, k=5)[-1] |
| T4 = kdtree_type(points, compact_nodes=False, |
| balanced_tree=False).query(points, k=5)[-1] |
| assert_array_equal(T1, T2) |
| assert_array_equal(T1, T3) |
| assert_array_equal(T1, T4) |
|
|
| def test_kdtree_pickle(kdtree_type): |
| |
| import pickle |
| np.random.seed(0) |
| n = 50 |
| k = 4 |
| points = np.random.randn(n, k) |
| T1 = kdtree_type(points) |
| tmp = pickle.dumps(T1) |
| T2 = pickle.loads(tmp) |
| T1 = T1.query(points, k=5)[-1] |
| T2 = T2.query(points, k=5)[-1] |
| assert_array_equal(T1, T2) |
|
|
| def test_kdtree_pickle_boxsize(kdtree_type): |
| |
| import pickle |
| np.random.seed(0) |
| n = 50 |
| k = 4 |
| points = np.random.uniform(size=(n, k)) |
| T1 = kdtree_type(points, boxsize=1.0) |
| tmp = pickle.dumps(T1) |
| T2 = pickle.loads(tmp) |
| T1 = T1.query(points, k=5)[-1] |
| T2 = T2.query(points, k=5)[-1] |
| assert_array_equal(T1, T2) |
|
|
| def test_kdtree_copy_data(kdtree_type): |
| |
| |
| |
| np.random.seed(0) |
| n = 5000 |
| k = 4 |
| points = np.random.randn(n, k) |
| T = kdtree_type(points, copy_data=True) |
| q = points.copy() |
| T1 = T.query(q, k=5)[-1] |
| points[...] = np.random.randn(n, k) |
| T2 = T.query(q, k=5)[-1] |
| assert_array_equal(T1, T2) |
|
|
| def test_ckdtree_parallel(kdtree_type, monkeypatch): |
| |
| np.random.seed(0) |
| n = 5000 |
| k = 4 |
| points = np.random.randn(n, k) |
| T = kdtree_type(points) |
| T1 = T.query(points, k=5, workers=64)[-1] |
| T2 = T.query(points, k=5, workers=-1)[-1] |
| T3 = T.query(points, k=5)[-1] |
| assert_array_equal(T1, T2) |
| assert_array_equal(T1, T3) |
|
|
| monkeypatch.setattr(os, 'cpu_count', lambda: None) |
| with pytest.raises(NotImplementedError, match="Cannot determine the"): |
| T.query(points, 1, workers=-1) |
|
|
|
|
| def test_ckdtree_view(): |
| |
| |
| |
| np.random.seed(0) |
| n = 100 |
| k = 4 |
| points = np.random.randn(n, k) |
| kdtree = cKDTree(points) |
|
|
| |
| def recurse_tree(n): |
| assert_(isinstance(n, cKDTreeNode)) |
| if n.split_dim == -1: |
| assert_(n.lesser is None) |
| assert_(n.greater is None) |
| assert_(n.indices.shape[0] <= kdtree.leafsize) |
| else: |
| recurse_tree(n.lesser) |
| recurse_tree(n.greater) |
| x = n.lesser.data_points[:, n.split_dim] |
| y = n.greater.data_points[:, n.split_dim] |
| assert_(x.max() < y.min()) |
|
|
| recurse_tree(kdtree.tree) |
| |
| n = kdtree.tree |
| assert_array_equal(np.sort(n.indices), range(100)) |
| |
| assert_array_equal(kdtree.data[n.indices, :], n.data_points) |
|
|
| |
| |
|
|
| def test_kdtree_list_k(kdtree_type): |
| |
| n = 200 |
| m = 2 |
| klist = [1, 2, 3] |
| kint = 3 |
|
|
| np.random.seed(1234) |
| data = np.random.uniform(size=(n, m)) |
| kdtree = kdtree_type(data, leafsize=1) |
|
|
| |
| dd, ii = kdtree.query(data, klist) |
| dd1, ii1 = kdtree.query(data, kint) |
| assert_equal(dd, dd1) |
| assert_equal(ii, ii1) |
|
|
| |
| klist = np.array([1, 3]) |
| kint = 3 |
| dd, ii = kdtree.query(data, kint) |
| dd1, ii1 = kdtree.query(data, klist) |
| assert_equal(dd1, dd[..., klist - 1]) |
| assert_equal(ii1, ii[..., klist - 1]) |
|
|
| |
| |
| dd, ii = kdtree.query(data, 1) |
| dd1, ii1 = kdtree.query(data, [1]) |
| assert_equal(len(dd.shape), 1) |
| assert_equal(len(dd1.shape), 2) |
| assert_equal(dd, np.ravel(dd1)) |
| assert_equal(ii, np.ravel(ii1)) |
|
|
| @pytest.mark.fail_slow(10) |
| def test_kdtree_box(kdtree_type): |
| |
| n = 2000 |
| m = 3 |
| k = 3 |
| np.random.seed(1234) |
| data = np.random.uniform(size=(n, m)) |
| kdtree = kdtree_type(data, leafsize=1, boxsize=1.0) |
|
|
| |
| kdtree2 = kdtree_type(data, leafsize=1) |
|
|
| for p in [1, 2, 3.0, np.inf]: |
| dd, ii = kdtree.query(data, k, p=p) |
|
|
| dd1, ii1 = kdtree.query(data + 1.0, k, p=p) |
| assert_almost_equal(dd, dd1) |
| assert_equal(ii, ii1) |
|
|
| dd1, ii1 = kdtree.query(data - 1.0, k, p=p) |
| assert_almost_equal(dd, dd1) |
| assert_equal(ii, ii1) |
|
|
| dd2, ii2 = simulate_periodic_box(kdtree2, data, k, boxsize=1.0, p=p) |
| assert_almost_equal(dd, dd2) |
| assert_equal(ii, ii2) |
|
|
| def test_kdtree_box_0boxsize(kdtree_type): |
| |
| n = 2000 |
| m = 2 |
| k = 3 |
| np.random.seed(1234) |
| data = np.random.uniform(size=(n, m)) |
| kdtree = kdtree_type(data, leafsize=1, boxsize=0.0) |
|
|
| |
| kdtree2 = kdtree_type(data, leafsize=1) |
|
|
| for p in [1, 2, np.inf]: |
| dd, ii = kdtree.query(data, k, p=p) |
|
|
| dd1, ii1 = kdtree2.query(data, k, p=p) |
| assert_almost_equal(dd, dd1) |
| assert_equal(ii, ii1) |
|
|
| def test_kdtree_box_upper_bounds(kdtree_type): |
| data = np.linspace(0, 2, 10).reshape(-1, 2) |
| data[:, 1] += 10 |
| with pytest.raises(ValueError): |
| kdtree_type(data, leafsize=1, boxsize=1.0) |
| with pytest.raises(ValueError): |
| kdtree_type(data, leafsize=1, boxsize=(0.0, 2.0)) |
| |
| kdtree_type(data, leafsize=1, boxsize=(2.0, 0.0)) |
|
|
| def test_kdtree_box_lower_bounds(kdtree_type): |
| data = np.linspace(-1, 1, 10) |
| assert_raises(ValueError, kdtree_type, data, leafsize=1, boxsize=1.0) |
|
|
| def simulate_periodic_box(kdtree, data, k, boxsize, p): |
| dd = [] |
| ii = [] |
| x = np.arange(3 ** data.shape[1]) |
| nn = np.array(np.unravel_index(x, [3] * data.shape[1])).T |
| nn = nn - 1.0 |
| for n in nn: |
| image = data + n * 1.0 * boxsize |
| dd2, ii2 = kdtree.query(image, k, p=p) |
| dd2 = dd2.reshape(-1, k) |
| ii2 = ii2.reshape(-1, k) |
| dd.append(dd2) |
| ii.append(ii2) |
| dd = np.concatenate(dd, axis=-1) |
| ii = np.concatenate(ii, axis=-1) |
|
|
| result = np.empty([len(data), len(nn) * k], dtype=[ |
| ('ii', 'i8'), |
| ('dd', 'f8')]) |
| result['ii'][:] = ii |
| result['dd'][:] = dd |
| result.sort(order='dd') |
| return result['dd'][:, :k], result['ii'][:, :k] |
|
|
|
|
| @pytest.mark.skipif(python_implementation() == 'PyPy', |
| reason="Fails on PyPy CI runs. See #9507") |
| def test_ckdtree_memuse(): |
| |
|
|
| |
| |
|
|
| try: |
| import resource |
| except ImportError: |
| |
| return |
| |
| dx, dy = 0.05, 0.05 |
| y, x = np.mgrid[slice(1, 5 + dy, dy), |
| slice(1, 5 + dx, dx)] |
| z = np.sin(x)**10 + np.cos(10 + y*x) * np.cos(x) |
| z_copy = np.empty_like(z) |
| z_copy[:] = z |
| |
| FILLVAL = 99. |
| mask = np.random.randint(0, z.size, np.random.randint(50) + 5) |
| z_copy.flat[mask] = FILLVAL |
| igood = np.vstack(np.nonzero(x != FILLVAL)).T |
| ibad = np.vstack(np.nonzero(x == FILLVAL)).T |
| mem_use = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| |
| for i in range(10): |
| tree = cKDTree(igood) |
| |
| num_leaks = 0 |
| for i in range(100): |
| mem_use = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| tree = cKDTree(igood) |
| dist, iquery = tree.query(ibad, k=4, p=2) |
| new_mem_use = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| if new_mem_use > mem_use: |
| num_leaks += 1 |
| |
| |
| assert_(num_leaks < 10) |
|
|
| def test_kdtree_weights(kdtree_type): |
|
|
| data = np.linspace(0, 1, 4).reshape(-1, 1) |
| tree1 = kdtree_type(data, leafsize=1) |
| weights = np.ones(len(data), dtype='f4') |
|
|
| nw = tree1._build_weights(weights) |
| assert_array_equal(nw, [4, 2, 1, 1, 2, 1, 1]) |
|
|
| assert_raises(ValueError, tree1._build_weights, weights[:-1]) |
|
|
| for i in range(10): |
| |
| c1 = tree1.count_neighbors(tree1, np.linspace(0, 10, i)) |
| c2 = tree1.count_neighbors(tree1, np.linspace(0, 10, i), |
| weights=(weights, weights)) |
| c3 = tree1.count_neighbors(tree1, np.linspace(0, 10, i), |
| weights=(weights, None)) |
| c4 = tree1.count_neighbors(tree1, np.linspace(0, 10, i), |
| weights=(None, weights)) |
| tree1.count_neighbors(tree1, np.linspace(0, 10, i), |
| weights=weights) |
|
|
| assert_array_equal(c1, c2) |
| assert_array_equal(c1, c3) |
| assert_array_equal(c1, c4) |
|
|
| for i in range(len(data)): |
| |
| w1 = weights.copy() |
| w1[i] = 0 |
| data2 = data[w1 != 0] |
| tree2 = kdtree_type(data2) |
|
|
| c1 = tree1.count_neighbors(tree1, np.linspace(0, 10, 100), |
| weights=(w1, w1)) |
| |
| c2 = tree2.count_neighbors(tree2, np.linspace(0, 10, 100)) |
|
|
| assert_array_equal(c1, c2) |
|
|
| |
| |
| assert_raises(ValueError, tree1.count_neighbors, |
| tree2, np.linspace(0, 10, 100), weights=w1) |
|
|
| @pytest.mark.fail_slow(10) |
| def test_kdtree_count_neighbous_multiple_r(kdtree_type): |
| n = 2000 |
| m = 2 |
| np.random.seed(1234) |
| data = np.random.normal(size=(n, m)) |
| kdtree = kdtree_type(data, leafsize=1) |
| r0 = [0, 0.01, 0.01, 0.02, 0.05] |
| i0 = np.arange(len(r0)) |
| n0 = kdtree.count_neighbors(kdtree, r0) |
| nnc = kdtree.count_neighbors(kdtree, r0, cumulative=False) |
| assert_equal(n0, nnc.cumsum()) |
|
|
| for i, r in zip(itertools.permutations(i0), |
| itertools.permutations(r0)): |
| |
| n = kdtree.count_neighbors(kdtree, r) |
| assert_array_equal(n, n0[list(i)]) |
|
|
| def test_len0_arrays(kdtree_type): |
| |
| |
| rng = np.random.RandomState(1234) |
| X = rng.rand(10, 2) |
| Y = rng.rand(10, 2) |
| tree = kdtree_type(X) |
| |
| d, i = tree.query([.5, .5], k=1) |
| z = tree.query_ball_point([.5, .5], 0.1*d) |
| assert_array_equal(z, []) |
| |
| d, i = tree.query(Y, k=1) |
| mind = d.min() |
| z = tree.query_ball_point(Y, 0.1*mind) |
| y = np.empty(shape=(10, ), dtype=object) |
| y.fill([]) |
| assert_array_equal(y, z) |
| |
| other = kdtree_type(Y) |
| y = tree.query_ball_tree(other, 0.1*mind) |
| assert_array_equal(10*[[]], y) |
| |
| y = tree.count_neighbors(other, 0.1*mind) |
| assert_(y == 0) |
| |
| y = tree.sparse_distance_matrix(other, 0.1*mind, output_type='dok_matrix') |
| assert_array_equal(y == np.zeros((10, 10)), True) |
| y = tree.sparse_distance_matrix(other, 0.1*mind, output_type='coo_matrix') |
| assert_array_equal(y == np.zeros((10, 10)), True) |
| y = tree.sparse_distance_matrix(other, 0.1*mind, output_type='dict') |
| assert_equal(y, {}) |
| y = tree.sparse_distance_matrix(other, 0.1*mind, output_type='ndarray') |
| _dtype = [('i', np.intp), ('j', np.intp), ('v', np.float64)] |
| res_dtype = np.dtype(_dtype, align=True) |
| z = np.empty(shape=(0, ), dtype=res_dtype) |
| assert_array_equal(y, z) |
| |
| d, i = tree.query(X, k=2) |
| mind = d[:, -1].min() |
| y = tree.query_pairs(0.1*mind, output_type='set') |
| assert_equal(y, set()) |
| y = tree.query_pairs(0.1*mind, output_type='ndarray') |
| z = np.empty(shape=(0, 2), dtype=np.intp) |
| assert_array_equal(y, z) |
|
|
| def test_kdtree_duplicated_inputs(kdtree_type): |
| |
| n = 1024 |
| for m in range(1, 8): |
| data = np.ones((n, m)) |
| data[n//2:] = 2 |
|
|
| for balanced, compact in itertools.product((False, True), repeat=2): |
| kdtree = kdtree_type(data, balanced_tree=balanced, |
| compact_nodes=compact, leafsize=1) |
| assert kdtree.size == 3 |
|
|
| tree = (kdtree.tree if kdtree_type is cKDTree else |
| kdtree.tree._node) |
|
|
| assert_equal( |
| np.sort(tree.lesser.indices), |
| np.arange(0, n // 2)) |
| assert_equal( |
| np.sort(tree.greater.indices), |
| np.arange(n // 2, n)) |
|
|
|
|
| def test_kdtree_noncumulative_nondecreasing(kdtree_type): |
| |
|
|
| |
| |
| kdtree = kdtree_type([[0]], leafsize=1) |
|
|
| assert_raises(ValueError, kdtree.count_neighbors, |
| kdtree, [0.1, 0], cumulative=False) |
|
|
| def test_short_knn(kdtree_type): |
|
|
| |
|
|
| xyz = np.array([ |
| [0., 0., 0.], |
| [1.01, 0., 0.], |
| [0., 1., 0.], |
| [0., 1.01, 0.], |
| [1., 0., 0.], |
| [1., 1., 0.]], |
| dtype='float64') |
|
|
| ckdt = kdtree_type(xyz) |
|
|
| deq, ieq = ckdt.query(xyz, k=4, distance_upper_bound=0.2) |
|
|
| assert_array_almost_equal(deq, |
| [[0., np.inf, np.inf, np.inf], |
| [0., 0.01, np.inf, np.inf], |
| [0., 0.01, np.inf, np.inf], |
| [0., 0.01, np.inf, np.inf], |
| [0., 0.01, np.inf, np.inf], |
| [0., np.inf, np.inf, np.inf]]) |
|
|
| def test_query_ball_point_vector_r(kdtree_type): |
|
|
| np.random.seed(1234) |
| data = np.random.normal(size=(100, 3)) |
| query = np.random.normal(size=(100, 3)) |
| tree = kdtree_type(data) |
| d = np.random.uniform(0, 0.3, size=len(query)) |
|
|
| rvector = tree.query_ball_point(query, d) |
| rscalar = [tree.query_ball_point(qi, di) for qi, di in zip(query, d)] |
| for a, b in zip(rvector, rscalar): |
| assert_array_equal(sorted(a), sorted(b)) |
|
|
| def test_query_ball_point_length(kdtree_type): |
|
|
| np.random.seed(1234) |
| data = np.random.normal(size=(100, 3)) |
| query = np.random.normal(size=(100, 3)) |
| tree = kdtree_type(data) |
| d = 0.3 |
|
|
| length = tree.query_ball_point(query, d, return_length=True) |
| length2 = [len(ind) for ind in tree.query_ball_point(query, d, return_length=False)] |
| length3 = [len(tree.query_ball_point(qi, d)) for qi in query] |
| length4 = [tree.query_ball_point(qi, d, return_length=True) for qi in query] |
| assert_array_equal(length, length2) |
| assert_array_equal(length, length3) |
| assert_array_equal(length, length4) |
|
|
| def test_discontiguous(kdtree_type): |
|
|
| np.random.seed(1234) |
| data = np.random.normal(size=(100, 3)) |
| d_contiguous = np.arange(100) * 0.04 |
| d_discontiguous = np.ascontiguousarray( |
| np.arange(100)[::-1] * 0.04)[::-1] |
| query_contiguous = np.random.normal(size=(100, 3)) |
| query_discontiguous = np.ascontiguousarray(query_contiguous.T).T |
| assert query_discontiguous.strides[-1] != query_contiguous.strides[-1] |
| assert d_discontiguous.strides[-1] != d_contiguous.strides[-1] |
|
|
| tree = kdtree_type(data) |
|
|
| length1 = tree.query_ball_point(query_contiguous, |
| d_contiguous, return_length=True) |
| length2 = tree.query_ball_point(query_discontiguous, |
| d_discontiguous, return_length=True) |
|
|
| assert_array_equal(length1, length2) |
|
|
| d1, i1 = tree.query(query_contiguous, 1) |
| d2, i2 = tree.query(query_discontiguous, 1) |
|
|
| assert_array_equal(d1, d2) |
| assert_array_equal(i1, i2) |
|
|
|
|
| @pytest.mark.parametrize("balanced_tree, compact_nodes", |
| [(True, False), |
| (True, True), |
| (False, False), |
| (False, True)]) |
| def test_kdtree_empty_input(kdtree_type, balanced_tree, compact_nodes): |
| |
| np.random.seed(1234) |
| empty_v3 = np.empty(shape=(0, 3)) |
| query_v3 = np.ones(shape=(1, 3)) |
| query_v2 = np.ones(shape=(2, 3)) |
|
|
| tree = kdtree_type(empty_v3, balanced_tree=balanced_tree, |
| compact_nodes=compact_nodes) |
| length = tree.query_ball_point(query_v3, 0.3, return_length=True) |
| assert length == 0 |
|
|
| dd, ii = tree.query(query_v2, 2) |
| assert ii.shape == (2, 2) |
| assert dd.shape == (2, 2) |
| assert np.isinf(dd).all() |
|
|
| N = tree.count_neighbors(tree, [0, 1]) |
| assert_array_equal(N, [0, 0]) |
|
|
| M = tree.sparse_distance_matrix(tree, 0.3) |
| assert M.shape == (0, 0) |
|
|
| @KDTreeTest |
| class _Test_sorted_query_ball_point: |
| def setup_method(self): |
| np.random.seed(1234) |
| self.x = np.random.randn(100, 1) |
| self.ckdt = self.kdtree_type(self.x) |
|
|
| def test_return_sorted_True(self): |
| idxs_list = self.ckdt.query_ball_point(self.x, 1., return_sorted=True) |
| for idxs in idxs_list: |
| assert_array_equal(idxs, sorted(idxs)) |
|
|
| for xi in self.x: |
| idxs = self.ckdt.query_ball_point(xi, 1., return_sorted=True) |
| assert_array_equal(idxs, sorted(idxs)) |
|
|
| def test_return_sorted_None(self): |
| """Previous behavior was to sort the returned indices if there were |
| multiple points per query but not sort them if there was a single point |
| per query.""" |
| idxs_list = self.ckdt.query_ball_point(self.x, 1.) |
| for idxs in idxs_list: |
| assert_array_equal(idxs, sorted(idxs)) |
|
|
| idxs_list_single = [self.ckdt.query_ball_point(xi, 1.) for xi in self.x] |
| idxs_list_False = self.ckdt.query_ball_point(self.x, 1., return_sorted=False) |
| for idxs0, idxs1 in zip(idxs_list_False, idxs_list_single): |
| assert_array_equal(idxs0, idxs1) |
|
|
|
|
| def test_kdtree_complex_data(): |
| |
| points = np.random.rand(10, 2).view(complex) |
|
|
| with pytest.raises(TypeError, match="complex data"): |
| t = KDTree(points) |
|
|
| t = KDTree(points.real) |
|
|
| with pytest.raises(TypeError, match="complex data"): |
| t.query(points) |
|
|
| with pytest.raises(TypeError, match="complex data"): |
| t.query_ball_point(points, r=1) |
|
|
|
|
| def test_kdtree_tree_access(): |
| |
| np.random.seed(1234) |
| points = np.random.rand(100, 4) |
| t = KDTree(points) |
| root = t.tree |
|
|
| assert isinstance(root, KDTree.innernode) |
| assert root.children == points.shape[0] |
|
|
| |
| nodes = [root] |
| while nodes: |
| n = nodes.pop(-1) |
|
|
| if isinstance(n, KDTree.leafnode): |
| assert isinstance(n.children, int) |
| assert n.children == len(n.idx) |
| assert_array_equal(points[n.idx], n._node.data_points) |
| else: |
| assert isinstance(n, KDTree.innernode) |
| assert isinstance(n.split_dim, int) |
| assert 0 <= n.split_dim < t.m |
| assert isinstance(n.split, float) |
| assert isinstance(n.children, int) |
| assert n.children == n.less.children + n.greater.children |
| nodes.append(n.greater) |
| nodes.append(n.less) |
|
|
|
|
| def test_kdtree_attributes(): |
| |
| np.random.seed(1234) |
| points = np.random.rand(100, 4) |
| t = KDTree(points) |
|
|
| assert isinstance(t.m, int) |
| assert t.n == points.shape[0] |
|
|
| assert isinstance(t.n, int) |
| assert t.m == points.shape[1] |
|
|
| assert isinstance(t.leafsize, int) |
| assert t.leafsize == 10 |
|
|
| assert_array_equal(t.maxes, np.amax(points, axis=0)) |
| assert_array_equal(t.mins, np.amin(points, axis=0)) |
| assert t.data is points |
|
|
|
|
| @pytest.mark.parametrize("kdtree_class", [KDTree, cKDTree]) |
| def test_kdtree_count_neighbors_weighted(kdtree_class): |
| rng = np.random.RandomState(1234) |
| r = np.arange(0.05, 1, 0.05) |
|
|
| A = rng.random(21).reshape((7,3)) |
| B = rng.random(45).reshape((15,3)) |
|
|
| wA = rng.random(7) |
| wB = rng.random(15) |
|
|
| kdA = kdtree_class(A) |
| kdB = kdtree_class(B) |
|
|
| nAB = kdA.count_neighbors(kdB, r, cumulative=False, weights=(wA,wB)) |
|
|
| |
| weights = wA[None, :] * wB[:, None] |
| dist = np.linalg.norm(A[None, :, :] - B[:, None, :], axis=-1) |
| expect = [np.sum(weights[(prev_radius < dist) & (dist <= radius)]) |
| for prev_radius, radius in zip(itertools.chain([0], r[:-1]), r)] |
| assert_allclose(nAB, expect) |
|
|
|
|
| def test_kdtree_nan(): |
| vals = [1, 5, -10, 7, -4, -16, -6, 6, 3, -11] |
| n = len(vals) |
| data = np.concatenate([vals, np.full(n, np.nan)])[:, None] |
| with pytest.raises(ValueError, match="must be finite"): |
| KDTree(data) |
|
|
|
|
| def test_nonfinite_inputs_gh_18223(): |
| rng = np.random.default_rng(12345) |
| coords = rng.uniform(size=(100, 3), low=0.0, high=0.1) |
| t = KDTree(coords, balanced_tree=False, compact_nodes=False) |
| bad_coord = [np.nan for _ in range(3)] |
|
|
| with pytest.raises(ValueError, match="must be finite"): |
| t.query(bad_coord) |
| with pytest.raises(ValueError, match="must be finite"): |
| t.query_ball_point(bad_coord, 1) |
|
|
| coords[0, :] = np.nan |
| with pytest.raises(ValueError, match="must be finite"): |
| KDTree(coords, balanced_tree=True, compact_nodes=False) |
| with pytest.raises(ValueError, match="must be finite"): |
| KDTree(coords, balanced_tree=False, compact_nodes=True) |
| with pytest.raises(ValueError, match="must be finite"): |
| KDTree(coords, balanced_tree=True, compact_nodes=True) |
| with pytest.raises(ValueError, match="must be finite"): |
| KDTree(coords, balanced_tree=False, compact_nodes=False) |
|
|
|
|
| @pytest.mark.parametrize("incantation", [cKDTree, KDTree]) |
| def test_gh_18800(incantation): |
| |
| |
| |
|
|
| class ArrLike(np.ndarray): |
| def __new__(cls, input_array): |
| obj = np.asarray(input_array).view(cls) |
| |
| |
| obj.all = None |
| return obj |
|
|
| def __array_finalize__(self, obj): |
| if obj is None: |
| return |
| self.all = getattr(obj, 'all', None) |
|
|
| points = [ |
| [66.22, 32.54], |
| [22.52, 22.39], |
| [31.01, 81.21], |
| ] |
| arr = np.array(points) |
| arr_like = ArrLike(arr) |
| tree = incantation(points, 10) |
| tree.query(arr_like, 1) |
| tree.query_ball_point(arr_like, 200) |
|
|