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 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_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)
