如何让Sympy支持含自定义函数的表达式的模运算求值？

你需要定义一个能处理模运算的自定义函数，且希望SymPy能自动拆解表达式（如Foo(n)+1）计算整体模值，而非返回未化简的Mod对象。以下是问题重现及解决方案：

问题重现 #

你的初始代码仅能处理单独的Foo(n) % mod，但组合表达式无法自动化简：

from sympy import Function

class Foo(Function):
    @classmethod
    def eval(cls, n):
        pass
    
    def __mod__(self, modulus: int) -> int:
        n, = self.args
        return (n + 1) % modulus

print(Foo(7) % 3)
# 2
print((Foo(7) + 1) % 3)
# Mod(Foo(7) + 1, 3)
print((Foo(7) + 1) % 3 == 0)
# False

解决方案 #

1. 适配SymPy的模运算钩子 #

SymPy为自定义函数提供了_eval_mod标准方法，替代直接重载__mod__，更适配SymPy的表达式系统：

from sympy import Function, Integer

class Foo(Function):
    @classmethod
    def eval(cls, n):
        pass
    
    def _eval_mod(self, modulus):
        # 定义Foo(n)的模运算逻辑
        n, = self.args
        return (n + 1) % modulus

2. 递归化简表达式模运算 #

实现一个递归函数，遍历表达式树，利用模运算的分配律（如(a+b)%m = [(a%m)+(b%m)]%m）自动拆解计算：

def simplify_mod(expr, modulus):
    if isinstance(expr, Integer):
        return expr % modulus
    elif isinstance(expr, Foo):
        return expr._eval_mod(modulus)
    elif expr.is_Add:
        total = 0
        for term in expr.args:
            total += simplify_mod(term, modulus)
        return total % modulus
    elif expr.is_Mul:
        product = 1
        for factor in expr.args:
            product *= simplify_mod(factor, modulus)
        return product % modulus
    elif expr.is_Mod:
        return simplify_mod(expr.args[0], expr.args[1])
    else:
        return expr % modulus

3. 测试验证 #

print(Foo(7) % 3)
# 输出: 2

expr = Foo(7) + 1
print(simplify_mod(expr, 3))
# 输出: 0

print(simplify_mod(expr, 3) == 0)
# 输出: True

原理说明 #

• _eval_mod是SymPy官方提供的自定义函数模运算接口，当SymPy处理Mod(Foo(n), m)时会自动调用该方法。
• simplify_mod递归遍历表达式的每个节点，对加法、乘法等运算应用模运算的分配规则，避免直接计算超大整数，完美适配你无法存储超大值的场景。

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