Numpy fastest way to create mask array around indexes

Question

I have a 1D array:

arr = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0, 0, 0, 0, 0, 2, 3, 0, 0, 1, ...], dtype='uint16')

I want to create a mask array that is True anywhere that is +/- N indexes of of a value greater than 2, yielding the following (for N=3)

mask = [F, T, T, T, T, T, T, T, F, F, T, T, T, T, T, T, T, F, F, F, T, T, T, T, T, T, T, ...]

(note, I used T/F for readability's sake)

I need this to be rather fast as my actual array is millions of points long, N will likely be something like 500.

Edit: similar problem

Note that arr = [0, 0, ...] is not an array but a list. Did you mean a list or it is really a Numpy array? If so, what is its dtype? — Jérôme Richard
– Jérôme Richard, Commented May 19, 2022 at 19:43
@JérômeRichard thanks, it's an ndarray with dtype=uint16 — joshp
– joshp, Commented May 19, 2022 at 19:45

Z Li · Accepted Answer · 2022-05-19 20:49:39Z

3

Find the elements larger than 2 then set the elements around them to True:

a = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0])

N = 3
mask = a > 2
center = np.where(mask)[0]
mask[np.maximum(np.ravel(center - np.arange(1, 1 + N).reshape(-1, 1)), 0)] = True
mask[np.minimum(np.ravel(center + np.arange(1, 1 + N).reshape(-1, 1)), len(a)-1)] = True

Thanks @Michael Szczesny for pointing out the edge case. The maximum and minimum ensures that the index would not (unintentionally) go out of bounds.

edited May 19, 2022 at 20:49

answered May 19, 2022 at 19:54

Z Li

4,3381 gold badge7 silver badges21 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Z Li Over a year ago

Yes you are right. I have updated the answer.

mozway · Accepted Answer · 2022-05-19 20:03:38Z

1

You can use a convolution with numpy.convolve and an array of ones of length 2*N+1 (N on each side plus self) as mask:

np.convolve(arr>2, np.ones(2*N+1), mode='same').astyoe(bool)

Output:

array([False,  True,  True,  True,  True,  True,  True,  True, False,
       False,  True,  True,  True,  True,  True,  True,  True, False,
       False, False,  True,  True,  True,  True,  True,  True,  True])

answered May 19, 2022 at 20:03

mozway

267k13 gold badges56 silver badges106 bronze badges

4 Comments

paime Over a year ago

It is also to be noted that if you don't cast to bool then you even have a count of the number of neighbors.

joshp Over a year ago

convolve is much slower... (my mistake), thank you both!!! also, there is a typeo astyoe

mozway Over a year ago

@paime yes, I love convolutions, so powerful ;)

Ben Grossmann Over a year ago

@joshp You might get faster performance by modifying the np.ones call to np.ones(2*N+1,dtype='int') or np.ones(2*N+1,dtype='bool')

Collectives™ on Stack Overflow

Numpy fastest way to create mask array around indexes

2 Answers 2

1 Comment

4 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

4 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related