如何利用3D掩码将4D numpy数组对应维度元素置0？

直接利用numpy的广播机制和布尔索引就能搞定，核心思路是把3D掩码扩展成和4D数组匹配的维度，再进行赋值操作：

方法1：直接修改原数组 #

import numpy as np

# 假设D是形状为(nx, ny, nz, nt)的4D数组，M是形状为(nx, ny, nz)的3D掩码
# 给掩码M增加一个维度，使其与D的最后一个轴匹配（自动广播到nt长度）
D[~M[..., np.newaxis]] = 0

这里M[..., np.newaxis]会把M的形状从(nx, ny, nz)变成(nx, ny, nz, 1)，numpy会自动将其广播到和D相同的(nx, ny, nz, nt)形状。~符号用来取反，筛选出所有M为False的位置，然后把D对应位置的元素设为0。

方法2：生成新数组（不修改原数组） #

如果不想改动原数组D，可以用np.where生成新数组：

D_new = np.where(M[..., np.newaxis], D, 0)

这个方法会保留M为True位置的D元素，其余位置设为0，返回一个新数组。

内容的提问来源于stack exchange，提问作者Omroth