在numpy中使用2d掩码掩盖3d数组

时间:2016-06-07 14:35:49

标签: python arrays numpy mask

我有一个三维数组,我想使用二维数组进行掩码,该数组的尺寸与三维数组中最右边的两个相同。有没有办法在不编写以下循环的情况下执行此操作?

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

2 个答案:

答案 0 :(得分:12)

numpy.broadcast_to(Numpy 1.10.0中的新内容):

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

答案 1 :(得分:8)

没有循环,你可以把它写成:

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

给出: