如何自定义SymPy打印函数以简化链式法则运算输出？

问题场景 #

使用SymPy计算链式法则时，输出会携带所有变量依赖（如(x,y,z,t)），导致结果冗长。示例代码：

import sympy as sym
def ddt(s):
    return sym.diff(s, t)
x, y, z, t = sym.symbols("x y z t")
a = sym.Function('a')(x,y,z,t)
b = sym.Function('b')(x,y,z,t)
c = sym.Function('c')(a, b)
print(ddt(c))

原始输出：

Derivative(a(x, y, z, t), t)*Derivative(c(a(x, y, z, t), b(x, y, z, t)), a(x, y, z, t)) + Derivative(b(x, y, z, t), t)*Derivative(c(a(x, y, z, t), b(x, y, z, t)), b(x, y, z, t))

期望得到简洁输出：

Derivative(a, t)*Derivative(c, a) + Derivative(b, t)*Derivative(c, b)

同时需要对应简洁的LaTeX格式：$\frac{\partial a}{\partial t}\frac{\partial c}{\partial a} + \frac{\partial b}{\partial t}\frac{\partial c}{\partial b}$

解决方案 #

可以通过自定义替换规则移除函数的变量依赖标注，以下两种方式可实现需求：

方法1：用replace批量替换输出文本

利用SymPy的表达式替换功能，将所有带参数的函数实例替换为对应的抽象符号：

import sympy as sym

def ddt(s):
    return sym.diff(s, t)

x, y, z, t = sym.symbols("x y z t")
# 定义带参数的函数实例
a_func = sym.Function('a')(x,y,z,t)
b_func = sym.Function('b')(x,y,z,t)
c_func = sym.Function('c')(a_func, b_func)

# 生成导数表达式
result = ddt(c_func)

# 定义对应抽象符号
a_sym = sym.symbols('a')
b_sym = sym.symbols('b')
c_sym = sym.symbols('c')

# 替换所有函数实例为抽象符号
simplified_result = result.replace(a_func, a_sym).replace(b_func, b_sym).replace(c_func, c_sym)

# 打印简洁形式
print(simplified_result)
# 输出：Derivative(a, t)*Derivative(c, a) + Derivative(b, t)*Derivative(c, b)

# 生成简洁LaTeX代码
latex_output = sym.latex(simplified_result)
print(latex_output)
# 输出：\frac{\partial a}{\partial t} \frac{\partial c}{\partial a} + \frac{\partial b}{\partial t} \frac{\partial c}{\partial b}

方法2：先定义抽象符号，再关联具体函数

先基于抽象符号构建导数关系，再替换为具体函数验证，最后反向替换回抽象符号输出：

import sympy as sym

# 先定义抽象符号
a, b, c, t = sym.symbols("a b c t")
# 构建抽象导数表达式
abstract_deriv = sym.diff(c, a)*sym.diff(a, t) + sym.diff(c, b)*sym.diff(b, t)

# 定义带参数的具体函数
x, y, z = sym.symbols("x y z")
a_func = sym.Function('a')(x,y,z,t)
b_func = sym.Function('b')(x,y,z,t)
c_func = sym.Function('c')(a_func, b_func)

# 计算具体函数的导数（可选验证步骤）
concrete_deriv = sym.diff(c_func, t)

# 将具体导数替换回抽象符号输出
simplified_result = concrete_deriv.subs({a_func: a, b_func: b, c_func: c})
print(simplified_result)
# 输出与方法1一致

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