在numpy中使用2D掩码对3D数组进行掩码处理

Question

我有一个三维数组，我想使用一个与三维数组最右边的两个维度相同的二维数组进行掩码处理。是否有一种方法可以在不编写以下循环的情况下完成此操作？

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

Answer 1

在编程中有numpy.broadcast_to（Numpy 1.10.0中新增）：

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

- user2379410

---

这个答案看起来比被接受的答案更好（更像是numpythonic？）。广播通常不涉及复制数据n次。它只是通过适当设置步幅来模拟它。 - Mad Physicist

实际上，@Bart答案中使用numpy.newaxis选项现在会发出警告：VisibleDeprecationWarning: boolean index did not match indexed array along dimension 2; dimension is 3 but corresponding boolean dimension is 1。 - Benjamin

@Benjamin；你使用的是Python/Numpy哪个版本？我已经使用了Python 3.6.1和Numpy 1.13.0里的'np.newaxis'很长时间了，没有出现任何警告。 - Bart

如果有多个广播选项会发生什么？例如，所有维度大小均为8时会怎样？ - Gulzar

这种方法并不总是有效的。例如，在上面的例子中，如果 field3d = np.random.rand(nx, nx, 4)（注意 nx 和 nt=4 的顺序不同），则 np.broadcast_to() 将会发出一个 ValueError: "operands could not be broadcast..."。在图像分析中，将二进制掩码应用于具有多个通道（例如 RGBA）的图像是一项常见任务。 - normanius

Answer 2

没有循环，你可以这样写：

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

输出：

真

- Bart

---

field3d_mask[...] = (field2d > 0.3)[np.newaxis, ...] - Mad Physicist