有效地在非矩形2D网格上找到最近点的索引

Question

我有一个不规则的(非矩形)lon/lat网格和lon/lat坐标中的一堆点,它们应该对应于网格上的点(尽管由于数值原因它们可能略微偏离).现在我需要相应的lon/lat点的索引.

我已经编写了一个函数来执行此操作,但它确实很慢.

def find_indices(lon,lat,x,y):
    lonlat = np.dstack([lon,lat])
    delta = np.abs(lonlat-[x,y])
    ij_1d = np.linalg.norm(delta,axis=2).argmin()
    i,j = np.unravel_index(ij_1d,lon.shape)
    return i,j

ind = [find_indices(lon,lat,p*) for p in points]

我很确定在numpy/scipy中有更好(更快)的解决方案.我已经谷歌搜索了很多,但到目前为止我的答案还没有找到.

有关如何有效找到相应(最近)点的指数的任何建议？

PS:这个问题来自另一个问题(点击).

编辑:解决方案

根据@Cong Ma的回答,我找到了以下解决方案:

def find_indices(points,lon,lat,tree=None):
    if tree is None:
        lon,lat = lon.T,lat.T
        lonlat = np.column_stack((lon.ravel(),lat.ravel()))
        tree = sp.spatial.cKDTree(lonlat)
    dist,idx = tree.query(points,k=1)
    ind = np.column_stack(np.unravel_index(idx,lon.shape))
    return [(i,j) for i,j in ind]

为了解决这个问题以及Divakar的答案中的问题,下面是我使用find_indices的功能的一些时间(以及它在速度方面的瓶颈)(参见上面的链接):

spatial_contour_frequency/pil0                :   331.9553
spatial_contour_frequency/pil1                :   104.5771
spatial_contour_frequency/pil2                :     2.3629
spatial_contour_frequency/pil3                :     0.3287

pil0是我最初的方法,pil1Divakar和pil2/ pil3上面的最终解决方案,其中树是在运行中创建的pil2(即,在其中find_indices调用的循环的每次迭代)并且只有一次pil3(参见其他线程以获取详细信息).尽管Divakar对我的初始方法的改进使我的速度提高了3倍,但cKDTree将其提升到一个全新的水平,再加上50倍速!将树的创建移出函数会使事情变得更快.

Answer 1

如果这些点足够本地化，您可以尝试直接scipy.spatial实现cKDTree，正如我在另一篇文章中讨论的那样。那篇文章是关于插值的，但您可以忽略它并仅使用查询部分。

TL;博士版本：

阅读 的文档scipy.sptial.cKDTree。(n, m)通过将- 形状的ndarray 对象传递给初始化器来创建树numpy，并且将从 - 维坐标创建树n m。

tree = scipy.spatial.cKDTree(array_of_coordinates)

之后，用于tree.query()检索k第一个最近的邻居（可能使用近似和并行化，请参阅文档），或用于tree.query_ball_point()查找给定距离容差内的所有邻居。

如果这些点没有很好地定位，并且球面曲率/非平凡拓扑开始起作用，您可以尝试将流形分成多个部分，每个部分都足够小以被视为局部。