# Copyright Anne M. Archibald 2008 # Released under the scipy license from numpy.testing import (assert_equal, assert_array_equal, assert_, assert_almost_equal, assert_array_almost_equal) 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 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, "near_d=%g should be less than %g" % (near_d,d)) 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, "near_d=%g should be less than %g" % (near_d,d)) 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, "near_d=%g should be less than %g" % (near_d,d)) 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))) 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 = KDTree(self.data,leafsize=2) self.x = np.random.randn(self.m) self.d = 0.2 self.k = 10 class Test_random_far(Test_random): def setup_method(self): Test_random.setup_method(self) self.x = np.random.randn(self.m)+10 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 = KDTree(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])) class Test_small_nonleaf(Test_small): def setup_method(self): Test_small.setup_method(self) self.kdtree = KDTree(self.data,leafsize=1) class Test_small_compiled(Test_small): def setup_method(self): Test_small.setup_method(self) self.kdtree = cKDTree(self.data) class Test_small_nonleaf_compiled(Test_small): def setup_method(self): Test_small.setup_method(self) self.kdtree = cKDTree(self.data,leafsize=1) class Test_random_compiled(Test_random): def setup_method(self): Test_random.setup_method(self) self.kdtree = cKDTree(self.data) class Test_random_far_compiled(Test_random_far): def setup_method(self): Test_random_far.setup_method(self) self.kdtree = cKDTree(self.data) class Test_vectorization: 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_single_query_all_neighbors(self): d, i = self.kdtree.query([0,0,0],k=None,distance_upper_bound=1.1) assert_(isinstance(d,list)) assert_(isinstance(i,list)) def test_vectorized_query_all_neighbors(self): d, i = self.kdtree.query(np.zeros((2,4,3)),k=None,distance_upper_bound=1.1) assert_equal(np.shape(d),(2,4)) assert_equal(np.shape(i),(2,4)) assert_(isinstance(d[0,0],list)) assert_(isinstance(i[0,0],list)) class Test_vectorization_compiled: 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) 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 = KDTree(self.data,leafsize=2) self.x = np.random.randn(m) self.p = 2. self.eps = 0 self.d = 0.2 class Test_random_ball_compiled(ball_consistency): def setup_method(self): n = 100 m = 4 np.random.seed(1234) self.data = np.random.randn(n,m) self.T = cKDTree(self.data,leafsize=2) self.x = np.random.randn(m) self.p = 2. self.eps = 0 self.d = 0.2 class Test_random_ball_compiled_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 = cKDTree(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))) class Test_random_ball_compiled_largep_issue9890(ball_consistency): # allow some roundoff errors due to numerical issues 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 = cKDTree(self.data) self.x = self.data self.p = 100 self.eps = 0 self.d = 10 class Test_random_ball_approx(Test_random_ball): def setup_method(self): Test_random_ball.setup_method(self) self.eps = 0.1 class Test_random_ball_approx_compiled(Test_random_ball_compiled): def setup_method(self): Test_random_ball_compiled.setup_method(self) self.eps = 0.1 class Test_random_ball_approx_compiled_periodic(Test_random_ball_compiled_periodic): def setup_method(self): Test_random_ball_compiled_periodic.setup_method(self) self.eps = 0.1 class Test_random_ball_far(Test_random_ball): def setup_method(self): Test_random_ball.setup_method(self) self.d = 2. class Test_random_ball_far_compiled(Test_random_ball_compiled): def setup_method(self): Test_random_ball_compiled.setup_method(self) self.d = 2. class Test_random_ball_far_compiled_periodic(Test_random_ball_compiled_periodic): def setup_method(self): Test_random_ball_compiled_periodic.setup_method(self) self.d = 2. class Test_random_ball_l1(Test_random_ball): def setup_method(self): Test_random_ball.setup_method(self) self.p = 1 class Test_random_ball_l1_compiled(Test_random_ball_compiled): def setup_method(self): Test_random_ball_compiled.setup_method(self) self.p = 1 class Test_random_ball_l1_compiled_periodic(Test_random_ball_compiled_periodic): def setup_method(self): Test_random_ball_compiled_periodic.setup_method(self) self.p = 1 class Test_random_ball_linf(Test_random_ball): def setup_method(self): Test_random_ball.setup_method(self) self.p = np.inf class Test_random_ball_linf_compiled_periodic(Test_random_ball_compiled_periodic): def setup_method(self): Test_random_ball_compiled_periodic.setup_method(self) self.p = np.inf def test_random_ball_vectorized(): n = 20 m = 5 T = KDTree(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_random_ball_vectorized_compiled(): n = 20 m = 5 np.random.seed(1234) T = cKDTree(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(): np.random.seed(0) n = 5000 k = 2 points = np.random.randn(n,k) T = cKDTree(points) l1 = T.query_ball_point(points,0.003,n_jobs=1) l2 = T.query_ball_point(points,0.003,n_jobs=64) l3 = T.query_ball_point(points,0.003,n_jobs=-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))) 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 = KDTree(self.data1,leafsize=2) self.data2 = np.random.randn(n,m) self.T2 = KDTree(self.data2,leafsize=2) self.p = 2. self.eps = 0 self.d = 0.2 class Test_two_random_trees_compiled(two_trees_consistency): def setup_method(self): n = 50 m = 4 np.random.seed(1234) self.data1 = np.random.randn(n,m) self.T1 = cKDTree(self.data1,leafsize=2) self.data2 = np.random.randn(n,m) self.T2 = cKDTree(self.data2,leafsize=2) self.p = 2. self.eps = 0 self.d = 0.2 class Test_two_random_trees_compiled_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 = cKDTree(self.data1,leafsize=2, boxsize=1.0) self.data2 = np.random.uniform(size=(n,m)) self.T2 = cKDTree(self.data2,leafsize=2, boxsize=1.0) self.p = 2. self.eps = 0 self.d = 0.2 class Test_two_random_trees_far(Test_two_random_trees): def setup_method(self): Test_two_random_trees.setup_method(self) self.d = 2 class Test_two_random_trees_far_compiled(Test_two_random_trees_compiled): def setup_method(self): Test_two_random_trees_compiled.setup_method(self) self.d = 2 class Test_two_random_trees_far_compiled_periodic(Test_two_random_trees_compiled_periodic): def setup_method(self): Test_two_random_trees_compiled_periodic.setup_method(self) self.d = 2 class Test_two_random_trees_linf(Test_two_random_trees): def setup_method(self): Test_two_random_trees.setup_method(self) self.p = np.inf class Test_two_random_trees_linf_compiled(Test_two_random_trees_compiled): def setup_method(self): Test_two_random_trees_compiled.setup_method(self) self.p = np.inf class Test_two_random_trees_linf_compiled_periodic(Test_two_random_trees_compiled_periodic): def setup_method(self): Test_two_random_trees_compiled_periodic.setup_method(self) 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) class Test_count_neighbors(count_neighbors_consistency): def setup_method(self): n = 50 m = 2 np.random.seed(1234) self.T1 = KDTree(np.random.randn(n,m),leafsize=2) self.T2 = KDTree(np.random.randn(n,m),leafsize=2) class Test_count_neighbors_compiled(count_neighbors_consistency): def setup_method(self): n = 50 m = 2 np.random.seed(1234) self.T1 = cKDTree(np.random.randn(n,m),leafsize=2) self.T2 = cKDTree(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): # raises an exception for bug 870 (FIXME: Does it?) self.T1.sparse_distance_matrix(self.T1, self.r) 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 = cKDTree(data1,leafsize=2) self.T2 = cKDTree(data2,leafsize=2) self.r = 0.5 self.p = 2 self.data1 = data1 self.data2 = data2 self.n = n self.m = m class Test_sparse_distance_matrix_compiled(sparse_distance_matrix_consistency): def setup_method(self): n = 50 m = 4 np.random.seed(0) data1 = np.random.randn(n,m) data2 = np.random.randn(n,m) self.T1 = cKDTree(data1,leafsize=2) self.T2 = cKDTree(data2,leafsize=2) self.ref_T1 = KDTree(data1, leafsize=2) self.ref_T2 = KDTree(data2, leafsize=2) self.r = 0.5 self.n = n self.m = m self.data1 = data1 self.data2 = data2 self.p = 2 def test_consistency_with_python(self): M1 = self.T1.sparse_distance_matrix(self.T2, self.r) M2 = self.ref_T1.sparse_distance_matrix(self.ref_T2, self.r) assert_array_almost_equal(M1.todense(), M2.todense(), decimal=14) def test_against_logic_error_regression(self): # regression test for gh-5077 logic error np.random.seed(0) too_many = np.array(np.random.randn(18, 2), dtype=int) tree = cKDTree(too_many, balanced_tree=False, compact_nodes=False) d = tree.sparse_distance_matrix(tree, 3).todense() assert_array_almost_equal(d, d.T, decimal=14) def test_ckdtree_return_types(self): # brute-force reference 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. # test return type 'dict' 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) # test return type 'ndarray' 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) # test return type 'dok_matrix' r = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dok_matrix') assert_array_almost_equal(ref, r.todense(), decimal=14) # test return type 'coo_matrix' r = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='coo_matrix') assert_array_almost_equal(ref, r.todense(), decimal=14) 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(): np.random.seed(0) n = 50 k = 4 points = np.random.randn(n,k) T = KDTree(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(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_onetree_query_compiled(): np.random.seed(0) n = 100 k = 4 points = np.random.randn(n,k) T = cKDTree(points) check_onetree_query(T, 0.1) points = np.random.randn(3*n,k) points[:n] *= 0.001 points[n:2*n] += 2 T = cKDTree(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(): tree = KDTree([[0, 1]]) assert_equal(tree.query_pairs(0.5), set()) def test_query_pairs_single_node_compiled(): tree = cKDTree([[0, 1]]) assert_equal(tree.query_pairs(0.5), set()) def test_ckdtree_query_pairs(): np.random.seed(0) n = 50 k = 2 r = 0.1 r2 = r**2 points = np.random.randn(n,k) T = cKDTree(points) # brute force reference 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) # test default return type s = T.query_pairs(r) l1 = sorted(s) assert_array_equal(l0,l1) # test return type 'set' s = T.query_pairs(r, output_type='set') l1 = sorted(s) assert_array_equal(l0,l1) # test return type 'ndarray' 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_ball_point_ints(): # Regression test for #1373. x, y = np.mgrid[0:4, 0:4] points = list(zip(x.ravel(), y.ravel())) tree = KDTree(points) assert_equal(sorted([4, 8, 9, 12]), sorted(tree.query_ball_point((2, 0), 1))) points = np.asarray(points, dtype=float) tree = KDTree(points) assert_equal(sorted([4, 8, 9, 12]), sorted(tree.query_ball_point((2, 0), 1))) def test_kdtree_comparisons(): # Regression test: node comparisons were done wrong in 0.12 w/Py3. nodes = [KDTree.node() for _ in range(3)] assert_equal(sorted(nodes), sorted(nodes[::-1])) def test_ckdtree_build_modes(): # check if different build modes for cKDTree give # similar query results np.random.seed(0) n = 5000 k = 4 points = np.random.randn(n, k) T1 = cKDTree(points).query(points, k=5)[-1] T2 = cKDTree(points, compact_nodes=False).query(points, k=5)[-1] T3 = cKDTree(points, balanced_tree=False).query(points, k=5)[-1] T4 = cKDTree(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_ckdtree_pickle(): # test if it is possible to pickle # a cKDTree try: import cPickle as pickle except ImportError: import pickle np.random.seed(0) n = 50 k = 4 points = np.random.randn(n, k) T1 = cKDTree(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_ckdtree_pickle_boxsize(): # test if it is possible to pickle a periodic # cKDTree try: import cPickle as pickle except ImportError: import pickle np.random.seed(0) n = 50 k = 4 points = np.random.uniform(size=(n, k)) T1 = cKDTree(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_ckdtree_copy_data(): # check if copy_data=True makes the kd-tree # impervious to data corruption by modification of # the data arrray np.random.seed(0) n = 5000 k = 4 points = np.random.randn(n, k) T = cKDTree(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(): # check if parallel=True also generates correct # query results np.random.seed(0) n = 5000 k = 4 points = np.random.randn(n, k) T = cKDTree(points) T1 = T.query(points, k=5, n_jobs=64)[-1] T2 = T.query(points, k=5, n_jobs=-1)[-1] T3 = T.query(points, k=5)[-1] assert_array_equal(T1, T2) assert_array_equal(T1, T3) def test_ckdtree_view(): # Check that the nodes can be correctly viewed from Python. # This test also sanity checks each node in the cKDTree, and # thus verifies the internal structure of the kd-tree. np.random.seed(0) n = 100 k = 4 points = np.random.randn(n, k) kdtree = cKDTree(points) # walk the whole kd-tree and sanity check each node 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) # check that indices are correctly retrieved n = kdtree.tree assert_array_equal(np.sort(n.indices), range(100)) # check that data_points are correctly retrieved assert_array_equal(kdtree.data[n.indices, :], n.data_points) # cKDTree is specialized to type double points, so no need to make # a unit test corresponding to test_ball_point_ints() def test_ckdtree_list_k(): # check ckdtree periodic boundary n = 200 m = 2 klist = [1, 2, 3] kint = 3 np.random.seed(1234) data = np.random.uniform(size=(n, m)) kdtree = cKDTree(data, leafsize=1) # check agreement between arange(1,k+1) and k dd, ii = kdtree.query(data, klist) dd1, ii1 = kdtree.query(data, kint) assert_equal(dd, dd1) assert_equal(ii, ii1) # now check skipping one element 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]) # check k == 1 special case # and k == [1] non-special case 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)) def test_ckdtree_box(): # check ckdtree periodic boundary n = 2000 m = 3 k = 3 np.random.seed(1234) data = np.random.uniform(size=(n, m)) kdtree = cKDTree(data, leafsize=1, boxsize=1.0) # use the standard python KDTree for the simulated periodic box kdtree2 = cKDTree(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_ckdtree_box_0boxsize(): # check ckdtree periodic boundary that mimics non-periodic n = 2000 m = 2 k = 3 np.random.seed(1234) data = np.random.uniform(size=(n, m)) kdtree = cKDTree(data, leafsize=1, boxsize=0.0) # use the standard python KDTree for the simulated periodic box kdtree2 = cKDTree(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_ckdtree_box_upper_bounds(): data = np.linspace(0, 2, 10).reshape(-1, 2) data[:, 1] += 10 assert_raises(ValueError, cKDTree, data, leafsize=1, boxsize=1.0) assert_raises(ValueError, cKDTree, data, leafsize=1, boxsize=(0.0, 2.0)) # skip a dimension. cKDTree(data, leafsize=1, boxsize=(2.0, 0.0)) def test_ckdtree_box_lower_bounds(): data = np.linspace(-1, 1, 10) assert_raises(ValueError, cKDTree, 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(): # unit test adaptation of gh-5630 # NOTE: this will fail when run via valgrind, # because rss is no longer a reliable memory usage indicator. try: import resource except ImportError: # resource is not available on Windows return # Make some data 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 # Place FILLVAL in z_copy at random number of random locations 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 # burn-in for i in range(10): tree = cKDTree(igood) # count memleaks while constructing and querying cKDTree 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 # ideally zero leaks, but errors might accidentally happen # outside cKDTree assert_(num_leaks < 10) def test_ckdtree_weights(): data = np.linspace(0, 1, 4).reshape(-1, 1) tree1 = cKDTree(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): # since weights are uniform, these shall agree: 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)): # this tests removal of one data point by setting weight to 0 w1 = weights.copy() w1[i] = 0 data2 = data[w1 != 0] tree2 = cKDTree(data2) c1 = tree1.count_neighbors(tree1, np.linspace(0, 10, 100), weights=(w1, w1)) # "c2 is correct" c2 = tree2.count_neighbors(tree2, np.linspace(0, 10, 100)) assert_array_equal(c1, c2) #this asserts for two different trees, singular weights # crashes assert_raises(ValueError, tree1.count_neighbors, tree2, np.linspace(0, 10, 100), weights=w1) def test_ckdtree_count_neighbous_multiple_r(): n = 2000 m = 2 np.random.seed(1234) data = np.random.normal(size=(n, m)) kdtree = cKDTree(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)): # permute n0 by i and it shall agree n = kdtree.count_neighbors(kdtree, r) assert_array_equal(n, n0[list(i)]) def test_len0_arrays(): # make sure len-0 arrays are handled correctly # in range queries (gh-5639) np.random.seed(1234) X = np.random.rand(10,2) Y = np.random.rand(10,2) tree = cKDTree(X) # query_ball_point (single) d,i = tree.query([.5, .5], k=1) z = tree.query_ball_point([.5, .5], 0.1*d) assert_array_equal(z, []) # query_ball_point (multiple) 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) # query_ball_tree other = cKDTree(Y) y = tree.query_ball_tree(other, 0.1*mind) assert_array_equal(10*[[]], y) # count_neighbors y = tree.count_neighbors(other, 0.1*mind) assert_(y == 0) # sparse_distance_matrix 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) # query_pairs 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_ckdtree_duplicated_inputs(): # check ckdtree with duplicated inputs n = 1024 for m in range(1, 8): data = np.concatenate([ np.full((n // 2, m), 1), np.full((n // 2, m), 2)], axis=0) # it shall not divide more than 3 nodes. # root left (1), and right (2) kdtree = cKDTree(data, leafsize=1) assert_equal(kdtree.size, 3) kdtree = cKDTree(data) assert_equal(kdtree.size, 3) # if compact_nodes are disabled, the number # of nodes is n (per leaf) + (m - 1)* 2 (splits per dimension) + 1 # and the root kdtree = cKDTree(data, compact_nodes=False, leafsize=1) assert_equal(kdtree.size, n + m * 2 - 1) def test_ckdtree_noncumulative_nondecreasing(): # check ckdtree with duplicated inputs # it shall not divide more than 3 nodes. # root left (1), and right (2) kdtree = cKDTree([[0]], leafsize=1) assert_raises(ValueError, kdtree.count_neighbors, kdtree, [0.1, 0], cumulative=False) def test_short_knn(): # The test case is based on github: #6425 by @SteveDoyle2 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 = cKDTree(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(): np.random.seed(1234) data = np.random.normal(size=(100, 3)) query = np.random.normal(size=(100, 3)) tree = cKDTree(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(): np.random.seed(1234) data = np.random.normal(size=(100, 3)) query = np.random.normal(size=(100, 3)) tree = cKDTree(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(): 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 = cKDTree(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_cdktree_empty_input(balanced_tree, compact_nodes): # https://github.com/scipy/scipy/issues/5040 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 = cKDTree(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) class Test_sorted_query_ball_point(object): def setup_method(self): np.random.seed(1234) self.x = np.random.randn(100, 1) self.ckdt = cKDTree(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(): # Test that KDTree rejects complex input points (gh-9108) 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)