Sort 3D eigenvector array using eigenvalues

Question

I am trying to re-order a tensor of eigenvalues and eigenvectors in numpy in descending order. I have partially used this answer to do so.
I am now trying to do the ordering on an array of eigenvalues of shape (32448,3) and eigenvectors of shape (32448,3,3) respectively. Here is the code that I am currently using:
# compute eigenvalues and eigenvectors using numpy
eigenvalues, eigenvectors = eigh(struct)

# sort eigenvalues and eigenvectors in descending order of eigenvalues
idx = np.argsort(eigenvalues, axis=1)[:,::-1]

print(idx.shape)
# (32448, 3)

eigenvalues = eigenvalues[:,idx]
eigenvectors = eigenvectors[:,:,idx]

Computing the indexes goes fine, however as soon as the line eigenvalues = eigenvalues[:,idx] executes, my program pretty much breaks; it hangs up and I have to terminate it.

yatu · Answer

It looks like your input array must be of shape (N,M,M), so you have N eigenvector matrices of with M eigenvectors each. To order them according to the eigenvalues, you'll need to leverage advanced indexing from the result of taking argsort on the eigenvalue array:
import numpy as np
import numpy.linalg as linalg

A = np.random.random((3,2,2))
eigenValues, eigenVectors = linalg.eig(A)

eigenValues
array([[ 0.93101687,  0.07290812],
       [ 0.00375849,  1.24016112],
       [ 1.61859044, -0.08976086]])

eigenVectors
array([[[ 0.93211183, -0.52806487],
        [ 0.36217059,  0.84920403]],

[[-0.86214688, -0.7177663 ],
        [ 0.50665843, -0.69628409]],

[[ 0.71581294, -0.62592751],
        [ 0.69829208,  0.77988124]]])

Find the indices that sort the eigenValues along the second axis:
eig_s = eigenValues.argsort(1)[:,::-1]
eig_s
array([[0, 1],
       [1, 0],
       [0, 1]], dtype=int64)

And use them to index the rows in eigenVectors:
np.take_along_axis(eigenVectors, eig_s[...,None], 1)
array([[[ 0.93211183, -0.52806487],
        [ 0.36217059,  0.84920403]],

[[ 0.50665843, -0.69628409],
        [-0.86214688, -0.7177663 ]],

[[ 0.71581294, -0.62592751],
        [ 0.69829208,  0.77988124]]])

Similarly for the eigenvalues:
np.take_along_axis(eigenValues, eig_s, axis=1)
array([[ 0.93101687,  0.07290812],
       [ 1.24016112,  0.00375849],
       [ 1.61859044, -0.08976086]])

Sort 3D eigenvector array using eigenvalues

One Answer

Add your own answers!

Ask a Question