如何基于Sympy创建自定义Symbol类？解决幂运算类型不符问题

问题描述 #

在Sympy符号计算场景下，通过继承Symbol类定义了CustomSymbol，同时自定义了CustomPow类，但执行幂运算后相乘时，结果被自动转换为Sympy原生Pow类型，而非自定义的CustomPow：

from sympy import *

class CustomSymbol(Symbol):
    def __pow__(self, other):
        return CustomPow(self, other)

class CustomPow(Pow):
    pass

a = CustomSymbol('a')
x = a**2 * a**3
type(x)  # 返回 <class 'sympy.core.power.Pow'>，而非预期的CustomPow

问题根源是Sympy内部会自动合并同底数幂运算，直接生成原生Pow(a, 2+3)，未触发自定义类的处理逻辑。以下是无需重写整个Sympy库的解决方案：

解决方案 #

方法1：重写_eval_*系列方法适配Sympy内部逻辑 #

Sympy内部运算优先调用_eval_*开头的方法，而非直接使用魔术方法。通过重写这些方法，让自定义类参与运算合并：

from sympy import *

class CustomSymbol(Symbol):
    def _eval_power(self, exp):
        # 替换__pow__，用Sympy标准的幂运算入口方法
        return CustomPow(self, exp)

class CustomPow(Pow):
    def _eval_mul(self, other):
        # 自定义同底数幂的乘法合并逻辑
        if isinstance(other, CustomPow) and self.base == other.base:
            return CustomPow(self.base, self.exp + other.exp)
        # 其他情况沿用父类逻辑
        return super()._eval_mul(other)

a = CustomSymbol('a')
x = a**2 * a**3
type(x)  # 返回 <class '__main__.CustomPow'>

方法2：直接继承Basic类完全控制符号行为 #

若需要彻底摆脱原生Symbol的默认行为，可直接继承Sympy的基础类Basic，手动实现运算逻辑：

from sympy import Basic, Integer

class CustomSymbol(Basic):
    def __new__(cls, name):
        return super().__new__(cls, name)
    
    def __pow__(self, other):
        return CustomPow(self, other)

class CustomPow(Basic):
    def __new__(cls, base, exp):
        # 自定义整数指数的幂实例创建逻辑
        if isinstance(base, CustomSymbol) and isinstance(exp, Integer):
            return super().__new__(cls, base, exp)
        return super().__new__(cls, base, exp)
    
    def __mul__(self, other):
        # 自定义同底数幂相乘的合并规则
        if isinstance(other, CustomPow) and self.args[0] == other.args[0]:
            return CustomPow(self.args[0], self.args[1] + other.args[1])
        return super().__mul__(other)

a = CustomSymbol('a')
x = a**2 * a**3
type(x)  # 返回 <class '__main__.CustomPow'>

方法3：注册自定义运算合并规则 #

通过Sympy的规则系统，为Mul运算添加自定义合并逻辑，指定处理CustomSymbol幂运算时优先使用CustomPow：

from sympy import *
from sympy.core.rules import Transform

class CustomSymbol(Symbol):
    pass

class CustomPow(Pow):
    pass

# 定义同底数CustomSymbol幂的合并规则
def custom_pow_merge(pow1, pow2):
    if (isinstance(pow1.base, CustomSymbol) and isinstance(pow2.base, CustomSymbol)
        and pow1.base == pow2.base):
        return CustomPow(pow1.base, pow1.exp + pow2.exp)
    return None

# 将规则注册到Mul的求值逻辑中
Mul._eval_rules.append(Transform(custom_pow_merge, predicate=lambda x: isinstance(x, Mul)))

a = CustomSymbol('a')
x = a**2 * a**3
type(x)  # 返回 <class '__main__.CustomPow'>

关键提示 #

• 优先使用_eval_power、_eval_mul这类Sympy标准入口方法，比直接重写魔术方法更兼容内部运算流程。
• 自定义类必须继承Sympy的Basic或其子类，才能融入Sympy的符号计算体系，无需重写整个库。
• 仅需自定义幂运算选方法1（轻量化），需完全控制符号行为选方法2，需扩展多个运算规则选方法3。

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