Python scipy.spatial.cKDTree() Examples

def cmi(x, y, z, k=3, base=2):
  """Mutual information of x and y, conditioned on z
  x,y,z should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
  if x is a one-dimensional scalar and we have four samples
  """
  assert len(x)==len(y), 'Lists should have same length.'
  assert k <= len(x) - 1, 'Set k smaller than num samples - 1.'
  intens = 1e-10 # Small noise to break degeneracy, see doc.
  x = [list(p + intens*nr.rand(len(x[0]))) for p in x]
  y = [list(p + intens*nr.rand(len(y[0]))) for p in y]
  z = [list(p + intens*nr.rand(len(z[0]))) for p in z]
  points = zip2(x,y,z)
  # Find nearest neighbors in joint space, p=inf means max-norm.
  tree = ss.cKDTree(points)
  dvec = [tree.query(point, k+1, p=float('inf'))[0][k] for point in points]
  a = avgdigamma(zip2(x,z), dvec)
  b = avgdigamma(zip2(y,z), dvec)
  c = avgdigamma(z,dvec)
  d = digamma(k)
  return (-a-b+c+d) / log(base) 

def mi(x, y, k=3, base=2):
  """Mutual information of x and y.
  x,y should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
  if x is a one-dimensional scalar and we have four samples.
  """
  assert len(x)==len(y), 'Lists should have same length.'
  assert k <= len(x) - 1, 'Set k smaller than num samples - 1.'
  intens = 1e-10 # Small noise to break degeneracy, see doc.
  x = [list(p + intens*nr.rand(len(x[0]))) for p in x]
  y = [list(p + intens*nr.rand(len(y[0]))) for p in y]
  points = zip2(x,y)
  # Find nearest neighbors in joint space, p=inf means max-norm.
  tree = ss.cKDTree(points)
  dvec = [tree.query(point, k+1, p=float('inf'))[0][k] for point in points]
  a = avgdigamma(x,dvec)
  b = avgdigamma(y,dvec)
  c = digamma(k)
  d = digamma(len(x))
  return (-a-b+c+d) / log(base) 

def adi(R_est, t_est, R_gt, t_gt, model):
    """
    Average Distance of Model Points for objects with indistinguishable views
    - by Hinterstoisser et al. (ACCV 2012).

    :param R_est, t_est: Estimated pose (3x3 rot. matrix and 3x1 trans. vector).
    :param R_gt, t_gt: GT pose (3x3 rot. matrix and 3x1 trans. vector).
    :param model: Object model given by a dictionary where item 'pts'
    is nx3 ndarray with 3D model points.
    :return: Error of pose_est w.r.t. pose_gt.
    """
    pts_est = misc.transform_pts_Rt(model['pts'], R_est, t_est)
    pts_gt = misc.transform_pts_Rt(model['pts'], R_gt, t_gt)

    # Calculate distances to the nearest neighbors from pts_gt to pts_est
    nn_index = spatial.cKDTree(pts_est)
    nn_dists, _ = nn_index.query(pts_gt, k=1)

    e = nn_dists.mean()
    return e 

def test_ridges(self, name):
        # Check that the ridges computed by Voronoi indeed separate
        # the regions of nearest neighborhood, by comparing the result
        # to KDTree.

        points = DATASETS[name]

        tree = KDTree(points)
        vor = qhull.Voronoi(points)

        for p, v in vor.ridge_dict.items():
            # consider only finite ridges
            if not np.all(np.asarray(v) >= 0):
                continue

            ridge_midpoint = vor.vertices[v].mean(axis=0)
            d = 1e-6 * (points[p[0]] - ridge_midpoint)

            dist, k = tree.query(ridge_midpoint + d, k=1)
            assert_equal(k, p[0])

            dist, k = tree.query(ridge_midpoint - d, k=1)
            assert_equal(k, p[1]) 

