如何用Python优雅处理元素为一维数组的向量与矩阵的线性代数运算？

方案1：用NumPy的einsum与索引简化操作 #

针对你的场景，无需复杂广播，利用NumPy的张量收缩工具einsum和对角索引就能轻松实现目标操作：

对角矩阵的指数运算

import numpy as np

# 构造示例3×3×N对角矩阵
N = 2
A = np.zeros((3, 3, N))
A[0, 0, :] = [1, 2]
A[1, 1, :] = [2, 4]
A[2, 2, :] = [3, 6]

# 计算矩阵指数：仅对对角元素取指数
exp_A = np.zeros_like(A)
# 获取3×3矩阵的对角索引
diag_indices = np.diag_indices(3)
exp_A[diag_indices] = np.exp(A[diag_indices])

# 输出结果
print(exp_A)

矩阵与向量的乘法（3×3×N 乘 3×N）

用einsum直接描述维度间的运算关系，避免手动调整轴：

# 构造3×N的波函数向量
psi = np.random.rand(3, N)
# 执行矩阵-向量乘法：result[i,k] = Σ_j A[i,j,k] * psi[j,k]
result = np.einsum('ijk,jk->ik', A, psi)

方案2：用Einops库实现更直观的维度操作 #

Einops通过字符串语法清晰描述维度变换/运算逻辑，完全避开广播规则的困扰，非常适合量子系统的张量操作：

首先安装Einops：

pip install einops

对角矩阵指数运算

from einops import einsum

exp_A = np.zeros_like(A)
# 仅对i=j的对角位置赋值指数结果
exp_A = einsum(exp_A, A, 'i j k -> i j k', where='i==j') + np.exp(einsum(A, 'i i k -> i k'))[:, None, :]

矩阵-向量乘法

result = einsum(A, psi, 'i j k, j k -> i k')

这两种方案都能让你的张量操作逻辑清晰、代码简洁，无需再纠结复杂的广播规则。

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