how to calculate a 2D array with numpy mask

Question

I have a 2 dimension array and based if the value is greater than 0 I want to do a operation (example with x+1). In plain python something like this:

a = [[2,5], [4,0], [0,2]]
for x in range(3):
    for y in range(2):
        if a[x][y] > 0:
            a[x][y] = a[x][y] + 1

Result for a is [[3, 6], [5, 0], [0, 3]]. This is what I want.

Now I want to prevent the nested loop and tried with numpy something like this:

a = np.array([[2,5], [4,0], [0,2]])
mask = (a > 0)
a[mask] + 1

The result is now a 1 dimension and the shape of the array [3 6 5 3]. How can I do this operation and don't loose the dimension like in the plain python example before?

Divakar · Accepted Answer · 2015-06-14 08:26:11Z

2

If a is a numpy array, you can simply do -

a[a>0] +=1

Sample run -

In [335]: a = np.array([[2,5], [4,0], [0,2]])

In [336]: a
Out[336]: 
array([[2, 5],
       [4, 0],
       [0, 2]])

In [337]: a[a>0] +=1

In [338]: a
Out[338]: 
array([[3, 6],
       [5, 0],
       [0, 3]])

answered Jun 14, 2015 at 8:26

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

2 Comments

gustavgans Over a year ago

Thank you! With a mask it would be a[mask] = a[mask]+ 1 right?

Divakar Over a year ago

@gustavgans Here a>0 is the mask. So it's equivalent to a[mask] += 1. That += is an in-situ change. Another more explicit way to put it would be as you did: a[mask] = a[mask]+ 1.

Collectives™ on Stack Overflow

how to calculate a 2D array with numpy mask

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related