通过矢量化更快地运行

Question

我正在寻找优化这个循环，但我没有想法/技能。 我尝试用 numba、jax 重写它，但没有成功。 实施不重新计算的条件（增益很小），但运行时间仍然在 30 秒左右。 我认为矢量化是一种方法，但我找不到正确实现它的方法。任何人都可以向我展示如何对其进行矢量化，因为知道所选输入将具有不同的长度？

#float, float, float, int
x, y, z, idx = df['x'].values, df['y'].values, df[z_col].values, df.index.values
azi_rad = np.radians(azi)
cos_azi_rad, sin_azi_rad, tan_ang_tol_rad, ang_tol_rad = np.cos(azi_rad), np.sin(azi_rad), np.radians(ang_tol), np.tan(ang_tol_rad)
band_width, max_distance = 100, 1000

#will receive results by indentation later
res_select=np.empty_like(idx)

for id in idx :
    #origin point x,y,z
    xi, yi, zi = x[id], y[id], z[id]
    
    #cutting to smaller chunk
    cond_xy_cut = ((xi-max_distance) <= x) & (x <= (xi+max_distance)) & ((yi-max_distance) <= y) & (y <= (yi+max_distance))
    x_cut, y_cut, z_cut, w_cut, idx_cut = x[cond_xy_cut], y[cond_xy_cut], z[cond_xy_cut], weights[cond_xy_cut], idx[cond_xy_cut]

    #computing delta from origin point once for the 2 're-use' under
    dx,dy = x_cut-xi, y_cut-yi

    #rotating x,y around origin point at azi angle
    x_rotate = cos_azi_rad*(dx) - sin_azi_rad*(dy) + xi
    y_rotate = sin_azi_rad*(dx) + cos_azi_rad*(dy) + yi

    #computing new delta with x_rotate, y_rotate
    dx, dy = x_rotate-xi, y_rotate-yi

    #computing absolute arrays
    abs_dx, abs_dy = np.abs(dx), np.abs(dy)

    #computing delta_z
    dz = z_cut-zi

    #filtering selected rows
    cond_x = abs_dx <= band_width
    cond_y = (abs_dy <= max_distance) & (y_rotate >= yi)
    cond_z = np.abs(dz) <= band_height
    cond_cone = (tan_ang_tol_rad*abs_dy) >= np.sqrt((dx)**2 + (dz)**2)

    #summing conditions
    cond_sum = cond_x & cond_y & cond_z & cond_cone

    # pushing selected index in the result array
    res_select[id] = idx_cut[cond_sum & (idx_cut!=id)]

~ 27s runtime with 69000 values

Answer 1

有多种方法可以提高代码速度，使用 numpy 数组而不是标准列表进行计算通常总是更快（正如您已经尝试过的那样）。

-其次，因为您的计算彼此独立（计算距离 a 不会影响距离 b 的计算），您可以尝试通过例如以下方式并行化您的代码：超线程。 \

-第三，从我的工业工程背景来看，我非常了解计算距离的必要性。这里的技巧是不要（如您所说）实施条件来不重新计算事物。相反，准确地查看哪些点确实需要重新计算通常会更快！