如何使用数组索引张量？

Question

假设我有一个 3D 数组（张量）

print a.shape
(100, 100, 100)

并希望使用一个数组索引其第一个轴：

print i.shape
(20,)

及其最后一个轴使用另一个数组

print j.shape
(30,)

我的目的是获得一个 (20, 100, 30) 数组，例如，用于

a[?!, :, ?!] = b

但是，我想不通。

print a[i.reshape(20, 1, 1), :, j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1, 1), :, j].shape
(20, 1, 30, 100)

print a[i.reshape(20, 1), :, j].shape
(20, 30, 100)

print a[i.reshape(20, 1, 1), j.reshape(1, 1, 30)].shape
(20, 1, 30, 100)

据我了解 "advanced indexing" 规则，第一次尝试应该成功，但我什至没有得到 3D 数组，最后是完整维度 (100) 而不是中间.

Answer 1

方法 #1： 您可以使用 np.ix_ 来获取此类索引网格，并且只需索引到输入数组就可以得到所需的输出。因此，一个实现就像这样 -

a[np.ix_(i,np.arange(a.shape[1]),j)]

方法 #2： 使用广播索引的更简单方法 -

a[i[:,None],:,j].swapaxes(1,2)

用样例验证-

In [24]: a = np.random.randint(0,99,(5,3,5))

In [25]: i = np.random.randint(0,5,(2))

In [26]: j = np.random.randint(0,5,(2))

In [27]: a[i[0],:,j[0]]
Out[27]: array([15,  7, 74])

In [28]: a[i[0],:,j[1]]
Out[28]: array([32, 19, 85])

In [29]: a[i[1],:,j[0]]
Out[29]: array([76, 65, 96])

In [30]: a[i[1],:,j[1]]
Out[30]: array([54, 65, 66])

In [31]: a[np.ix_(i,np.arange(a.shape[1]),j)]
Out[31]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

In [50]: a[i[:,None],:,j].swapaxes(1,2)
Out[50]: 
array([[[15, 32],
        [ 7, 19],
        [74, 85]],

       [[76, 54],
        [65, 65],
        [96, 66]]])

使用索引赋值

对于方法 #1，它很简单 -

a[np.ix_(i,np.arange(a.shape[1]),j)] = b

对于方法 #2，如果 b 是标量，它也应该是直截了当的 -

a[i[:,None],:,j] = b

再次对于方法#2，如果您要分配给形状为 (20,100,30) 的 ndarray b，我们需要在分配之前交换 b 的轴，就像这样 -

a[i[:,None],:,j] = np.swapaxes(b,1,2)

如何使用数组索引张量？

How to index a tensor using arrays?

python

numpy

slice

python-2.7