def test_query_ball_point_multithreading():
    np.random.seed(0)
    n = 5000
    k = 2
    points = np.random.randn(n,k)
    T = cKDTree(points)
    l1 = T.query_ball_point(points,0.003,n_jobs=1)
    l2 = T.query_ball_point(points,0.003,n_jobs=64)
    l3 = T.query_ball_point(points,0.003,n_jobs=-1)
    
    for i in range(n):
        if l1[i] or l2[i]:
            assert_array_equal(l1[i],l2[i])
        
    for i in range(n):
        if l1[i] or l3[i]:
            assert_array_equal(l1[i],l3[i]) 

def adi(pose_est, pose_gt, model):
    """
    Average Distance of Model Points for objects with indistinguishable views
    - by Hinterstoisser et al. (ACCV 2012).

    :param pose_est: Estimated pose given by a dictionary:
    {'R': 3x3 rotation matrix, 't': 3x1 translation vector}.
    :param pose_gt: The ground truth pose given by a dictionary (as pose_est).
    :param model: Object model given by a dictionary where item 'pts'
    is nx3 ndarray with 3D model points.
    :return: Error of pose_est w.r.t. pose_gt.
    """
    pts_gt = misc.transform_pts_Rt(model['pts'], pose_gt['R'], pose_gt['t'])
    pts_est = misc.transform_pts_Rt(model['pts'], pose_est['R'], pose_est['t'])

    # Calculate distances to the nearest neighbors from pts_gt to pts_est
    nn_index = spatial.cKDTree(pts_est)
    nn_dists, _ = nn_index.query(pts_gt, k=1)

    e = nn_dists.mean()
    return e 

def adi(R_est, t_est, R_gt, t_gt, pts):
  """Average Distance of Model Points for objects with indistinguishable views
  - by Hinterstoisser et al. (ACCV'12).

  :param R_est: 3x3 ndarray with the estimated rotation matrix.
  :param t_est: 3x1 ndarray with the estimated translation vector.
  :param R_gt: 3x3 ndarray with the ground-truth rotation matrix.
  :param t_gt: 3x1 ndarray with the ground-truth translation vector.
  :param pts: nx3 ndarray with 3D model points.
  :return: The calculated error.
  """
  pts_est = misc.transform_pts_Rt(pts, R_est, t_est)
  pts_gt = misc.transform_pts_Rt(pts, R_gt, t_gt)

  # Calculate distances to the nearest neighbors from vertices in the
  # ground-truth pose to vertices in the estimated pose.
  nn_index = spatial.cKDTree(pts_est)
  nn_dists, _ = nn_index.query(pts_gt, k=1)

  e = nn_dists.mean()
  return e 

def _view_kd_2d(cls, data, data_view, target_view, x_tolerance=1e-3, y_tolerance=1e-3):
        t_xmin, t_ymin, t_xmax, t_ymax = target_view.bbox
        d_xmin, d_ymin, d_xmax, d_ymax = data_view.bbox
        nodata = target_view.nodata
        view = np.full(target_view.shape, nodata)
        viewcoords = target_view.coords
        datacoords = data_view.coords
        yx_tolerance = np.sqrt(x_tolerance**2 + y_tolerance**2)
        row_bool = ((datacoords[:,0] <= t_ymax + y_tolerance) &
                    (datacoords[:,0] >= t_ymin - y_tolerance))
        col_bool = ((datacoords[:,1] <= t_xmax + x_tolerance) &
                    (datacoords[:,1] >= t_xmin - x_tolerance))
        yx_tree = datacoords[row_bool & col_bool]
        tree = spatial.cKDTree(yx_tree)
        yx_dist, yx_ix = tree.query(viewcoords)
        yx_passed = yx_dist <= yx_tolerance
        view.flat[yx_passed] = data.flat[row_bool & col_bool].flat[yx_ix[yx_passed]]
        return view 

