TheDemonTelevision
TheDemonTelevision

Reputation: 31

Sum in square inside np.array

I have variable:

Q = np.asarray([
   [0.20, 0.  , 0.3 , 0.1 ],
   [0.10, 0.  , 0.2 , 0.  ],
   [0.05, 0.  , 0.  , 0.  ],
   [0.  , 0.  , 0.  , 0.05],
], dtype=float)

And I want for each element to find a sum of itself and upper, upper-left, left elements.
If there is no left or upper-left or upper elements we need to make the sum of what we have.
In return I need to get J with these sums.
For example:J[3][3] = Q[3][3]+Q[3][2]+Q[2][2]+Q[2][3]
J[0][1] = Q[0][1]+Q[0][0]


Can I do this via numpy functions only? I want it to be very fast, so NO cycles.
Do you have any ideas?

UPD: The resulting array should be the same size as Q

J = np.array([
   [0.2 , 0.2 , 0.3 , 0.4 ],
   [0.3 , 0.3 , 0.5  , 0.6 ],
   [0.15, 0.15, 0.2 , 0.2 ],
   [0.05 , 0.05 , 0.  , 0.05 ],
])

Upvotes: 1

Views: 263

Answers (2)

Jan Christoph Terasa
Jan Christoph Terasa

Reputation: 5935

To sum up all elements in a sliding 2x2 window, you can use scipy.signal.convolve2d to do a 2d-convolution of your array with a unit kernel of shape (2, 2):

import numpy as np
from scipy.signal import convolve2d

kernel = np.ones((2, 2))  

Q = np.asarray([
   [0.20, 0.  , 0.3 , 0.1 ],
   [0.10, 0.  , 0.2 , 0.  ],
   [0.05, 0.  , 0.  , 0.  ],
   [0.  , 0.  , 0.  , 0.05],
], dtype=float)

result = convolve2d(Q, kernel, 'same')
# array([[0.2 , 0.2 , 0.3 , 0.4 ],
#        [0.3 , 0.3 , 0.5 , 0.6 ],
#        [0.15, 0.15, 0.2 , 0.2 ],
#        [0.05, 0.05, 0.  , 0.05]])

Upvotes: 4

mozway
mozway

Reputation: 260500

Just for info, as @JanChristoph's answer is much shorter and more elegant, here it the way to implement it using numpy only:

sub_shape = (2, 2)
view_shape = tuple(np.subtract(Q.shape, sub_shape) + 1) + sub_shape
strides = Q.strides + Q.strides

sub_matrices = np.lib.stride_tricks.as_strided(Q, view_shape, strides)

sub_matrices.sum(axis=(2,3))

output:

array([[0.3 , 0.5 , 0.6 ],
       [0.15, 0.2 , 0.2 ],
       [0.05, 0.  , 0.05]])

Upvotes: 2

Related Questions