vectorising linalg.eig() in numpy

Question

I have an m*m*n numpy array (call it A) and I would like to find the eigenvalues of every submatrix A[:,:,n] in this array. I could do it with linalg.eig() in a loop with relative ease, but there really ought to be a way to vectorise this. Something like a ufunc, but that can process subvectors instead of individual elements. Is this possible?

Answer 1

The computation of the eigenvalues and eigenvectors can not be vectorised in the sense that there's no way in general to share work for different matrices. np.linalg.eig (for real input) is just a wrapper for dgeev, which according to the docs only accepts a single matrix per call and the computation is fairly expensive, so for matrices that are not small the overhead of a python loop will be negligible.

Though, if you're doing this for many very small matrices it can become too slow. There are several questions related to this and the solution usually ends up being a compiled extension. As enigmaticPhysicist says in a comment, the idea of processing subvectors and submatrices in the same way as ufuncs would be useful in general. These are called generalised ufuncs and are already in numpy's development version. I find it around 8 times faster for matrices of shape (1000, 3, 3):

In [2]: np.__version__
Out[2]: '1.8.0.dev-dcf7cac'

In [3]: A = np.random.rand(1000, 3, 3)

In [4]: timeit np.linalg.eig(A)
P100 loops, best of 3: 9.65 ms per loop

In [5]: timeit [np.linalg.eig(Ai) for Ai in A]
10 loops, best of 3: 74.6 ms per loop

In [6]: a1 = np.linalg.eig(A)

In [7]: a2 = [np.linalg.eig(Ai) for Ai in A]

In [8]: all(np.allclose(a1[i][j], a2[j][i]) for j in xrange(1000) for i in xrange(2))
Out[8]: True