def _view_kd_2d(cls, data, data_view, target_view, x_tolerance=1e-3, y_tolerance=1e-3):
        t_xmin, t_ymin, t_xmax, t_ymax = target_view.bbox
        d_xmin, d_ymin, d_xmax, d_ymax = data_view.bbox
        nodata = target_view.nodata
        view = np.full(target_view.shape, nodata)
        yx_tolerance = np.sqrt(x_tolerance**2 + y_tolerance**2)
        viewrows, viewcols = target_view.grid_indices()
        rows, cols = data_view.grid_indices()
        row_bool = (rows <= t_ymax + y_tolerance) & (rows >= t_ymin - y_tolerance)
        col_bool = (cols <= t_xmax + x_tolerance) & (cols >= t_xmin - x_tolerance)
        yx_tree = np.vstack(np.dstack(np.meshgrid(rows[row_bool], cols[col_bool], indexing='ij')))
        yx_query = np.vstack(np.dstack(np.meshgrid(viewrows, viewcols, indexing='ij')))
        tree = spatial.cKDTree(yx_tree)
        yx_dist, yx_ix = tree.query(yx_query)
        yx_passed = yx_dist < yx_tolerance
        view.flat[yx_passed] = data[np.ix_(row_bool, col_bool)].flat[yx_ix[yx_passed]]
        return view 

def setup_method(self):
        n = 50
        m = 4
        np.random.seed(0)
        data1 = np.random.randn(n,m)
        data2 = np.random.randn(n,m)
        self.T1 = cKDTree(data1,leafsize=2)
        self.T2 = cKDTree(data2,leafsize=2)
        self.ref_T1 = KDTree(data1, leafsize=2)
        self.ref_T2 = KDTree(data2, leafsize=2)
        self.r = 0.5
        self.n = n
        self.m = m
        self.data1 = data1
        self.data2 = data2
        self.p = 2 

def _create_mesh(self):
        """ Mesh assignment method

            Based on a target value, determine a mesh size for the testlines
            that is compatible with the simulation box.
            Create the grid and initialize a cKDTree object with it to
            facilitate fast searching of the gridpoints touched by molecules.
        """
        box = utilities.get_box(self.universe, self.normal)
        n, d = utilities.compute_compatible_mesh_params(self.target_mesh, box)
        self.mesh_nx = n[0]
        self.mesh_ny = n[1]
        self.mesh_dx = d[0]
        self.mesh_dy = d[1]

        _x = np.linspace(0, box[0], num=int(self.mesh_nx), endpoint=False)
        _y = np.linspace(0, box[1], num=int(self.mesh_ny), endpoint=False)
        _X, _Y = np.meshgrid(_x, _y)
        self.meshpoints = np.array([_X.ravel(), _Y.ravel()]).T
        self.meshtree = cKDTree(self.meshpoints, boxsize=box[:2]) 

def _view_kd(cls, data, data_view, target_view, x_tolerance=1e-3, y_tolerance=1e-3):
        """
        Appropriate if:
            - Grid is regular
            - Data is regular
            - Grid and data have same cellsize OR no interpolation is needed
        """
        nodata = target_view.nodata
        view = np.full(target_view.shape, nodata)
        viewrows, viewcols = target_view.grid_indices()
        rows, cols = data_view.grid_indices()
        ytree = spatial.cKDTree(rows[:, None])
        xtree = spatial.cKDTree(cols[:, None])
        ydist, y_ix = ytree.query(viewrows[:, None])
        xdist, x_ix = xtree.query(viewcols[:, None])
        y_passed = ydist < y_tolerance
        x_passed = xdist < x_tolerance
        view[np.ix_(y_passed, x_passed)] = data[y_ix[y_passed]][:, x_ix[x_passed]]
        return view 

def test_onetree_query_compiled():
    np.random.seed(0)
    n = 100
    k = 4
    points = np.random.randn(n,k)
    T = cKDTree(points)
    check_onetree_query(T, 0.1)

    points = np.random.randn(3*n,k)
    points[:n] *= 0.001
    points[n:2*n] += 2
    T = cKDTree(points)
    check_onetree_query(T, 0.1)
    check_onetree_query(T, 0.001)
    check_onetree_query(T, 0.00001)
    check_onetree_query(T, 1e-6) 

