[英]Indexing and retrieving numpy rows from a multidimensional array
我有一個形狀的源多維數組(a,b,c,c,d)
存儲大小為d
的向量/數據,以及另一個形狀數組(a,b,e,2)
存儲大小為 2 的e
索引。二維值對應於數據數組的索引 2-3(大小為c
的兩個維度)。 請注意,兩個 arrays 共享相同a,b
尺寸大小。
我想要做的是使用這些索引從第一個數組中檢索大小為d
的行。 因此,output 數組的大小應為(a,b,e,d)
,即沿a,b
維度的大小為d
的e
向量。
a, b, c, d = 3,5,7,9
e = 11
data = np.random.rand(a,b,c,c,d)
inds = np.random.randint(0,c, size=(a,b,e,2))
res = data[:, :, inds[:,:,:,0], inds[:,:,:,1],:]
print(' - Obtained shape:', res.shape)
print(' - Desired shape:', (a,b,e,d))
# - Obtained shape: (3, 5, 3, 5, 11, 9)
# - Desired shape: (3, 5, 11, 9)
我現在能想到的唯一方法是通過在所有三個主要維度中生成類似范圍的索引來強制執行完整的索引:
import numpy as np
rng = np.random.default_rng()
a, b, c, d = 3, 5, 7, 9
e = 11
data = rng.uniform(size=(a, b, c, c, d))
inds = rng.integers(0, c, size=(a, b, e, 2))
# generate open index meshes to reduce memory for at least here
aind, bind, _, = np.ogrid[:a, :b, :e]
res = data[aind, bind, inds[..., 0], inds[..., 1], :]
print(' - Obtained shape:', res.shape)
print(' - Desired shape:', (a, b, e, d))
隨機檢查以查看值是否正確:
sample_index_pos = (1, 1, 8) # <-> (a, b, e)
c_inds = inds[sample_index_pos] # <-> (c, c)
expected = data[sample_index_pos[:2] + tuple(c_inds)]
have = res[sample_index_pos]
print(np.array_equal(expected, have))
# True
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.