Function _pquery

reverse_geocoder/cKDTree_MP.py:17–36

Function that parallelly queries the K-D tree based on chunks of data returned by the scheduler

(scheduler, data, ndata, ndim, leafsize,
            x, nx, d, i, k, eps, p, dub, ierr)

    return np.frombuffer(shmem_array.get_obj())

def _pquery(scheduler, data, ndata, ndim, leafsize,
            x, nx, d, i, k, eps, p, dub, ierr):
    """
    Function that parallelly queries the K-D tree based on chunks of data returned by the scheduler
    """
    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

def num_cpus():
    """