test.py 文件源码

def test_point_correctness():
    import itertools
    stencil = [-1, 0, 1]
    ndim = 3
    n = 2000
    stencil = itertools.product(*[stencil]*ndim)
    stencil = np.array(list(stencil)).astype(np.int32)

    points = (np.random.rand(n, ndim) * [1, 2, 3]).astype(np.float32)
    scale = 0.1

    spec = GridSpec(points, float(scale))
    offsets = spec.stencil(stencil).astype(np.int32)
    grid = PointGrid(spec, points, offsets)

    pairs = grid.pairs()

    from scipy.spatial import cKDTree
    tree = cKDTree(points)
    tree_pairs = tree.query_pairs(scale, output_type='ndarray')
    print(tree_pairs)
    print(pairs)

    assert np.alltrue(npi.unique(tree_pairs) == npi.unique(np.sort(pairs, axis=1)))