24

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
Chiel
  • 6,006
  • 2
  • 32
  • 57

2 Answers2

22

There's numpy.broadcast_to (new in Numpy 1.10.0):

field3d_mask = np.broadcast_to(field2d > 0.3, field3d.shape)
  • 1
    This seems much better (more numpythonic?) than the accepted answer. A broadcast usually does not involve copying the data n-time. It just simulates it by setting a stride appropriately. – Mad Physicist Jun 30 '16 at 17:22
  • Yes, in fact the option with `numpy.newaxis` in @Bart's answer now warns: `VisibleDeprecationWarning: boolean index did not match indexed array along dimension 2; dimension is 3 but corresponding boolean dimension is 1` – Benjamin Nov 24 '16 at 19:49
  • @Benjamin; with which Python/Numpy version? I've been using the `np.newaxis` for months++ now with Python 3.6.1 and Numpy 1.13.0, without any warning. – Bart Jul 03 '17 at 21:34
  • 3
    What happens if there is more than one broadcasting option? for example all dimensions are of size 8? – Gulzar May 26 '20 at 10:51
  • 1
    This doesn't always work. For instance, if `field3d = np.random.rand(nx, nx, 4)` in the above example (note the different order of `nx` and `nt=4`) `np.broadcast_to()` issues a ValueError: "operands could not be broadcast..." It's a common task in image analysis to apply a binary mask to an image with multiple channels (RGBA for example). – normanius May 26 '21 at 17:48
14

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

Bart
  • 9,825
  • 5
  • 47
  • 73