Masking a 3D NumPy array --> returns 5D array

Question

I have a huge 3d NumPy array and a same shaped mask array filled with ones or zeros. Let's use this as an example:

arr = np.random.randint(1, 100, size=(4, 4, 4))
mask = np.zeros(shape=(4,4,4), dtype='int')
mask[0,2:3,:] = 1

My goal is to get a same shaped array (in this example (4,4,4)), let's name it new_arr, with the values from arr where mask == 1. All other values in new_arr should get be -2. (not 0)

However, doing this:

new_arr = arr[mask]

results in a (4, 4, 4, 4, 4) shaped array.

So, I've got pretty much 2 question: 1:) How do I get my desired result and 2.) Why does masking a 3d NumPy array this way result in a 5d array.

Quang Hoang · Accepted Answer · 2020-11-25 15:44:01Z

2

I believe you just need np.where:

out = np.where(mask, arr, -2)

to answer your other question

Why does masking a 3d NumPy array this way result in a 5d array.

The idea is that mask has dtype int. So, arr[mask] actually works as slicing, not masking. So a 0 in mask would result in arr[0], which is (4,4). And then you do that for all values in mask, and you get an array with the same dimension of mask whose each value is a (4,4) array.

Another way to fix your code is:

out = np.full(arr.shape, -2)
out[mask.astype(bool)] = arr[mask.astype(bool)]

answered Nov 25, 2020 at 15:44

Quang Hoang

151k11 gold badges63 silver badges86 bronze badges

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

Collectives™ on Stack Overflow

Masking a 3D NumPy array --> returns 5D array

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related