python - 用于在两个数组上搜索 ckdtree 的向量化循环

Question

我有一个 csv 文件，其中包含随机位置的纬度、经度和海拔值。我想应用 IDW 插值来生成规则网格。我使用 scipy.spatial.cKDTree 进行最近邻搜索并查找未知点的高程值。当输出网格具有尺寸 (z < 1000 X1000) 时，以下代码可以正常工作。如果维度增加，代码运行速度会非常慢。请帮助我矢量化 for 循环，而不使用 cKDTree 进行删除。谢谢。

## Inverse distance weighted function
def idw(p, dist, values):

    dist_pow = np.power(dist, 2)
    nominator = np.sum(values/dist_pow)
    denominator = np.sum(1/dist_pow)
    if denominator > 0:
        return nominator/denominator
    else:
        return none
## Reading the lat/lon and elevation values from file
lat = []
lon = []
ele = []

with open('VSKP_ground_dat.csv') as read:
    csvreader = csv.DictReader(read)
    for row in csvreader:
        lat.append(float(row['LAT']))
        lon.append(float(row['LON']))
        ele.append(float(row['ALT']))
xycoord = np.c_[lon,lat]
ele_arr = np.array(ele)

## ------------- Creating KDTree
point_tree = spatial.cKDTree(xycoord, leafsize=25)
## ------------- Creating empty grid matrix with np.zeros
xmin, xmax, ymin, ymax  = 81.903158, 83.352158, 17.25856, 18.40056
## --------- Defining resolution
xres, yres = 0.01, 0.01

x = np.arange(xmin, xmax, xres)
y = np.arange(ymin, ymax, yres)
z = np.zeros((x.shape[0], y.shape[0]), dtype=np.float16)


for i, val1 in enumerate(x):
    for j, val2 in enumerate(y):
        p = np.array([val1, val2])
        # points_idx = point_tree.query_ball_point(p, dist_2)
        distances, points_idx = point_tree.query(p, k=6, eps=0)
        ele_vals = ele_arr[points_idx]
        value = idw(p, distances, ele_vals)
        z[i,j] = value

Answer 1

首先，修复您的 idw 函数以处理最后一个索引:

def idw(dist, values, p = 2):
    out = np.empty(dist.shape[:-1])
    mask = np.isclose(dist, 0).any(-1)
    out[mask] = values[np.isclose(dist, 0)]                   # should be only one per point
    dist_pow = np.power(dist[~mask], -p)                      # division is costly, do it once
    nominator = np.sum(values[~mask] * dist_pow, axis = -1)  # over mask to prevent divide by zero
    denominator = np.sum(dist_pow, index = -1)
    out[~mask] = nominator / denominator
    return out

然后根据np.meshgrid输出完成剩下的工作

x = np.arange(xmin, xmax, xres)                           # len i
y = np.arange(ymin, ymax, yres)                           # len j
xy = np.stack(np.meshgrid(x, y), axis = -1)               # shape(i, j, 2)
distances, points_idx = point_tree.query(xy, k=6, eps=0)  # shape (i, j, 6)
ele_vals = ele_arr[points_idx]                            # shape (i, j, 6)
z = idw(distances, ele_vals)                              # shape (i, j)