Vectorized Eigen Decomposition in NumPy: Batched 4x4 Matrices with np.linalg.eigh, No Python Loops

Vectorize NumPy eigen decomposition for batched 4x4 matrices with np.linalg.eigh: compute all eigenvalues/eigenvectors at once and keep a consistent ordering.

Vectorizing linear algebra workloads in NumPy pays off when the same operation must be repeated for many inputs. A common case is computing eigenvalues and eigenvectors for a structured 4×4 matrix that depends on several parameters. The goal here is to run the whole batch at once without an explicit Python loop.

Baseline: single-call eigen decomposition

The following function builds a 4×4 matrix from inputs (x1, y1, x2, y2), then computes and sorts its eigenvalues and eigenvectors using np.linalg.eigh.

import numpy as np

def phi(px, py):
    return px + 1j * py

def phi_bar(px, py):
    return np.conj(phi(px, py))

def solve_block(u1, v1, u2, v2):
    alpha = 10
    u = u1 + u2
    v = v1 + v2
    M = np.array([
        [alpha,           phi(u, v),     phi(u, v),     phi_bar(u, v)],
        [phi_bar(u, v),   alpha,         0,             phi(u, v)],
        [phi_bar(u, v),   0,             -alpha,        phi(u, v)],
        [phi(u, v),       phi_bar(u, v), phi_bar(u, v), -alpha]
    ])

    w, V = np.linalg.eigh(M)
    order = np.argsort(w)
    return w[order], V[:, order]

What blocks vectorization

This implementation operates on scalars and returns results for a single matrix. When inputs arrive as 1‑d arrays of the same shape, calling it in a loop is undesirable. The task is to compute all eigenvalues/eigenvectors in one go, in a batched manner, while keeping the same matrix construction and sorting behavior.

Vectorized approach

The idea is to assemble a stacked array of shape “batch × 4 × 4” and pass that stack directly to np.linalg.eigh. It also helps to avoid repeated evaluation of identical subexpressions by computing them once and reusing them during the assembly of the batch.

def solve_batched(u1_arr, v1_arr, u2_arr, v2_arr):
    alpha   = 10
    u_arr   = u1_arr + u2_arr
    v_arr   = v1_arr + v2_arr
    z       = phi(u_arr, v_arr)
    z_conj  = phi_bar(u_arr, v_arr)

    H = np.stack([
        np.stack([np.full_like(z, alpha), z,       z,       z_conj],              axis=-1),
        np.stack([z_conj,                 np.full_like(z, alpha), np.zeros_like(z), z], axis=-1),
        np.stack([z_conj,                 np.zeros_like(z),       -np.full_like(z, alpha), z], axis=-1),
        np.stack([z,                      z_conj,                 z_conj,                 -np.full_like(z, alpha)], axis=-1)
    ], axis=-2)

    evals, evecs = np.linalg.eigh(H)

    idx = np.argsort(evals, axis=-1)
    evals = np.take_along_axis(evals, idx, axis=-1)
    evecs = np.take_along_axis(evecs, idx[..., None], axis=-1)

    return evals, evecs

Why this works

By stacking the matrices into a single array, the eigen decomposition is computed for the whole batch in one call. Precomputing phi(u, v) and its conjugate prevents redundant work when populating the blocks. The final step mirrors the baseline: eigenvalues and eigenvectors are sorted consistently along the last axis.

Why you should care

This pattern lets you compute all eigenvalues and eigenvectors across many input tuples without writing an explicit Python loop. The construction of the matrix remains the same; only the data layout changes to a batched form that is processed in one shot.

Takeaways

When a matrix depends on elementwise expressions of input arrays, first build the shared intermediates, then assemble a stacked 4×4 block per input, and finally call np.linalg.eigh once on the batch. Keep the postprocessing identical to the scalar version by sorting eigenpairs along the last axis so downstream code sees consistent ordering.