如何用Numpy数组arr2索引arr1并获取后续20个时间点数据？

问题描述 #

我有两个Numpy数组arr1和arr2，arr2包含arr1的索引值：

• arr1形状为(100, 8, 96, 192)，对应维度为 时间×模式×纬度×经度
• arr2形状为(8, 96, 192)，其中的数值是arr1时间维度的索引

需要实现：通过arr2索引arr1，获取每个位置指定时间点的数据，同时获取该时间点之后的20个数据点。

示例数据 #

import numpy as np

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

解决方案 #

步骤说明 #

1. 修正索引范围：避免后续索引越界，因为要取当前时间点+20个后续点（共21个数据），所以arr2中的最大索引需限制为100 - 21 = 79
2. 构造时间偏移量：生成0到20的整数，对应原时间点到之后第20个点的偏移
3. 广播索引匹配维度：将偏移量和arr2的维度对齐，生成每个位置的时间索引数组
4. 花式索引取数：结合模式、纬度、经度的维度索引，从arr1中提取目标数据

代码实现 #

# 修正索引范围，防止越界
arr2_clipped = np.clip(arr2, 0, 100 - 21)  # 确保最大索引+20 <= 99（arr1时间维度的最大索引）

# 构造时间偏移数组，扩展维度以匹配arr2的广播要求
time_offsets = np.arange(21)[:, np.newaxis, np.newaxis, np.newaxis]  # 形状(21,1,1,1)

# 计算每个位置的完整时间索引：原索引+偏移量，形状(21,8,96,192)
time_indices = arr2_clipped + time_offsets

# 构造模式、纬度、经度的维度索引，用于花式索引
mode_idx = np.arange(8)[np.newaxis, :, np.newaxis, np.newaxis]    # 形状(1,8,1,1)
lat_idx = np.arange(96)[np.newaxis, np.newaxis, :, np.newaxis]    # 形状(1,1,96,1)
lon_idx = np.arange(192)[np.newaxis, np.newaxis, np.newaxis, :]    # 形状(1,1,1,192)

# 提取目标数据，结果形状为(21,8,96,192)
result = arr1[time_indices, mode_idx, lat_idx, lon_idx]

结果说明 #

• result的第一个维度对应时间偏移：索引0是arr2指定的原时间点，索引1到20是该时间点之后的20个数据点
• 若只需要后续20个点（不包含原时间点），只需将time_offsets改为np.arange(1,21)，并调整arr2_clipped的上限为100 - 20 = 80

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