numpy大数组索引使解释器崩溃

Question

我想引用具有索引i和j的两个数组的矩阵的numpy数组。 我在下面使用的方法效果很好，但是在处理非常大的数组时会使解释器崩溃。 我知道为什么会这样，但是我对numpy来说还太陌生，以至于不知道一种更好的方法。

有什么办法可以使用大型数组有效地实现下面的代码？

import numpy as np
np.set_printoptions(precision=4,suppress=True)

def test(COUNT):
    M = np.random.random_sample((COUNT,4,4,)) # Many matrices
    i = np.random.randint(4, size=COUNT)
    j = np.random.randint(4, size=COUNT)

    # Debug prints
    print M # Print the source matrices for reference
    print i # Print the i indices for reference
    print j # Print the j indices, for reference

    # return the diagonal, this is where the code fails because
    # M[:,i,j] gets incredibly large. This is what i'm trying to solve
    return  M[:,i,j].diagonal() 
    #return np.einsum('ii->i', M[:,i,j])

一些例子：

# test 1 item, easy
print test(1)

[[[ 0.4158  0.2146  0.0371  0.4449]
  [ 0.8894  0.9889  0.0961  0.7343]
  [ 0.8905  0.2062  0.1663  0.04  ]
  [ 0.691   0.1203  0.6524  0.636 ]]]
[1]    
[0]
[ 0.8894]

完美，第一个（也是唯一一个）矩阵的索引[1] [0]为0.884

# test 2 items
print test(2)

[[[ 0.0697  0.434   0.8456  0.592 ]
  [ 0.4413  0.8893  0.9973  0.9184]
  [ 0.7951  0.7392  0.8603  0.8069]
  [ 0.5054  0.3846  0.7708  0.0563]]

 [[ 0.7414  0.2676  0.4796  0.1424]
  [ 0.1203  0.9183  0.1341  0.074 ]
  [ 0.2375  0.3475  0.2298  0.9879]
  [ 0.7814  0.0262  0.4498  0.9864]]]
[2 3]
[1 1]
[ 0.7392  0.0262]

不出所料，第一个矩阵的索引[2] [1]和第二个矩阵的[3] [1]的值为[0.7392 0.0262]，一切都很好！

# too many items!
print test(1000000)

机器停顿是因为M [：，i，j]与所有抛弃值都太大（我只关心对角线）。

我用np.einsum摸索了一下，看是否有帮助。 但这再次对我来说太新了，所以现在我正在寻找一点帮助！ :)

Answer 1

我认为einsum不会为您做任何事情-您只是在用它代替diagonal 。 但是尝试：

M[np.arange(COUNT),i,j]

这应该返回所需的元素，而无需收集额外的内容。

之所以有效，是因为它等效于索引：

M[[0 1], [2 3], [1 1]]

即元素

M[0,2,1] and M[1,3,1]

另一个生成一个(COUNT,COUNT)矩阵，并从中提取对角线(COUNT,)数组。