如何以矢量化方式多次索引多维数组numpy？

Question

我正在尝试在 numpy 中索引多维数组（4 维）。数组的形状为 (125,125,125,3)。我有 3 个单独的二维索引数组列表。列表的大小分别为 (N,4)、(M,4) 和 (1,4)。 3 个单独的列表代表我尝试索引的 4D 数组中的行、列和深度值。例如考虑以下内容：

ix = [[0,1,2,3],
     [3,4,5,6]]
iy = [[2,3,4,5],
     [5,6,7,8]]
iz = [[1,2,3,4]]

weights.shape = (125,125,125,3)

我想用 ix、iy 和 iz 中行、列和深度索引数组的所有可能组合来索引 weights。例如，如果我取每个索引矩阵的第一行，这意味着我想要 select 行 [0,1,2,3]、列 [2,3,4,5] 和深度值 [1,2,3,4] weights。我一直想 select weights 的第 4 个维度中的所有元素。这意味着我基本上 selecting (4,4,4,3) 切片 weights.

现在，我已经使用以下代码通过循环索引实现了这一点

w = np.empty(shape=(X,Y,Z,4,4,4,weights.ndim-1))
for i in range(X):
    for j in range(Y):
        w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
        w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]

我的最终目标是尽快构建形状为 (N,M,1,4,4,4,3) 的多维数组 w。这部分代码将运行多次，因此如果有一种使用内置 numpy 函数的矢量化方式来执行此操作，那将是理想的。

如果有任何需要澄清的问题，请告诉我。这是我第一次提出有关堆栈溢出的问题，所以如果有任何不清楚或混淆的地方，我深表歉意！

Answer 1

您可以使用索引和广播来实现这一点。

import numpy as np

weights = np.random.rand(125, 125, 125, 3)

ix = np.array([[0,1,2,3], [3,4,5,6]])
iy = np.array([[2,3,4,5], [5,6,7,8]])
iz = np.array([[1,2,3,4]])

X = len(ix)
Y = len(iy)
Z = len(iz)

def compute1(weights):
    w = np.empty(shape=(X, Y, Z, 4, 4, 4, weights.ndim-1))
    for i in range(X):
        for j in range(Y):
            w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
            w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]
    return w

def compute2(weights):
    return weights[ix[:, None, None, :, None, None], iy[None, :, None, None, :, None], iz[None, None, :, None, None, :]]

print(np.allclose(compute1(weights), compute2(weights)))

给出 True.

基准测试 -

%timeit compute1(weights)
%timeit compute2(weights)

给出-

36.7 µs ± 897 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
6.28 µs ± 62.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

如您所见，对于这种大小的数据，广播解决方案的速度提高了大约 6 倍。

如何以矢量化方式多次索引多维数组numpy？

How to index multidimensional array multiple times in a vectorized way numpy?

python

numpy

numpy-indexing

numpy-ndarray