pythonarraysnumpy

Create new matrix with sum of all n*n submatrixes sum without using loops


I have a starting array (a) and I'm using loops to create a new matrix (b) where every element represents the sum of every unique subarray n*n (0+1+5+6=12, 1+2+6+7=16, ...). Can I achieve the same result without using loops by inbuilt Numpy methods? I tried np.reshape, but I can't count intersecting subarrays with it.

Code using loops:

import numpy as np
a = np.array([[ 0,  1,  2,  3, 4], [5,  6,  7,  8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19], [15, 16, 17, 18, 20]])
n = 2

b = np.empty([0, np.size(a[0])-n+1], int)
for row in range(0, np.size(a[0])-n+1):
    temp_list = []
    for col in range(0, np.size(a[0])-n+1):
        temp_list.append(np.sum(a[row:row+n, col:col+n]))
    b = np.vstack([b, [np.array(temp_list)]])
    
print('Input (a):')
print(a)
print('Output (b):')
print(b)

Result:

Input (a):
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [15 16 17 18 20]]
Output (b):
[[12 16 20 24]
 [32 36 40 44]
 [52 56 60 64]
 [62 66 70 75]]

Solution

  • This is essentially a 2d convolution with a square filter:

    import numpy as np
    from scipy.signal import convolve2d
    
    a = np.array([[ 0,  1,  2,  3, 4], [5,  6,  7,  8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19], [15, 16, 17, 18, 20]])
    n = 2
    
    f = np.ones((n,n))
    
    convolve2d(a, f, mode='valid')
    

    output:

    array([[12., 16., 20., 24.],
           [32., 36., 40., 44.],
           [52., 56., 60., 64.],
           [62., 66., 70., 75.]])
    

    If you care about integer output, just use that instead f = np.ones((n,n), dtype=int)