def adi(R_est, t_est, R_gt, t_gt, model):
    """
    Average Distance of Model Points for objects with indistinguishable views
    - by Hinterstoisser et al. (ACCV 2012).

    :param R_est, t_est: Estimated pose (3x3 rot. matrix and 3x1 trans. vector).
    :param R_gt, t_gt: GT pose (3x3 rot. matrix and 3x1 trans. vector).
    :param model: Object model given by a dictionary where item 'pts'
    is nx3 ndarray with 3D model points.
    :return: Error of pose_est w.r.t. pose_gt.
    """
    pts_est = misc.transform_pts_Rt(model['pts'], R_est, t_est)
    pts_gt = misc.transform_pts_Rt(model['pts'], R_gt, t_gt)

    # Calculate distances to the nearest neighbors from pts_gt to pts_est
    nn_index = spatial.cKDTree(pts_est)
    nn_dists, _ = nn_index.query(pts_gt, k=1)

    e = nn_dists.mean()
    return e 

def evaluate_pbc_fast(self, points):
        grid = points
        pos = self.pos
        box = self.box
        d = self.sigma * 2.5
        results = np.zeros(grid.shape[1], dtype=float)
        gridT = grid[::-1].T[:]
        tree = cKDTree(gridT, boxsize=box)
        # the indices of grid elements within distane d from each of the pos
        scale = 2. * self.sigma**2
        indlist = tree.query_ball_point(pos, d)
        for n, ind in enumerate(indlist):
            dr = gridT[ind, :] - pos[n]
            cond = np.where(dr > box / 2.)
            dr[cond] -= box[cond[1]]
            cond = np.where(dr < -box / 2.)
            dr[cond] += box[cond[1]]
            dens = np.exp(-np.sum(dr * dr, axis=1) / scale)
            results[ind] += dens

        return results 

def test_ckdtree_pickle():
    # test if it is possible to pickle
    # a cKDTree
    try:
        import cPickle as pickle
    except ImportError:
        import pickle
    np.random.seed(0)
    n = 50
    k = 4
    points = np.random.randn(n, k)
    T1 = cKDTree(points)
    tmp = pickle.dumps(T1)
    T2 = pickle.loads(tmp)
    T1 = T1.query(points, k=5)[-1]
    T2 = T2.query(points, k=5)[-1]
    assert_array_equal(T1, T2) 

def kldiv(x, xp, k=3, base=2):
  """KL Divergence between p and q for x~p(x),xp~q(x).
  x, xp should be a list of vectors, e.g. x = [[1.3],[3.7],[5.1],[2.4]]
  if x is a one-dimensional scalar and we have four samples
  """
  assert k <= len(x) - 1, 'Set k smaller than num samples - 1.'
  assert k <= len(xp) - 1, 'Set k smaller than num samples - 1.'
  assert len(x[0]) == len(xp[0]), 'Two distributions must have same dim.'
  d = len(x[0])
  n = len(x)
  m = len(xp)
  const = log(m) - log(n-1)
  tree = ss.cKDTree(x)
  treep = ss.cKDTree(xp)
  nn = [tree.query(point, k+1, p=float('inf'))[0][k] for point in x]
  nnp = [treep.query(point, k, p=float('inf'))[0][k-1] for point in x]
  return (const + d * np.mean(map(log, nnp)) \
    - d * np.mean(map(log, nn))) / log(base)

# DISCRETE ESTIMATORS 

def test_ckdtree_pickle_boxsize():
    # test if it is possible to pickle a periodic
    # cKDTree
    try:
        import cPickle as pickle
    except ImportError:
        import pickle
    np.random.seed(0)
    n = 50
    k = 4
    points = np.random.uniform(size=(n, k))
    T1 = cKDTree(points, boxsize=1.0)
    tmp = pickle.dumps(T1)
    T2 = pickle.loads(tmp)
    T1 = T1.query(points, k=5)[-1]
    T2 = T2.query(points, k=5)[-1]
    assert_array_equal(T1, T2) 

