Numpy 使用einsum从5d矩阵中提取对角线(横向向下)

Numpy 使用einsum从5d矩阵中提取对角线(横向向下),numpy,diagonal,numpy-einsum,Numpy,Diagonal,Numpy Einsum,我只用Numpy-einsum提取了一条对角线。在einsum的帮助下,我如何获得像[6,37,68,99]这样的其他对角线 x = np.arange(1, 26 ).reshape(5,5) y = np.arange(26, 51).reshape(5,5) z = np.arange(51, 76).reshape(5,5) t = np.arange(76, 101).reshape(5,5) p = np.arange(101, 126).reshape(5,5) a4

我只用Numpy-einsum提取了一条对角线。在einsum的帮助下,我如何获得像[6,37,68,99]这样的其他对角线

x =  np.arange(1, 26 ).reshape(5,5)
y =  np.arange(26, 51).reshape(5,5)
z =  np.arange(51, 76).reshape(5,5)
t =  np.arange(76, 101).reshape(5,5)
p =  np.arange(101, 126).reshape(5,5)

a4 = np.array([x, y, z, t, p]
提取一条对角线:

>>>np.einsum('iii->i', a4)
>>>[  1  32  63  94 125]
我没有任何使用
einsum
的“简单”解决方案,但使用for循环非常简单:

import numpy as np

# Generation of a 3x3x3 matrix
x =  np.arange(1 , 10).reshape(3,3)
y =  np.arange(11, 20).reshape(3,3)
z =  np.arange(21, 30).reshape(3,3)

M = np.array([x, y, z])

# Generation of the index
I = np.arange(0,len(M))

# Generation of all the possible diagonals
for ii in [1,-1]:
    for jj in [1,-1]:
        print(M[I[::ii],I[::jj],I])

# OUTPUT:
# [ 1 15 29]
# [ 7 15 23]
# [21 15  9]
# [27 15  3]

我们修复最后一个维度的索引,并找到其他维度的所有可能的向后和向前索引组合。

您是否意识到此
einsum
与以下内容相同:

In [64]: a4=np.arange(1,126).reshape(5,5,5)
In [65]: i=np.arange(5)
In [66]: a4[i,i,i]
Out[66]: array([  1,  32,  63,  94, 125])
调整索引以获得其他对角线应该很容易

In [73]: a4[np.arange(4),np.arange(1,5),np.arange(4)]
Out[73]: array([ 6, 37, 68, 99])
“iii->i”生成主对角线与其说是一个设计的特征,不如说是一个愉快的意外。不要试图推它