def check_neighbors(dualtree, breadth_first, k, metric, kwargs): kdt = KDTree(X, leaf_size=1, metric=metric, **kwargs) dist1, ind1 = kdt.query(Y, k, dualtree=dualtree, breadth_first=breadth_first) dist2, ind2 = brute_force_neighbors(X, Y, k, metric, **kwargs) # don't check indices here: if there are any duplicate distances, # the indices may not match. Distances should not have this problem. assert_allclose(dist1, dist2)
def test_kd_tree_pickle(): import pickle np.random.seed(0) X = np.random.random((10, 3)) kdt1 = KDTree(X, leaf_size=1) ind1, dist1 = kdt1.query(X) def check_pickle_protocol(protocol): s = pickle.dumps(kdt1, protocol=protocol) kdt2 = pickle.loads(s) ind2, dist2 = kdt2.query(X) assert_allclose(ind1, ind2) assert_allclose(dist1, dist2) for protocol in (0, 1, 2): yield check_pickle_protocol, protocol
def test_should_find_closest_1(): kdtree = KDTree(dataset_2d) dist, ind = kdtree.query(np.array([9, 2]), n=1) assert ind == [4]
def test_should_find_n_closest_9(): kdtree = KDTree(dataset_3d) dist, ind = kdtree.query(np.array([4, 5, 6]), n=3) assert ind == [2, 3, 1] # ????? падает!
def test_should_find_n_closest_8(): kdtree = KDTree(dataset_2d) dist, ind = kdtree.query(np.array([3, 2]), n=3) assert ind == [0, 1, 5]
def test_should_find_n_closest_7(): kdtree = KDTree(dataset_2d) dist, ind = kdtree.query(np.array([7, 1]), n=3) assert ind == [5, 4, 1]
def test_should_find_closest_6(): kdtree = KDTree(dataset_3d) dist, ind = kdtree.query(np.array([8, 8, 8]), n=1) assert ind == [4]
def test_should_find_closest_5(): kdtree = KDTree(dataset_3d) dist, ind = kdtree.query(np.array([4, 5, 6]), n=1) assert ind == [2]
def test_should_find_closest_4(): kdtree = KDTree(dataset_3d) dist, ind = kdtree.query(np.array([1, 2, 3]), n=1) assert ind == [3]
def main(): dataset_2d = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]]) kdtree = KDTree(dataset_2d) dist, ind = kdtree.query(np.array([9, 2]), n=1)