The function numpy.linalg.eig supports calculating the eigenvalues and eigenvectors of a stack of matrices at once. In my case, I have 500000 2x2 matrices, organised in a 1000x500x2x2 numpy array, and calling numpy.linalg.eig on this returns 1000x500x2 eigenvalues and 1000x500x2 (2-component) eigenvectors. Just what I need. However, I need the eigenvalues and eigenvectors to be sorted. I tried to sort the eigenvalues with
vals, vecs = np.linalg.eig(array)
vals = vals[np.argsort(vals, axis = -1)]
yet this returns an array of shape 1000x500x2x500x2, instead of 1000x500x2, which I would expect. Is there a simple way to fix this, without resorting to looping over the array? And is there a similarly simple way to sort the eigenvectors according to the eigenvalues?
This question has sort of been answered here: Sorted eigenvalues and eigenvectors on a grid but the answer is tailored to a (another) special case, and not very well explained, and I haven't been successful in adapting it to my case. I considered posting a comment to that answer, asking for details, but I am currently at 49 reputation points, and you need 50 to comment.
Solution developed in conjunction with Tor. It will swap columns of the 2x2 matrices depending on "eigenvalue" order (fake eigenvectors generated by np.random).
import numpy as np
import itertools
Nx = 1000
Ny = 500
vals = np.random.rand(Nx, Ny, 2)
vecs = np.random.rand(Nx, Ny, 2, 2)
sort_order = np.argsort(vals, -1)
L = np.array(list(itertools.product(range(Nx), range(Ny), range(2))))
d0 = L[:,0]
d1 = L[:,1]
d2 = L[:,2]
vecsort = vecs[ np.c_[d0, d0], np.c_[d1, d1], np.c_[d2, d2], sort_order[d0, d1] ].reshape(Nx, Ny, 2, 2)
# vals can be sorted in place using:
vals.sort(axis=2)