Search code examples
pythonnumpymatrixlinear-algebralapack

Eigen values for matrix at each grid point


I have a 3d grid. At each grid point, I have a matrix. I would like to find eigen values and eigen vectors of this matrix at each grid point using python. I am doing something like this.

import numpy as np
from numpy import linalg as LA
n = 256
s = np.zeros((3,3,n,n,n))
#s is calculated by a formula here, this part is correct
e = np.zeros((3,n,n,n))
e[0:3,:,:,:] = LA.eigvals(s[0:3,0:3,:,:,:])

It gives following error,

1 n = 256
2 e = np.zeros(((3),n,n,n))
----> 3 e[0:3,:,:,:] = LA.eigvals(s[0:3,0:3,:,:,:])

ValueError: could not broadcast input array from shape (3,3,256,256) into shape (3,256,256,256)

Since it is a large array, using loops is taking a lot of time. I actually have to do it for many such cubic grids. Is there a way to avoid loops ? The code has to understand that the at each grid point, there is a matrix whose eigen values are needed, it is not a 5 cross 5 matrix.


Solution

  • To expand on my comment and offer some actual code, you could get this to work with your existing matrices with some array reshaping:

    eigs = np.linalg.eigvals(s.swapaxes(0, -1).swapaxes(1,-2))
    e = eigs.swapaxes(0,-1)
    

    will give you a matrix with

    >>> e.shape()
    (3, 256, 256, 256)