sum of squares (4 values shaping a square) within a 2d numpy array. Python

Question

I am looking to sum each 4 point combination in a 2d array that forms squares within the matrix

in4x4 = np.array(([1,2,3,4],[2,3,4,1],[3,4,1,2],[4,3,1,2]))
print(in4x4)
array([[1, 2, 3, 4],
       [2, 3, 4, 1],
       [3, 4, 1, 2],
       [4, 3, 1, 2]])

Expected output:

print(out3x3)
array([[ 8, 12, 12],
       [12, 12,  8],
       [14,  9,  6]]

Currently I am using numpy.diff in a several step process. Feel like there should be a cleaner approach. Example of the calculation I currently do:

diff3x4 = np.diff(a4x4,axis=0)
a3x4 = a4x4[0:3,:] * 2 + d3x4
d3x3 = np.diff(a3x4)
a3x3 = a3x4[:,0:3] * 2 + d3x3

Is there a clean possibly vectorized approach? Speed is of concern. Looked at scipy.convolve but does not seem to suit my purposes, but maybe I am just configuring the kernel wrong.

Answer 1

You can do 2-D convolution with a kerenl of ones to achieve the desired result.

Simple code that implements 2D-convolution and test on the input you gave:

import numpy as np

def conv2d(a, f):
    s = f.shape + tuple(np.subtract(a.shape, f.shape) + 1)
    strd = np.lib.stride_tricks.as_strided
    subM = strd(a, shape = s, strides = a.strides * 2)
    return np.einsum('ij,ijkl->kl', f, subM)

in4x4 = np.array(([1,2,3,4],[2,3,4,1],[3,4,1,2],[4,3,1,2]))
k = np.ones((2,2))

print(conv2d(in4x4, k))

output:

[[ 8. 12. 12.]
 [12. 12.  8.]
 [14.  9.  6.]]

In case you can use built-in function you can use like in the following:

import numpy as np
from scipy import signal
in4x4 = np.array(([1,2,3,4],[2,3,4,1],[3,4,1,2],[4,3,1,2]))
k = np.ones((2,2))
signal.convolve2d(in4x4, k, 'valid')

Which output:

array([[ 8., 12., 12.],
       [12., 12.,  8.],
       [14.,  9.,  6.]])

signal.convolve2d(in4x4, k, 'valid')