Computing eigenvalues along particular axes in python

[tobeornottobealttfan]

Hello,

 

I have a numpy array X in python which is of the shape (M,N,N). This means that for every i in range(M), X is a (N,N) matrix. 

Now, I want to diagonalize all such Xs without using for loop. I was wondering where there exists a numpy/ scipy function in python which does this. 

For example, while doing most operations likesum, multiply, mean etc. using numpy functions, there is an argument for axes. Do you know of any function for diagonalization which support the axes argument?

 

Thanks in advance.

 

[hirushaadi]

import numpy as np

# example array of x shape
M = 3
N = 4
X = np.random.rand(M, N, N)

# diagonalize a 2-D matrix here
def diagonalize_matrix(matrix):
    eigenvalues, eigenvectors = np.linalg.eig(matrix)
    diagonal_matrix = np.diag(eigenvalues)
    return eigenvectors @ diagonal_matrix @ np.linalg.inv(eigenvectors)

# diagonalize_matrix each 2-D slice of X along the first axis
diagonalized_X = np.apply_along_axis(diagonalize_matrix, axis=0, arr=X)

this should work

note that the @ operator is used for matrix multiplication