无循环计算对角减法的numpy有效方法

Question

我有以下 2d numpy 数组：

array([[5, 2, 6],
       [4, 5, 6],
       [7, 12, 3]])

我有一个（梯度）方程，它使用对角线进行如下计算：

(array[i , j] - array[i + 1, j + 1]) **2 + (array[i+1 , j] - array[i, j + 1]) **2

方程很简单。它正在获取每个 2x2 子矩阵，减去对角线，将它们平方，然后添加每个结果。在上面用于第一个 2x2 矩阵的示例中，等式将产生：

(5 - 5) ^ 2 + (4 - 2) ^ 2 = 4

然后，该值将成为新 numpy 数组的第一个元素。通过移动到下一个 2x2 子矩阵，其余元素将以相同的方式填充。

第一个明显的解决方案是在所有行和列上使用嵌套的 for 循环。但是有没有更有效的方法使用numpy（或其他库）来避免for循环。或者，如果不直接避免 for 循环，那么对于较大的 2D 数组，计算效率会更高。

Answer 1

这里有一些想法：

import numpy as np
from scipy.signal import convolve2d
from numba import njit, prange


arr0 = np.array(
    [[5, 2, 6],
     [4, 5, 6],
     [7, 12, 3]])

np.random.seed(0)
arr1 = np.random.randint(0, 20, (20, 20))

def weird_quazi_det_op(arr):
    H, W = arr.shape
    out_arr = np.zeros((H-1, W-1))
    for h in prange(H-1):
        for w in prange(W-1):
            out_arr[h, w] = (arr[h, w] - arr[h+1, w+1])**2 + (arr[h+1, w] - arr[h, w+1])**2
    return out_arr


@njit(parallel=True)
def weird_quazi_det_patallel(arr):
    H, W = arr.shape
    out_arr = np.zeros((H-1, W-1))
    for h in prange(H-1):
        for w in prange(W-1):
            out_arr[h, w] = (arr[h, w] - arr[h+1, w+1])**2 + (arr[h+1, w] - arr[h, w+1])**2
    return out_arr

@njit(fastmath=True)
def weird_quazi_det_fastmath(arr):
    H, W = arr.shape
    out_arr = np.zeros((H-1, W-1))
    for h in prange(H-1):
        for w in prange(W-1):
            out_arr[h, w] = (arr[h, w] - arr[h+1, w+1])**2 + (arr[h+1, w] - arr[h, w+1])**2
    return out_arr


def weird_quazi_det_conv(arr):
    sub0kern = np.array([[1, 0,], [0, -1]])
    sub1kern = np.array([[0, -1,], [1, 0]])

    sub0arr = convolve2d(arr, sub0kern, mode = 'valid')
    sub1arr = convolve2d(arr, sub1kern, mode = 'valid')

    out_arr = sub0arr**2+sub1arr**2
    return out_arr

# init runs
weird_quazi_det_patallel(arr0)
weird_quazi_det_fastmath(arr0)




print('===== op example: =====')
print(' - op code:')
%timeit weird_quazi_det_op(arr0)
print(' - op code + parallel:')
%timeit weird_quazi_det_patallel(arr0)
print(' - op code + fastmath:')
%timeit weird_quazi_det_fastmath(arr0)
print(' - conv solution:')
%timeit weird_quazi_det_conv(arr0)

print('===== larger example: =====')
print(' - op code:')
%timeit weird_quazi_det_op(arr1)
print(' - op code + parallel:')
%timeit weird_quazi_det_patallel(arr1)
print(' - op code + fastmath:')
%timeit weird_quazi_det_fastmath(arr1)
print(' - conv solution:')
%timeit weird_quazi_det_conv(arr1)

输出：

===== op example: =====
 - op code:
The slowest run took 6.15 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 5: 10.6 µs per loop
 - op code + parallel:
The slowest run took 251.69 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 5: 5.71 µs per loop
 - op code + fastmath:
The slowest run took 32.22 times longer than the fastest. This could mean that an intermediate result is being cached.
1000000 loops, best of 5: 678 ns per loop
 - conv solution:
The slowest run took 6.65 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 5: 18.7 µs per loop
===== larger example: =====
 - op code:
1000 loops, best of 5: 714 µs per loop
 - op code + parallel:
The slowest run took 180.35 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 5: 6.33 µs per loop
 - op code + fastmath:
The slowest run took 14.04 times longer than the fastest. This could mean that an intermediate result is being cached.
1000000 loops, best of 5: 1.42 µs per loop
 - conv solution:
10000 loops, best of 5: 41 µs per loop

我认为使用卷积来执行局部对角线减法，然后对差异求和将是最快的解决方案，但似乎没有超过 numba 编译代码的速度，即使使用最简单的实现也是如此。