如何在多维Numpy数组中按索引将后续值设为np.nan

需求说明 #

现有两个Numpy数组：

• arr1：形状为(100, 8, 96, 192)，维度对应时间×模式×纬度×经度
• arr2：形状为(8, 96, 192)，每个元素对应arr1时间维度的索引位置

需要实现：将arr1中每个(模式, 纬度, 经度)位置下，时间索引大于arr2对应值的所有元素设为np.nan。

示例数据 #

import numpy as np

arr1 = np.random.rand(*(100, 8, 96, 192))
arr2 = np.random.randint(low=0, high=80, size=(8, 96, 192))

尝试的错误代码及报错 #

用户尝试以下代码：

arr1 = np.random.rand(100, 8, 96, 192)
arr2 = np.random.randint(low=0, high=80, size=(8, 96, 192))
I, J, K = np.indices((8, 96, 192), sparse=True)
out = arr1[arr2:, I, J, K]

报错信息：

TypeError: 只有整数标量数组可以转换为标量索引

错误原因 #

原代码的索引方式不符合Numpy的索引规则：

• arr2:, I, J, K试图将多维数组arr2作为时间维度的切片起始点，但Numpy不支持直接用多维数组进行这种切片操作
• 稀疏索引I, J, K的使用方式错误，无法和arr2的维度匹配完成正确的索引定位

解决方案 #

利用Numpy的广播机制生成掩码，直接定位需要设为NaN的位置，高效且简洁：

import numpy as np

# 示例数据
arr1 = np.random.rand(*(100, 8, 96, 192))
arr2 = np.random.randint(low=0, high=80, size=(8, 96, 192))

# 创建时间维度的索引数组，扩展维度以匹配arr2的广播要求
time_indices = np.arange(arr1.shape[0])[:, np.newaxis, np.newaxis, np.newaxis]

# 生成掩码：时间索引大于arr2对应位置的元素标记为True
mask = time_indices > arr2

# 将掩码对应的位置设为np.nan
arr1[mask] = np.nan

代码解释 #

1. time_indices：生成形状为(100, 1, 1, 1)的数组，代表时间维度的所有索引
2. 广播匹配：time_indices和arr2广播后，得到形状为(100, 8, 96, 192)的掩码数组，每个位置标记该时间点是否需要设为NaN
3. 赋值操作：直接通过掩码索引arr1，将符合条件的元素赋值为np.nan

这种方法避免了生成庞大的索引数组，利用Numpy的广播特性实现高效计算，适合处理大尺寸数组。

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