rpb
rpb

Reputation: 3299

Efficiently assign a value within predefined range in Numpy

The objective is to assign new value within certain range (b_top,b_low).

The code below able to achieve the intended objective

b_top=np.array([1,7])
b_low=np.array([3,9])+1
Mask=np.zeros((1,11), dtype=bool)
for x,y in zip(b_top,b_low):
    Mask[0,x:y]=True

However, I wonder there is single line approach, or more efficient way of doing this?

Upvotes: 1

Views: 384

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114440

You can turn b_top and b_low into a mask using np.cumsum and the fact that bool and int8 are the same itemsize.

header = np.zeros(M.shape[1], np.uint8)
header[b_top] = 1
header[b_low if b_low[-1] < header.size else b_low[:-1]] = -1
header.cumsum(out=Mask[0].view(np.int8))

I've implemented this function in a little utility library I made. The function is called haggis.math.runs2mask. You would call it as

from haggis.math import runs2mask

Mask[0] = runs2mask(np.stack((b_top, b_low), -1), Mask.shape[1])

Upvotes: 1

Related Questions