I only managed to extract one diagonal using Numpy einsum. How do I get the other diagonals like [6, 37, 68, 99] with help of einsum?
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]
Extracting one diagonal:
>>>np.einsum('iii->i', a4)
>>>[ 1 32 63 94 125]
I don't have any "easy" solution using einsum
but it is quite simple with a for loop:
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]
We fix the index of the last dimension and we find all the possible combinations of backward and forward indexing for the other dimensions.