How to apply a function to each element of a matrix?

Viewed 113

For each element of a Matrix I want to count the sum of the k-closest cells. From this function

def sum_neighbors(a, radius, row_number, column_number):
    A = [[a[i][j] if  i >= 0 and i < len(a) and j >= 0 and j < len(a[0]) else 0
                for j in range(column_number-1-radius, column_number+radius)]
                    for i in range(row_number-1-radius, row_number+radius)]
    return np.sum(A)

I can do this. For instance, If I have the following matrix

A
array([[1, 8, 2, 2, 8, 2, 4, 2, 6],
       [3, 6, 9, 9, 1, 8, 0, 3, 9],
       [7, 4, 2, 7, 8, 5, 8, 8, 1],
       [8, 3, 2, 3, 6, 9, 9, 1, 6],
       [1, 1, 3, 5, 6, 2, 9, 1, 8],
       [7, 0, 8, 7, 9, 8, 4, 4, 6],
       [1, 8, 5, 4, 7, 4, 8, 7, 0],
       [2, 8, 0, 3, 5, 6, 4, 3, 9],
       [1, 5, 1, 9, 9, 4, 5, 7, 9]])

I can run the following to get, for each values, the sum of the all the values a radius of 2 cells.

s = np.shape(A)
As = np.zeros(s)
for i in range(0,s[0]):
    for j in range(0,s[1]):
        As[i][j]=neighbors(A, 2, i, j)

As
array([[ 18.,  29.,  40.,  49.,  55.,  45.,  39.,  43.,  34.],
       [ 29.,  42.,  60.,  77.,  81.,  75.,  75.,  73.,  56.],
       [ 40.,  55.,  76.,  99., 104., 104., 103., 104.,  81.],
       [ 42.,  60.,  86., 115., 121., 129., 126., 130., 101.],
       [ 40.,  64.,  95., 125., 131., 147., 140., 139., 109.],
       [ 40.,  60.,  86., 122., 126., 148., 149., 144., 108.],
       [ 39.,  57.,  79., 112., 122., 136., 134., 141., 108.],
       [ 34.,  51.,  79., 115., 127., 135., 140., 144., 108.],
       [ 32.,  46.,  69.,  99., 110., 110., 117., 118.,  88.]])

However, for large matrices the loop takes infinite time. Is there a way to apply the function to each element of the matrix without a loop?

1 Answers

Problem

The way to deal with the border issue is to pad the array with k zeros. Since you are summing adding these values to your array has no effect -- it simply makes life easier. Consider the top-left corner of your original array:

1 8 2 ...
3 6 9 ...
7 4 2 ...
.     .
.      .
.       .

You want to look a distance k = 2 in every direction from the 1 there at the very corner, for instance. By padding A with k = 2 zeros, we get

0 0 0 0 0 ...
0 0 0 0 0 ...
0 0 1 8 2 ...
0 0 3 6 9 ...
0 0 7 4 2 ...
.         .
.          .
.           .

The padding helps with the borders, but the trick to this process is taking advantage of numpy's stride_tricks sub-module. We are essentially performing a 2D convolution, and yes, you could use scipy for this, if you happen to have it installed and it's already a dependency on your project.

Solution

from numpy.lib.stride_tricks import sliding_window_view as windows

k = 2  # radius; k >= 1
w = 2*k + 1  # window size; 1x1, 3x3, 5x5, ...

# A = np.array(...)  # from question
neighbor_sums = windows(np.pad(A, k), (w, w)).sum(axis=(2, 3))

Output for k = 1:

array([[18, 29, 36, 31, 30, 23, 19, 24, 20],
       [29, 42, 49, 48, 50, 44, 40, 41, 29],
       [31, 44, 45, 47, 56, 54, 51, 45, 28],
       [24, 31, 30, 42, 51, 62, 52, 51, 25],
       [20, 33, 32, 49, 55, 62, 47, 48, 26],
       [18, 34, 41, 54, 52, 57, 47, 47, 26],
       [26, 39, 43, 48, 53, 55, 48, 45, 29],
       [25, 31, 43, 43, 51, 52, 48, 52, 35],
       [16, 17, 26, 27, 36, 33, 29, 37, 28]])

For k = 2:

array([[ 42,  60,  77,  81,  75,  75,  73,  56,  41],
       [ 55,  76,  99, 104, 104, 103, 104,  81,  57],
       [ 60,  86, 115, 121, 129, 126, 130, 101,  75],
       [ 64,  95, 125, 131, 147, 140, 139, 109,  77],
       [ 60,  86, 122, 126, 148, 149, 144, 108,  80],
       [ 57,  79, 112, 122, 136, 134, 141, 108,  79],
       [ 51,  79, 115, 127, 135, 140, 144, 108,  84],
       [ 46,  69,  99, 110, 110, 117, 118,  88,  66],
       [ 31,  47,  68,  78,  74,  85,  87,  66,  52]])

With scipy

If you're already using scipy as a dependency in your project, you may as well take advantage of its specialized convolve2d() function:

k = 2  # radius; k >= 1
w = 2*k + 1  # window size; 1x1, 3x3, 5x5, ...

neighbor_sums = scipy.signal.convolve2d(A, np.ones((w, w), int), mode='same')

Performance

A quick benchmark shows that scipy's convolve2D() function is indeed faster:

In [18]: ones = np.ones((w, w), int)

In [19]: padded_A = np.pad(A, k)

In [20]: %timeit convolve2d(A, ones, mode='same')
9.79 µs ± 786 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [21]: %timeit windows(padded_A, (w, w)).sum(axis=(2, 3))
34.5 µs ± 6.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Related