Numpy多维索引

Question

我有两个numpy数组- arr1和arr2。Arr2包含arr1的索引值。arr1的形状为(100,8,96,192)，arr2的形状为(8,96,192)。我想做的是用arr2索引arr1，然后得到之后的20个点。

对于上下文，arr1为时间x模型x纬度x长度，并且arr2中的所有索引值对应于arr1数组中的一个时间点。我想要得到arr1在arr2时间点的值，然后在那之后的二十个点。

示例数据

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

in: print(arr1)
out: array([[[[0.61718651, 0.24426295, 0.9165573 , ..., 0.24155022,
0.22327592, 0.9533857 ],
[0.21922781, 0.87948651, 0.926359  , ..., 0.64281931,
...,
[0.09342961, 0.29533331, 0.11398662, ..., 0.36239606,
0.40228814, 0.87284515]]]])

in: print(arr2)
out: array([[[22,  5, 64, ...,  0, 37,  6],
[71, 48, 33, ...,  8, 38, 32],
[15, 41, 61, ..., 56, 32, 48],
...,
...,
[66, 31, 32, ...,  0, 10,  6],
[ 9, 28, 72, ..., 71, 29, 34],
[65, 22, 50, ..., 58, 49, 35]]])

Answer 1

根据评论推断;您可以在一个操作中混合使用高级索引和广播来完成此操作:

import numpy as np
arr1 = np.random.rand(100, 8, 96, 192)
arr2 = np.random.randint(low=0, high=80, size=(8, 96, 192))
H, I, J, K = np.indices((20, 8, 96, 192), sparse=True)
out = arr1[H + arr2, I, J, K].mean(axis=0)