如何在Python中向量化批量计算矩阵特征值与特征向量的函数？

要实现无显式循环的批量计算，核心是利用numpy广播机制和np.linalg.eigh的批量处理能力（该函数支持最后两个维度为矩阵的高维数组输入），以下是具体改造方案：

向量化改造步骤 #

• 重构矩阵H为三维数组：将每个元素的计算适配数组输入，最终H的形状为(N, 4, 4)，其中N是输入数组的长度
• 批量计算特征值与特征向量：直接对三维H调用np.linalg.eigh，自动完成所有样本的特征计算
• 批量排序特征值与特征向量：利用np.argsort和np.take_along_axis完成批量排序，替代循环操作

改造后的完整代码 #

import numpy as np

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

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

def eig_vectorized(x1_array, y1_array, x2_array, y2_array):
    a = 10
    # 计算批量的x、y，形状为(N,)
    x = x1_array + x2_array
    y = y1_array + y2_array
    N = x.shape[0]
    
    # 初始化三维H矩阵，形状(N,4,4)
    H = np.zeros((N, 4, 4), dtype=np.complex128)
    # 批量填充矩阵元素，numpy自动广播适配批量维度
    H[:, 0, 0] = a
    H[:, 0, 1] = f(x, y)
    H[:, 0, 2] = f(x, y)
    H[:, 0, 3] = fs(x, y)
    
    H[:, 1, 0] = fs(x, y)
    H[:, 1, 1] = a
    H[:, 1, 3] = f(x, y)
    
    H[:, 2, 0] = fs(x, y)
    H[:, 2, 2] = -a
    H[:, 2, 3] = f(x, y)
    
    H[:, 3, 0] = f(x, y)
    H[:, 3, 1] = fs(x, y)
    H[:, 3, 2] = fs(x, y)
    H[:, 3, 3] = -a
    
    # 批量计算特征值和特征向量
    Eval, Evec = np.linalg.eigh(H)
    # 对每个样本的特征值单独排序，获取排序索引
    sorted_indices = np.argsort(Eval, axis=1)
    # 按索引重新排列特征值
    sorted_Eval = np.take_along_axis(Eval, sorted_indices, axis=1)
    # 按索引重新排列特征向量（适配三维数组的轴操作）
    sorted_Evec = np.take_along_axis(Evec, sorted_indices[:, None, :], axis=2)
    
    return sorted_Eval, sorted_Evec

代码说明 #

1. 三维矩阵构建：通过np.zeros初始化三维数组后，逐个位置赋值，numpy会自动将一维的f(x,y)等结果广播到批量维度，无需手动循环
2. 批量特征计算：np.linalg.eigh直接处理三维数组，返回的Eval形状为(N,4)（每个样本对应4个特征值），Evec形状为(N,4,4)（每个样本对应4×4的特征向量矩阵）
3. 批量排序：np.argsort指定axis=1实现对每个样本的特征值单独排序，np.take_along_axis高效完成批量索引重排，避免显式Python循环带来的性能损耗

测试示例 #

# 生成5个样本的测试输入
N = 5
x1_array = np.random.rand(N)
y1_array = np.random.rand(N)
x2_array = np.random.rand(N)
y2_array = np.random.rand(N)

# 调用向量化函数
evals, evecs = eig_vectorized(x1_array, y1_array, x2_array, y2_array)
print("特征值形状：", evals.shape)  # 输出 (5,4)
print("特征向量形状：", evecs.shape)  # 输出 (5,4,4)

内容的提问来源于stack exchange，提问作者Rinchen Sherpa