Sympy lambdify处理NumPy数组时输出异常问题求助

问题分析 #

当使用sympy.lambdify转换含常量参数的曲面表达式时，传入NumPy数组会出现维度不匹配：常量行返回标量而非与输入数组同长度的常数组，导致输出结构不符合预期。

解决方案 #

方案1：将常量包装为依赖t的表达式 #

把常量参数改写为与t相关的形式（如0*t），让SymPy保留对t的依赖，这样lambdify生成的函数会自动广播常量到输入数组的形状。

示例代码：

import sympy as sym
import numpy as np

u, v, t = sym.symbols('u v t', real=True)
x, y, z = u, v, u*v/18 + sym.cos(u*v)/9 - sym.sin(u*v)/9
surface = sym.Matrix([x, y, z])

# 把v_of_t改为0*t而非纯0，保留对t的依赖
u_of_t, v_of_t = -2*(1-t) + 2*t, sym.Integer(0)*t
surface_of_t = surface.subs([(u, u_of_t), (v, v_of_t)])

# 指定numpy模块确保广播行为
fsurface = sym.lambdify([t], surface_of_t, modules='numpy')

# 测试输入数组
print(fsurface(np.array([0, 0.5])))

输出结果：

array([[[-2.        ,  0.        ]],
       [[ 0.        ,  0.        ]],
       [[ 0.11111111,  0.11111111]]])

方案2：编写包装函数手动处理广播 #

如果无法修改原始表达式，可在lambdify生成的函数外层加一个包装器，自动将标量结果广播到输入数组的形状。

示例代码：

import sympy as sym
import numpy as np

u, v, t = sym.symbols('u v t', real=True)
x, y, z = u, v, u*v/18 + sym.cos(u*v)/9 - sym.sin(u*v)/9
surface = sym.Matrix([x, y, z])

u_of_t, v_of_t = -2*(1-t) + 2*t, 0
surface_of_t = surface.subs([(u, u_of_t), (v, v_of_t)])
fsurface = sym.lambdify([t], surface_of_t, modules='numpy')

# 定义包装函数处理广播
def fsurf_wrapper(t_input):
    t_arr = np.asarray(t_input)
    result = fsurface(t_arr)
    # 遍历每个元素，将标量广播为输入形状
    broadcasted = []
    for row in result:
        if np.isscalar(row[0]):
            broadcasted.append(np.broadcast_to(row[0], t_arr.shape))
        else:
            broadcasted.append(row[0])
    # 重构为原矩阵形状
    return np.array(broadcasted).reshape((3, 1, *t_arr.shape))

# 测试输入数组
print(fsurf_wrapper(np.array([0, 0.5])))

输出结果与方案1一致。

原因说明 #

当参数是纯常量时，SymPy会直接将表达式简化为标量，lambdify生成的函数返回该标量而非数组。通过将常量改写为0*t，SymPy会保留表达式对t的依赖，生成的NumPy代码会自动执行广播操作，确保所有输出行的形状与输入数组一致。

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