python - 复杂 Numpy 脚本的矢量化实现

Question

我正在从事一个个人项目，该项目涉及通过雷达预测天气模式的移动。我有三个 n x m numpy 数组；一种具有降水强度值，一种具有该降水在 X 方向上的运动(以像素为单位)，一种具有该降水在 Y 方向上的运动(以像素为单位)。我想使用这三个数组使用其他两个数组中的偏移量来确定降水像素的位置。

xMax = currentReflectivity.shape[0]
yMax = currentReflectivity.shape[1]                            
for x in xrange(currentReflectivity.shape[0]):
    for y in xrange(currentReflectivity.shape[1]):
        targetPixelX = xOffsetArray[x,y] + x
        targetPixelY = yOffsetArray[x,y] + y
        targetPixelX = int(targetPixelX)
        targetPixelY = int(targetPixelY)
        if targetPixelX < xMax and targetPixelY < yMax:
            interpolatedReflectivity[targetPixelX,targetPixelY] = currentReflectivity[x,y]

我想不出一种方法来对其进行矢量化；有什么想法吗？

Answer 1

这是一个使用 broadcasting 的向量化方法-

x_arr = np.arange(currentReflectivity.shape[0])[:,None]
y_arr = np.arange(currentReflectivity.shape[1])

targetPixelX_arr = (xOffsetArray[x_arr, y_arr] + x_arr).astype(int)
targetPixelY_arr = (yOffsetArray[x_arr, y_arr] + y_arr).astype(int)

valid_mask = (targetPixelX_arr < xMax) & (targetPixelY_arr < yMax)

R = targetPixelX_arr[valid_mask]
C = targetPixelY_arr[valid_mask]
interpolatedReflectivity[R,C] = currentReflectivity[valid_mask]

运行时测试

方法-

def org_app(currentReflectivity, xOffsetArray, yOffsetArray):
    m,n = currentReflectivity.shape
    interpolatedReflectivity = np.zeros((m,n))

    xMax = currentReflectivity.shape[0]
    yMax = currentReflectivity.shape[1]                            
    for x in xrange(currentReflectivity.shape[0]):
        for y in xrange(currentReflectivity.shape[1]):
            targetPixelX = xOffsetArray[x,y] + x
            targetPixelY = yOffsetArray[x,y] + y
            targetPixelX = int(targetPixelX)
            targetPixelY = int(targetPixelY)

            if targetPixelX < xMax and targetPixelY < yMax:
                interpolatedReflectivity[targetPixelX,targetPixelY] = \
                currentReflectivity[x,y]
    return interpolatedReflectivity

def broadcasting_app(currentReflectivity, xOffsetArray, yOffsetArray):
    m,n = currentReflectivity.shape
    interpolatedReflectivity = np.zeros((m,n))

    xMax, yMax = m,n        
    x_arr = np.arange(currentReflectivity.shape[0])[:,None]
    y_arr = np.arange(currentReflectivity.shape[1])

    targetPixelX_arr = (xOffsetArray[x_arr, y_arr] + x_arr).astype(int)
    targetPixelY_arr = (yOffsetArray[x_arr, y_arr] + y_arr).astype(int)

    valid_mask = (targetPixelX_arr < xMax) & (targetPixelY_arr < yMax)     
    R = targetPixelX_arr[valid_mask]
    C = targetPixelY_arr[valid_mask]
    interpolatedReflectivity[R,C] = currentReflectivity[valid_mask]
    return interpolatedReflectivity

时间和验证-

In [276]: # Setup inputs
     ...: m,n = 100,110  # currentReflectivity.shape
     ...: max_r = 120  # xOffsetArray's extent
     ...: max_c = 130  # yOffsetArray's extent
     ...: 
     ...: currentReflectivity = np.random.rand(m, n)
     ...: xOffsetArray = np.random.randint(0,max_r,(m, n))
     ...: yOffsetArray = np.random.randint(0,max_c,(m, n))
     ...: 

In [277]: out1 = org_app(currentReflectivity, xOffsetArray, yOffsetArray)
     ...: out2 = broadcasting_app(currentReflectivity, xOffsetArray, yOffsetArray)
     ...: print np.allclose(out1, out2)
     ...: 
True

In [278]: %timeit org_app(currentReflectivity, xOffsetArray, yOffsetArray)
100 loops, best of 3: 6.86 ms per loop

In [279]: %timeit broadcasting_app(currentReflectivity, xOffsetArray, yOffsetArray)
1000 loops, best of 3: 212 µs per loop

In [280]: 6860.0/212        # Speedup number
Out[280]: 32.35849056603774