Mask a 3d array with a 2d mask in numpy

python arrays numpy mask

12,740

Without the loop you could write it as:

field3d_mask[:,:,:] = field2d[np.newaxis,:,:] > 0.3

For example:

field3d_mask_1 = np.zeros(field3d.shape, dtype=bool)
field3d_mask_2 = np.zeros(field3d.shape, dtype=bool)

for t in range(nt):
    field3d_mask_1[t,:,:] = field2d > 0.3

field3d_mask_2[:,:,:] = field2d[np.newaxis,:,:] > 0.3

print((field3d_mask_1 == field3d_mask_2).all())

gives:

True

12,740

Author by

Chiel

I am a scientist that uses mostly C++, Python, and Julia. I have a keen interest in learning how to improve my codes from professional and/or more experienced programmers at SO. In my work it is important to write code that is fast, but also understandable by scientists. Most of our codes run on supercomputers and use Message Passing Interface, OpenMP and CUDA. I work mostly on MicroHH to do three dimensional simulations of turbulent atmospheric flows. I try to improve this code based on what I learn at SO.

Updated on June 06, 2022

Comments

Chiel almost 2 years

I have a 3-dimensional array that I want to mask using a 2-dimensional array that has the same dimensions as the two rightmost of the 3-dimensional array. Is there a way to do this without writing the following loop?

import numpy as np

nx = 2
nt = 4

field3d = np.random.rand(nt, nx, nx)
field2d = np.random.rand(nx, nx)

field3d_mask = np.zeros(field3d.shape, dtype=bool)

for t in range(nt):
    field3d_mask[t,:,:] = field2d > 0.3

field3d = np.ma.array(field3d, mask=field3d_mask)

print field2d
print field3d