Subsetting 2D array based on condition in numpy python

Question

I have a numpy 2D array of size 3600 * 7200. I have another array of same shape which I want to use as a mask.

The problem is that when I do something like this:

import numpy as np
N = 10
arr_a = np.random.random((N,N))
arr_b = np.random.random((N,N))
arr_a[arr_b > 0.0]

The resulting array is no longer 2D, it is 1D. How do I get a 2D array in return?

How do you want to fill the values that are not greater than 0? With nan? If you don't fill them what kind of a shape are you expecting? What if only one of them is greater than 0? — user2285236
– user2285236, Commented May 12, 2016 at 5:13

score 2 · Accepted Answer · 2016-05-12 05:24:18Z

2

You can use np.where to preserve the shape:

np.where(arr_b > 0.0, arr_a, np.nan)

It will take the corresponding values from arr_a when arr_b's value is greater than 0, otherwise it will use np.nan.

import numpy as np
N = 5
arr_a = np.random.randn(N,N)
arr_b = np.random.randn(N,N)
np.where(arr_b > 0.0, arr_a, np.nan)

Out[107]: 
array([[ 0.5743081 ,         nan, -1.69559034,         nan,  0.4987268 ],
       [ 0.33038264,         nan, -0.27151598,         nan, -0.73145628],
       [        nan,  0.46741932,  0.61225086,         nan,  1.08327459],
       [        nan, -1.20244926,  1.5834266 , -0.04675223, -1.14904974],
       [        nan,  1.20307104, -0.86777899,         nan,         nan]])

Collectives™ on Stack Overflow

Subsetting 2D array based on condition in numpy python

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related