如何在SymPy中定义支持运算与绘图的自定义分段函数

原代码存在的核心问题 #

• 未提前定义符号变量x，直接引用会触发未定义名称报错
• 数值计算分支写法冗余：不需要借助未定义的x做subs操作，直接对传入的参数计算返回值即可
• 冗余的isinstance(arg, exp)判断分支：完全不符合自定义Function类的运行逻辑，没有实际作用
• 缺少必要的类属性声明、微分方法实现，无法正常支持符号运算、绘图等SymPy内置能力

最简实现方案：使用内置Piecewise构造 #

这是SymPy官方推荐的分段函数实现方式，天然兼容所有内置计算、绘图能力，不需要额外开发自定义类，和内置函数行为完全一致：

from sympy import symbols, Piecewise, plot
# 定义实数域符号变量，避免逻辑歧义
x = symbols('x', real=True)
# 构造目标分段函数：x>=0时返回x，x<0时返回1/x²
f = Piecewise(
    (x, x >= 0),
    (1 / x**2, x < 0)
)
# 数值计算验证
print(f.subs(x, -2))  # 输出1/4
print(f.subs(x, 3))   # 输出3
# 直接调用绘图功能
plot(f, (x, -5, 5))

该写法支持符号求导、积分、表达式化简、变量代入等所有SymPy常规操作。 #

自定义Function类封装实现 #

如果需要将函数封装为类似sin、cos的独立可调用类（比如需要自定义打印规则、专属化简逻辑时），可按以下规则修正原代码：

1. eval方法仅处理纯数值输入的即时计算逻辑，符号输入时不强行返回展开表达式，保留函数原生形式
2. 实现fdiff方法支持符号求导
3. 移除无意义的类型判断分支
   修正后的完整可运行代码：

from sympy import Function, sympify, symbols, plot, Piecewise
x = symbols('x', real=True)

class CustomFunction(Function):
    @classmethod
    def eval(cls, arg):
        arg = sympify(arg)
        # 仅当传入是纯数值时，直接计算返回结果
        if arg.is_number:
            if arg >= 0:
                return arg
            else:
                return 1 / (arg ** 2)

    # 实现求导逻辑，支持符号微分
    def fdiff(self, argindex=1):
        arg = self.args[0]
        return Piecewise(
            (1, arg >= 0),
            (-2 / arg**3, arg < 0)
        )

# 使用方式与SymPy内置函数完全一致
f = CustomFunction(x)
# 数值计算验证
print(f.subs(x, -2))  # 输出1/4
print(f.subs(x, 3))   # 输出3
# 符号求导
print(f.diff(x))
# 直接调用绘图功能
plot(f, (x, -5, 5))

注意事项 #

• 定义符号变量时指定real=True，可以避免条件判断时出现符号域的逻辑歧义
• 自定义Function的eval方法如果对符号参数返回None，SymPy会自动保留函数表达式形式，不会提前展开
• 绘图功能无需额外适配，SymPy会自动传入数值采样点调用eval方法计算对应纵坐标值

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