Python – compute average absolute difference of element and neighbors in NumPy array

I’m looking for a way to calculate the average absolute difference between neighboring elements in a NumPy array. Namely, given an array like

[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9]]

The value for the middle square will be 2.5 (aka (4+3+2+1+1+2+3+4)/8). I know with SciPy’s correlate2d you can compute the average difference, but, as far as I know, not the average absolute difference (i.e. for the example above, correlate2d would give 0 – (-4+-3+-2+-1+1+2+3+4)/8 – not 2.5).

Is there a fast way to do this in Python? I don’t want to iterate over the elements since this will be running for very large arrays many times.

Answer

Why not write it by hand and numba.jit the result?

arr = (np.arange(10**6)+1).reshape(10**3,10**3)

@nb.njit
def get_avg_abs_diff(arr):
    n,m = arr.shape
    out = np.empty((n,m))
    for i in range(n):
        for j in range(m):
            a = max(0,i-1)
            b = i+2
            c = max(0,j-1)
            d = j+2
            neighbours = arr[a:b,c:d]
            out[i,j] = np.sum(np.abs(neighbours-arr[i,j]))/(neighbours.shape[0]*neighbours.shape[1]-1)
    return out

get_avg_abs_diff(arr)

This takes about 160ms for the 1000 by 1000 array (down from non jited version which takes 10.6s).

For your array I get

array([[2.66666667, 2.2       , 2.        ],
       [2.6       , 2.5       , 2.6       ],
       [2.        , 2.2       , 2.66666667]])