Ejemplo n.º 1
0
def test_kdtree_k1():
    kdtree = KDTree([X1, X2])
    dists, inds = kdtree.query([
        [0] * 10,
        [10] * 10,
        [0] * 10
    ])

    assert len(inds) == 3
    for subind in inds:
        assert len(subind) == 2

    # traj i
    assert inds[0][0] == 0
    assert inds[1][0] == 1
    assert inds[2][0] == 0

    # framei
    assert 0 <= inds[0][1] < 500
    assert 0 <= inds[1][1] < 1000
    assert 0 <= inds[2][1] < 500

    # distances
    assert len(dists) == 3
    for d in dists:
        assert 0 <= d < 0.5
Ejemplo n.º 2
0
def test_kdtree_k2():
    kdtree = KDTree([X1, X2])
    dists, inds = kdtree.query([
        [0] * 10,
        [10] * 10,
        [0] * 10
    ], k=2)

    assert len(inds) == 3

    # traj i
    for qp in inds[0]: assert qp[0] == 0
    for qp in inds[1]: assert qp[0] == 1
    for qp in inds[2]: assert qp[0] == 0

    # frame i
    for qp in inds[0]: assert 0 <= qp[1] < 500
    for qp in inds[1]: assert 0 <= qp[1] < 1000
    for qp in inds[2]: assert 0 <= qp[1] < 500

    # distances
    assert len(dists) == 3
    for d in dists:
        assert 0 <= d[0] < 0.5
        assert 0 <= d[1] < 0.5
Ejemplo n.º 3
0
def test_kdtree_k2():
    kdtree = KDTree([X1, X2])
    dists, inds = kdtree.query([[0] * 10, [10] * 10, [0] * 10], k=2)

    assert len(inds) == 3

    # traj i
    for qp in inds[0]:
        assert qp[0] == 0
    for qp in inds[1]:
        assert qp[0] == 1
    for qp in inds[2]:
        assert qp[0] == 0

    # frame i
    for qp in inds[0]:
        assert 0 <= qp[1] < 500
    for qp in inds[1]:
        assert 0 <= qp[1] < 1000
    for qp in inds[2]:
        assert 0 <= qp[1] < 500

    # distances
    assert len(dists) == 3
    for d in dists:
        assert 0 <= d[0] < 0.5
        assert 0 <= d[1] < 0.5
Ejemplo n.º 4
0
def test_kdtree_k1():
    kdtree = KDTree([X1, X2])
    dists, inds = kdtree.query([[0] * 10, [10] * 10, [0] * 10])

    assert len(inds) == 3
    for subind in inds:
        assert len(subind) == 2

    # traj i
    assert inds[0][0] == 0
    assert inds[1][0] == 1
    assert inds[2][0] == 0

    # framei
    assert 0 <= inds[0][1] < 500
    assert 0 <= inds[1][1] < 1000
    assert 0 <= inds[2][1] < 500

    # distances
    assert len(dists) == 3
    for d in dists:
        assert 0 <= d < 0.5
Ejemplo n.º 5
0
def find_closest_indices_to_cluster_center(tica_coords, clusterer_file, k=1):
    tica = compat_verboseload(tica_coords)
    clusterer = compat_verboseload(clusterer_file)
    kd = KDTree(tica)
    dist, inds = kd.query(clusterer.cluster_centers_, k=k)
    return inds