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,)數組。