How can I vectorize a function that returns eigenvalues and eigenvectors of a matrix in python?

Question

I'm working with a function in Python that constructs a 4×4 matrix based on inputs (x1, y1, x2, y2), and computes its eigenvalues and eigenvectors using np.linalg.eigh.

Here is a simplified version of my code:

import numpy as np

def f(kx, ky):  
    return kx + 1j * ky

def fs(kx, ky):  
    return np.conj(f(kx, ky))

def eig(x1, y1, x2, y2):
    a = 10
    x = x1 + x2
    y = y1 + y2
    H = np.array([
        [a, f(x, y), f(x, y), fs(x, y)],
        [fs(x, y), a, 0, f(x, y)],
        [fs(x, y), 0, -a, f(x, y)],
        [f(x, y), fs(x, y), fs(x, y), -a]
    ])
    
    Eval, Evec = np.linalg.eigh(H)
    sorted_indices = np.argsort(Eval)
    return Eval[sorted_indices], Evec[:, sorted_indices]

Now, I have 1-d arrays of input values:

x1_array, y1_array, x2_array, y2_array  # all same shape

I want to efficiently vectorize this function across those arrays — i.e., compute all eigenvalues/eigenvectors in a batched way without writing an explicit Python loop if possible.

Answer 1

It is a bit of stacking, but you can create a matrix that is your batch size times 4x4 and pass it to np.linalg.eigh. Also note the slight optimization avoiding multiple evaluations of f(x, y) and fs(x, y).

def eig_vectorized(x1_array, y1_array, x2_array, y2_array):
    a       = 10
    x_array = x1_array + x2_array
    y_array = y1_array + y2_array
    f_xy    = f(x_array, y_array)   # Optimization, you don't want to recompute
    fs_xy   = fs(x_array, y_array)  # these two again and again
    
    # Create H as an array-sizex4x4 matrix
    H = np.stack([
        np.stack([np.full_like(f_xy, a), f_xy, f_xy, fs_xy],                 axis=-1),
        np.stack([fs_xy, np.full_like(f_xy, a), np.zeros_like(f_xy), f_xy],  axis=-1),
        np.stack([fs_xy, np.zeros_like(f_xy), -np.full_like(f_xy, a), f_xy], axis=-1),
        np.stack([f_xy, fs_xy, fs_xy, -np.full_like(f_xy, a)],               axis=-1)
    ], axis=-2)
    
    Evals, Evecs = np.linalg.eigh(H)  # Compute eigenvalues and -vectors

    # Sort eigenvalues and eigenvectors
    sorted_indices = np.argsort(Evals, axis=-1)
    Evals = np.take_along_axis(Evals, sorted_indices, axis=-1)
    Evecs = np.take_along_axis(Evecs, sorted_indices[..., None], axis=-1)

    return Evals, Evecs