Python kdtree.KDTree方法代码示例

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def get_tree(name):
        """
        Get KD-tree for the tile with the given name.

        Args:
            name(str): The name of the tile.
        """
        from pykdtree.kdtree import KDTree
        lat_grid, lon_grid = SRTM30.get_grids(name)
        lat_grid, lon_grid = np.meshgrid(lat_grid, lon_grid, indexing = "ij")
        x, y, z = _latlon_to_cart(lat_grid, lon_grid)
        X = np.concatenate([x.reshape(-1, 1, order = "C"),
                            y.reshape(-1, 1, order = "C"),
                            z.reshape(-1, 1, order = "C")], axis = 1)
        tree = KDTree(X.astype(np.float32))
        return tree 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d():


    #7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])


    kdtree = KDTree(data_pts_real)
    dist, idx = kdtree.query(query_pts, sqr_dists=True)

    epsilon = 1e-5
    assert idx[0] == 7
    assert idx[1] == 93
    assert idx[2] == 45
    assert dist[0] == 0
    assert abs(dist[1] - 3.) < epsilon * dist[1]
    assert abs(dist[2] - 20001.) < epsilon * dist[2] 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_float32():


    #7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]], dtype=np.float32)


    kdtree = KDTree(data_pts_real.astype(np.float32))
    dist, idx = kdtree.query(query_pts, sqr_dists=True)
    epsilon = 1e-5
    assert idx[0] == 7
    assert idx[1] == 93
    assert idx[2] == 45
    assert dist[0] == 0
    assert abs(dist[1] - 3.) < epsilon * dist[1]
    assert abs(dist[2] - 20001.) < epsilon * dist[2]
    assert kdtree.data_pts.dtype == np.float32 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_8n_ub():
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real)
    dist, idx = kdtree.query(query_pts, k=8, distance_upper_bound=10e3, sqr_dists=False)

    exp_dist = np.array([[  0.00000000e+00,   4.05250235e+03,   4.07389794e+03,   8.08201128e+03,
                            8.17063009e+03,   np.Inf,   np.Inf,   np.Inf],
                        [  1.73205081e+00,   2.70216896e+03,   2.71431274e+03,   5.39537066e+03,
                            5.43793210e+03,   8.07855631e+03,   8.17119970e+03,   np.Inf],
                        [  1.41424892e+02,   3.25500021e+03,   3.44284958e+03,   6.58019346e+03,
                            6.81038455e+03,   9.89140135e+03,   np.Inf,   np.Inf]])
    n = 100
    exp_idx = np.array([[ 7,  8,  6,  9,  5, n,  n, n],
                        [93, 94, 92, 95, 91, 96, 90, n],
                        [45, 46, 44, 47, 43, 48, n, n]])

    assert np.array_equal(idx, exp_idx)
    assert np.allclose(dist, exp_dist) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_8n_ub_leaf20():
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real, leafsize=20)
    dist, idx = kdtree.query(query_pts, k=8, distance_upper_bound=10e3, sqr_dists=False)

    exp_dist = np.array([[  0.00000000e+00,   4.05250235e+03,   4.07389794e+03,   8.08201128e+03,
                            8.17063009e+03,   np.Inf,   np.Inf,   np.Inf],
                        [  1.73205081e+00,   2.70216896e+03,   2.71431274e+03,   5.39537066e+03,
                            5.43793210e+03,   8.07855631e+03,   8.17119970e+03,   np.Inf],
                        [  1.41424892e+02,   3.25500021e+03,   3.44284958e+03,   6.58019346e+03,
                            6.81038455e+03,   9.89140135e+03,   np.Inf,   np.Inf]])
    n = 100
    exp_idx = np.array([[ 7,  8,  6,  9,  5, n,  n, n],
                        [93, 94, 92, 95, 91, 96, 90, n],
                        [45, 46, 44, 47, 43, 48, n, n]])

    assert np.array_equal(idx, exp_idx)
    assert np.allclose(dist, exp_dist) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_8n_ub_eps():
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real)
    dist, idx = kdtree.query(query_pts, k=8, eps=0.1, distance_upper_bound=10e3, sqr_dists=False)

    exp_dist = np.array([[  0.00000000e+00,   4.05250235e+03,   4.07389794e+03,   8.08201128e+03,
                            8.17063009e+03,   np.Inf,   np.Inf,   np.Inf],
                        [  1.73205081e+00,   2.70216896e+03,   2.71431274e+03,   5.39537066e+03,
                            5.43793210e+03,   8.07855631e+03,   8.17119970e+03,   np.Inf],
                        [  1.41424892e+02,   3.25500021e+03,   3.44284958e+03,   6.58019346e+03,
                            6.81038455e+03,   9.89140135e+03,   np.Inf,   np.Inf]])
    n = 100
    exp_idx = np.array([[ 7,  8,  6,  9,  5, n,  n, n],
                        [93, 94, 92, 95, 91, 96, 90, n],
                        [45, 46, 44, 47, 43, 48, n, n]])

    assert np.array_equal(idx, exp_idx)
    assert np.allclose(dist, exp_dist) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_large_query():
    # Target idxs: 7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    # Repeat the same points multiple times to get 60000 query points
    n = 20000
    query_pts = np.repeat(query_pts, n, axis=0)

    kdtree = KDTree(data_pts_real)
    dist, idx = kdtree.query(query_pts, sqr_dists=True)

    epsilon = 1e-5
    assert np.all(idx[:n] == 7)
    assert np.all(idx[n:2*n] == 93)
    assert np.all(idx[2*n:] == 45)
    assert np.all(dist[:n] == 0)
    assert np.all(abs(dist[n:2*n] - 3.) < epsilon * dist[n:2*n])
    assert np.all(abs(dist[2*n:] - 20001.) < epsilon * dist[2*n:]) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_mask():
    #7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real)
    query_mask = np.zeros(data_pts_real.shape[0])
    query_mask[6:10] = True
    dist, idx = kdtree.query(query_pts, sqr_dists=True, mask=query_mask)

    epsilon = 1e-5
    assert idx[0] == 5  # would be 7 if no mask
    assert idx[1] == 93
    assert idx[2] == 45
    # would be 0 if no mask
    assert abs(dist[0] - 66759196.1053) < epsilon * dist[0]
    assert abs(dist[1] - 3.) < epsilon * dist[1]
    assert abs(dist[2] - 20001.) < epsilon * dist[2] 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def _create_resample_kdtree(self):
        """Set up kd tree on input."""
        # Get input information
        valid_input_index, source_lons, source_lats = \
            _get_valid_input_index_dask(self.source_geo_def,
                                        self.target_geo_def,
                                        self.reduce_data,
                                        self.radius_of_influence)

        # FIXME: Is dask smart enough to only compute the pixels we end up
        #        using even with this complicated indexing
        input_coords = lonlat2xyz(source_lons, source_lats)
        valid_input_index = da.ravel(valid_input_index)
        input_coords = input_coords[valid_input_index, :]
        input_coords = input_coords.compute()
        # Build kd-tree on input
        input_coords = input_coords.astype(np.float)
        valid_input_index, input_coords = da.compute(valid_input_index,
                                                     input_coords)
        return valid_input_index, KDTree(input_coords) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def nearest_neighbour(a, b):
    """
    find the nearest neighbours of a in b using KDTree
    Parameters
    ----------
    a : (n, ) numpy.ndarray
    b : (n, ) numpy.ndarray

    Returns
    ----------
    dist : n float
      Euclidian distance of the closest neighbour in b to a
    index : n float
      The index of the closest neighbour in b to a in terms of Euclidian distance
    """
    tree = KDTree(b)
    dist, index = tree.query(a)
    return (dist, index) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def kdtree(coordinates, use_pykdtree=True, **kwargs):
    """
    Create a KD-Tree object with the given coordinate arrays.

    Automatically transposes and flattens the coordinate arrays into a single
    matrix for use in the KD-Tree classes.

    All other keyword arguments are passed to the KD-Tree class.

    If installed, package ``pykdtree`` will be used instead of
    :class:`scipy.spatial.cKDTree` for better performance. Not all features are
    available in ``pykdtree`` so if you require the scipy version set
    ``use_pykdtee=False``.

    Parameters
    ----------
    coordinates : tuple of arrays
        Arrays with the coordinates of each data point. Should be in the
        following order: (easting, northing, vertical, ...). All coordinate
        arrays are used.
    use_pykdtree : bool
        If True, will prefer ``pykdtree`` (if installed) over
        :class:`scipy.spatial.cKDTree`. Otherwise, always use the scipy
        version.

    Returns
    -------
    tree : :class:`scipy.spatial.cKDTree` or ``pykdtree.kdtree.KDTree``
        The tree instance initialized with the given coordinates and arguments.

    """
    points = np.transpose(n_1d_arrays(coordinates, len(coordinates)))
    if pyKDTree is not None and use_pykdtree:
        tree = pyKDTree(points, **kwargs)
    else:
        tree = cKDTree(points, **kwargs)
    return tree 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test1d():

    data_pts = np.arange(1000)
    kdtree = KDTree(data_pts, leafsize=15)
    query_pts = np.arange(400, 300, -10)
    dist, idx = kdtree.query(query_pts)
    assert idx[0] == 400
    assert dist[0] == 0
    assert idx[1] == 390 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_float32_mismatch():


    #7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]], dtype=np.float32)

    kdtree = KDTree(data_pts_real)
    dist, idx = kdtree.query(query_pts, sqr_dists=True) 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test3d_float32_mismatch2():


    #7, 93, 45
    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real.astype(np.float32))
    try:
        dist, idx = kdtree.query(query_pts, sqr_dists=True)
        assert False
    except TypeError:
        assert True 

# 需要导入模块: from pykdtree import kdtree [as 别名]
# 或者: from pykdtree.kdtree import KDTree [as 别名]
def test_scipy_comp():

    query_pts = np.array([[  787014.438,  -340616.906,  6313018.],
                          [751763.125, -59925.969, 6326205.5],
                          [769957.188, -202418.125, 6321069.5]])

    kdtree = KDTree(data_pts_real)
    assert id(kdtree.data) == id(kdtree.data_pts)