def _find_correspondence(surf, ref_surf, eps=0, n_jobs=1, use_cell=False):

    if use_cell:
        points = compute_cell_center(surf)
        ref_points = compute_cell_center(ref_surf)
    else:
        points = get_points(surf)
        ref_points = get_points(ref_surf)

    tree = cKDTree(ref_points, leafsize=20, compact_nodes=False,
                   copy_data=False, balanced_tree=False)
    d, idx = tree.query(points, k=1, eps=0, n_jobs=n_jobs,
                        distance_upper_bound=eps+np.finfo(np.float).eps)

    if np.isinf(d).any():
        raise ValueError('Cannot find correspondences. Try increasing '
                         'tolerance.')

    return idx 

def bench_build(self):
        print()
        print('        Constructing kd-tree')
        print('=====================================')
        print(' dim | # points |  KDTree  | cKDTree ')

        for (m, n, repeat) in [(3,10000,3), (8,10000,3), (16,10000,3)]:
            print('%4s | %7s ' % (m, n), end=' ')
            sys.stdout.flush()

            data = np.concatenate((np.random.randn(n//2,m),
                                   np.random.randn(n-n//2,m)+np.ones(m)))

            print('| %6.3fs ' % (measure('T1 = KDTree(data)', repeat) / repeat), end=' ')
            sys.stdout.flush()
            print('| %6.3fs' % (measure('T2 = cKDTree(data)', repeat) / repeat), end=' ')
            sys.stdout.flush()
            print('') 

def test_against_logic_error_regression(self):
        # regression test for gh-5077 logic error
        np.random.seed(0)
        too_many = np.array(np.random.randn(18, 2), dtype=int)
        tree = cKDTree(too_many, balanced_tree=False, compact_nodes=False)
        d = tree.sparse_distance_matrix(tree, 3).todense()
        assert_array_almost_equal(d, d.T, decimal=14) 

def setup_method(self):
        Test_small.setup_method(self)
        self.kdtree = cKDTree(self.data) 

def setup_method(self):
        n = 10000
        m = 4
        np.random.seed(1234)
        self.data = np.random.uniform(size=(n,m))
        self.T = cKDTree(self.data,leafsize=2, boxsize=1)
        self.x = np.ones(m) * 0.1
        self.p = 2.
        self.eps = 0
        self.d = 0.2 

def build_kdtree(self, grid='node'):
        """Builds the kdtree for the specified grid"""
        try:
            from scipy.spatial import cKDTree
        except ImportError:
            raise ImportError("The scipy package is required to use "
                              "SGrid.locate_nearest\n"
                              " -- nearest neighbor interpolation")
        lon, lat = self._get_grid_vars(grid)
        if lon is None or lat is None:
            raise ValueError("{0}_lon and {0}_lat must be defined in order to "
                             "create and use KDTree for this grid".format(grid))
        lin_points = np.column_stack((lon.ravel(), lat.ravel()))
        self._kd_trees[grid] = cKDTree(lin_points, leafsize=4) 

def _build_kdtree(self):
        # Only import if it's used.
        try:
            from scipy.spatial import cKDTree
        except ImportError:
            raise ImportError("The scipy package is required to use "
                              "UGrid.locate_nodes\n"
                              " -- nearest neighbor interpolation")

        self._kdtree = cKDTree(self.nodes) 

def update_jname():

    from grizli import utils

    res = grizli_db.from_sql("select p_root, p_id, p_ra, p_dec from photometry_apcorr", engine)

    jn = [utils.radec_to_targname(ra=ra, dec=dec, round_arcsec=(0.001, 0.001), precision=2, targstr='j{rah}{ram}{ras}.{rass}{sign}{ded}{dem}{des}.{dess}') for ra, dec in zip(res['p_ra'], res['p_dec'])]

    for c in res.colnames:
        res.rename_column(c, c.replace('p_', 'j_'))

    zres = grizli_db.from_sql("select root, phot_root, id, ra, dec, z_map,"
                              "q_z, t_g800l, t_g102, t_g141, status from "
                              "redshift_fit where ra is not null and "
                              "status > 5", engine)

    # Find duplicates
    from scipy.spatial import cKDTree
    data = np.array([zres['ra'], zres['dec']]).T

    ok = zres['q_z'].filled(-100) > -0.7

    tree = cKDTree(data[ok])
    dr, ix = tree.query(data[ok], k=2)
    cosd = np.cos(data[:, 1]/180*np.pi)
    dup = (dr[:, 1] < 0.01/3600)  # & (zres['phot_root'][ix[:,0]] != zres['phot_root'][ix[:,1]])

    ix0 = ix[:, 0]
    ix1 = ix[:, 1]

    dup = (dr[:, 1] < 0.01/3600)
    dup &= (zres['phot_root'][ok][ix0] == zres['phot_root'][ok][ix1])
    dup &= (zres['id'][ok][ix0] == zres['id'][ok][ix1])

    # second is G800L
    dup &= zres['t_g800l'].filled(0)[ok][ix1] > 10

    plt.scatter(zres['z_map'][ok][ix0[dup]], zres['z_map'][ok][ix1[dup]],
                marker='.', alpha=0.1) 

def reweight(self, uv_radius, weightdist=1.0):
        """Reweight the sigmas based on the local  density of uv points

           Args:
               uv_radius (float): radius in uv-plane to look for nearby points
               weightdist (float): ??

           Returns:
               (Obsdata): a new reweighted observation object.
        """

        obs_new = self.copy()
        npts = len(obs_new.data)

        uvpoints = np.vstack((obs_new.data['u'], obs_new.data['v'])).transpose()
        uvpoints_tree1 = spatial.cKDTree(uvpoints)
        uvpoints_tree2 = spatial.cKDTree(-uvpoints)

        for i in range(npts):
            matches1 = uvpoints_tree1.query_ball_point(uvpoints[i, :], uv_radius)
            matches2 = uvpoints_tree2.query_ball_point(uvpoints[i, :], uv_radius)
            nmatches = len(matches1) + len(matches2)

            for sigma in ['sigma', 'qsigma', 'usigma', 'vsigma']:
                obs_new.data[sigma][i] = np.sqrt(nmatches)

        scale = np.mean(self.data['sigma']) / np.mean(obs_new.data['sigma'])
        for sigma in ['sigma', 'qsigma', 'usigma', 'vsigma']:
            obs_new.data[sigma] *= scale * weightdist

        if weightdist < 1.0:
            for i in range(npts):
                for sigma in ['sigma', 'qsigma', 'usigma', 'vsigma']:
                    obs_new.data[sigma][i] += (1 - weightdist) * self.data[sigma][i]

        return obs_new 

def _get_pids_naive(source, target, k=1, source_mask=None, target_mask=None,
                    return_weights=True, n_jobs=1):
    """Resampling based on k nearest points."""

    sp = me.get_points(source, mask=source_mask)
    tp = me.get_points(target, mask=target_mask)

    tree = cKDTree(sp, leafsize=20, compact_nodes=False, copy_data=False,
                   balanced_tree=False)
    dist, pids = tree.query(tp, k=k, eps=0, n_jobs=n_jobs)

    if return_weights:
        return pids, 1 / dist
    return pids 

def __init__(self, x, y):
        x = _ndim_coords_from_arrays(x)
        self._check_init_shape(x, y)
        self.tree = cKDTree(x)
        self.points = x
        self.values = y 

def test_ball_point_ints():
    """Regression test for #1373."""
    x, y = np.mgrid[0:4, 0:4]
    points = list(zip(x.ravel(), y.ravel()))
    tree = KDTree(points)
    assert_equal(sorted([4, 8, 9, 12]),
                 sorted(tree.query_ball_point((2, 0), 1)))
    points = np.asarray(points, dtype=np.float)
    tree = KDTree(points)
    assert_equal(sorted([4, 8, 9, 12]),
                 sorted(tree.query_ball_point((2, 0), 1)))

# cKDTree is specialized to type double points, so no need to make
# a unit test corresponding to test_ball_point_ints()