Search code examples
numpydiagonalnumpy-einsum

extracting diagonals (sideway down) from 5d matrices using einsum


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]

Solution

  • 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.