使用NumPy einsum实现球坐标梯度表达式报错问题排查

需求与代码背景 #

要实现的数学表达式：
$$\left(\partial_q n_b \partial_q n_b \partial_s n_c \partial_s n_c-\partial_q n_b \partial_s n_b \partial_s n_c \partial_q n_c\right)$$

其中$n_b$的定义：

nb = np.array([cos(theta)*sin(B*phi), cos(theta)*cos(B*phi), sin(theta)])

$\partial_s$为球坐标下的梯度，计划用einsum实现，写出的代码：

expr = np.einsum('i,i,j,j->', grad_theta, grad_theta, grad_phi, grad_phi) 
       - np.einsum('i,j,i,j->', grad_theta, grad_phi, grad_phi, grad_theta)

运行时报错：

ValueError: operand has more dimensions than subscripts given in einstein sum, but no '...' ellipsis provided to broadcast the extra dimensions.

补充的$n$定义与梯度形式：

nb = Matrix([sin(theta)*sin(B*phi), sin(theta)*cos(B*phi), cos(theta)])

grad_theta = diff(nb, theta) * (1/r)
grad_phi = diff(nb, phi) * (1/(r*sin(theta)))

grad_theta的具体结构：

grad_theta = Matrix([
[sin(B*phi)*cos(theta)/r],
[cos(theta)*cos(B*phi)/r],
[          -sin(theta)/r]])

错误原因 #

1. 维度与下标不匹配
   传入einsum的grad_theta和grad_phi是**(3,1)的2维列向量**，但代码中写的下标（i/j）仅对应1维数组，einsum无法解析多出来的维度，因此抛出错误。

2. 数据类型混用
   代码同时使用了NumPy数组和SymPy的Matrix/diff，最终得到的grad_theta是SymPy矩阵对象而非NumPy数组。np.einsum仅支持NumPy数组，直接传入SymPy矩阵会导致维度解析异常。

修复方案 #

方案1：统一为NumPy数组处理 #

将SymPy矩阵转为NumPy数组（若需保留符号计算，需确保变量为SymPy符号），同时修正einsum下标：

# 先将SymPy矩阵转为NumPy数组，提取列向量为1维数组
gt = np.array(grad_theta)[:, 0]
gp = np.array(grad_phi)[:, 0]

expr = np.einsum('i,i,j,j->', gt, gt, gp, gp) \
       - np.einsum('i,j,i,j->', gt, gp, gp, gt)

方案2：用...忽略额外维度 #

如果保留2维数组结构，用...让einsum自动广播多余维度：

expr = np.einsum('...i,...i,...j,...j->', grad_theta, grad_theta, grad_phi, grad_phi) \
       - np.einsum('...i,...j,...i,...j->', grad_theta, grad_phi, grad_phi, grad_theta)

方案3：改用SymPy符号计算的einsum #

如果全程做符号推导，直接使用SymPy的einsum：

from sympy import einsum

expr = einsum('i,i,j,j->', grad_theta[:,0], grad_theta[:,0], grad_phi[:,0], grad_phi[:,0]) \
       - einsum('i,j,i,j->', grad_theta[:,0], grad_phi[:,0], grad_phi[:,0], grad_theta[:,0])

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