Пользовательские функции SymPy: реализуем _eval_Mod для Mod

Разбираем, почему перегрузка % не работает в выражениях SymPy, и показываем интеграцию через _eval_Mod для Mod: вычисление остатков по модулю в формулах.

Пользовательские хуки вычисления в символьных системах отделяют выражения, которые лишь «красиво печатаются», от тех, что действительно считаются. Если вы определяете собственную функцию SymPy и хотите, чтобы она участвовала в модульной арифметике внутри больших выражений, наивной перегрузки оператора недостаточно. Правильный путь — подключиться к протоколу вычисления SymPy для Mod.

Problem setup

Цель — определить символьную функцию, которую нельзя вычислить напрямую, но которая всё же позволяет находить определённые свойства, например остатки по модулю. Прямая перегрузка оператора вроде бы справляется с самым простым случаем, но как только функция попадает в более крупное выражение, всё перестаёт работать.

from sympy import Function
class Buzz(Function):
    @classmethod
    def eval(cls, k):
        pass
    def __mod__(self, m: int) -> int:
        val, = self.args
        return (val + 1) % m
print(Buzz(7) % 3)
# 2
print((Buzz(7) + 1) % 3)
# Mod(Buzz(7) + 1, 3)
print((Buzz(7) + 1) % 3 == 0)
# False

Применение % к отдельному экземпляру работает, но как только обернуть его в сумму, вычисление блокируется. Казалось бы, остаток от суммы можно получить из остатков её слагаемых, однако выражение остаётся символьным.

Why this happens

В SymPy оператор % для выражений не вызывает ваш __mod__ так, как вы могли бы ожидать. Вместо этого Expr.__mod__ строит символьный Mod(self, other). Далее SymPy обращается к самому объекту Mod за правилом вычисления, а Mod, в свою очередь, ищет метод _eval_Mod у составляющих выражения. Если ваша пользовательская функция не реализует _eval_Mod, то в составных контекстах выражение обычно остаётся невычисленным.

The right way: implement _eval_Mod

Чтобы пользовательская символьная функция корректно работала с модульной арифметикой внутри больших выражений, реализуйте метод _eval_Mod. Перегружать __mod__ для этого случая не нужно: протокол на уровне выражений уже направляет вычисление остатка через Mod, который и вызывает _eval_Mod.

from sympy import Function, Expr, Integer
class Zed(Function):
    @classmethod
    def eval(cls, k):
        pass
    def _eval_Mod(self, modv: Expr) -> Integer:
        if not isinstance(modv, Integer):
            raise TypeError
        arg0, = self.args
        return (arg0 + 1) % modv
print(Zed(7) % 3)
# 2
print((Zed(7) + 1) % 3)
# 0

Так функция начинает участвовать в модульных вычислениях на уровне выражений. Отдельный вызов по‑прежнему работает, и что важно — суммы, содержащие функцию, теперь сокращаются по модулю так, как и ожидалось.

Scope and observed behavior

Улучшение выше распространяется на случаи вроде (Zed(7) + 1) % 3. При этом выражение вида Zed(7) * 2 % 3 может остаться как Mod(2*Zed(7), 3). Показанная интеграция делает вычислимыми аддитивные контексты, тогда как другие формы могут по‑прежнему оставаться символьными.

Related tools inside SymPy

Полиномиальные процедуры SymPy напрямую поддерживают арифметику по модулю. Ряд функций принимает параметр modulus — это полезно при разложении или упрощении многочленов над конечными кольцами.

from sympy import symbols, factor
x = symbols('x')
print(factor(x**2 + 1))
# x**2 + 1
print(factor(x**2 + 1, modulus=2))
# (x + 1)**2

Why this matters

В некоторых сценариях целые числа настолько велики, что их нельзя материализовать; числа существуют лишь как символьные конструкции. Даже в таком режиме может понадобиться вычислять отдельные свойства, например остатки по модулю. Интеграция с вычислительными хуками SymPy позволяет вашей функции работать внутри составных выражений, не вынуждая выполнять ранние полные вычисления объекта.

Takeaways

Чтобы пользовательские функции SymPy корректно взаимодействовали с модульной арифметикой в составе выражений, опирайтесь на протокол выражений, а не на перегрузку операторов. Реализуйте у функции _eval_Mod, чтобы Mod мог вычислять её вклад в больших формулах. Для работы с многочленами помните, что многие процедуры принимают аргумент modulus. Такой подход сохраняет объекты символьными, но позволяет выполнять точечные вычисления — именно то, что нужно, когда базовые числа слишком велики, чтобы хранить их целиком.