在 numpy 中用 2d 掩码屏蔽 3d 数组

Question

我有一个 3 维数组，我想使用与 3 维数组最右边的两个维度相同的 2 维数组来屏蔽它。有没有办法不用写下面的循环就可以做到这一点？

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

Answer 1

没有循环你可以这样写：

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

例如：

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())

给出：

True

Answer 2

numpy.broadcast_to（Numpy 1.10.0 中的新功能）：

field3d_mask = np.broadcast_to(field2d > 0.3, field3d.shape)

在 numpy 中用 2d 掩码屏蔽 3d 数组

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

python

arrays

numpy

mask