cKDTree_MP.py 文件源码

def _pquery(scheduler, data, ndata, ndim, leafsize,
            x, nx, d, i, k, eps, p, dub, ierr):
    try:
        _data = shmem_as_nparray(data).reshape((ndata, ndim))
        _x = shmem_as_nparray(x).reshape((nx, ndim))
        _d = shmem_as_nparray(d).reshape((nx, k))
        _i = shmem_as_nparray(i).reshape((nx, k))

        kdtree = cKDTree(_data, leafsize=leafsize)

        for s in scheduler:
            d_out, i_out = kdtree.query(_x[s, :], k=k, eps=eps, p=p,
                                        distance_upper_bound=dub)
            m_d = d_out.shape[0]
            m_i = i_out.shape[0]
            _d[s, :], _i[s, :] = d_out.reshape(m_d, 1), i_out.reshape(m_i, 1)
    except:
        ierr.value += 1