Своя арифметика в NumPy: dtype=object, магические методы и цепочки операций

Как в NumPy с dtype=object реализовать свою арифметику: магические методы Python‑класса, смешанные операции с int и цепочки выражений в массивах и со скалярами.

Создание собственного численного поведения в Python обычно начинается с магических методов вроде __add__ и __mul__. Сложность возникает, когда такие объекты оказываются внутри NumPy ndarray и участвуют в вычислениях вместе с массивами и скалярами. Будут ли ndarray по‑прежнему учитывать операторную логику вашего класса? Можно ли смешивать пользовательские объекты с обычными int? Короткий ответ, показанный ниже: при dtype=object NumPy делегирует операции поэлементно методам вашего класса, а с парой небольших доработок можно добиться согласованности для смешанных и цепочечных выражений.

Постановка задачи

Цель — использовать пользовательский класс, переопределяющий + и *, и добиться корректной работы операторов как в скалярных операциях, так и при хранении экземпляров в обычном ndarray. Требуемое поведение включает операции между массивами пользовательских объектов, между массивом и одиночным пользовательским объектом, между массивами пользовательских объектов и массивами обычных Python int, а также «перевёрнутые» операции вроде int * custom_object. И наконец, цепочки операций должны сохранять заданную семантику от начала до конца.

Минимальный пример, показывающий поведение

Далее — минимальный пример: класс преднамеренно «меняет местами» смысл + и *, чтобы результаты в демонстрациях были однозначными. Массив создаётся явно с dtype=object, чтобы элементы оставались Python‑объектами.

import numpy as np

class Quark:
    def __init__(self, payload):
        self.payload = payload

    def __add__(self, other):
        return self.payload * other.payload

    def __mul__(self, other):
        return self.payload + other.payload

x = Quark(5)
y = Quark(3)

box = np.array([x, y], dtype=object)

res = box * Quark(10)
print(res)  # [15 13], как и ожидалось

res = box * np.array([y, x], dtype=object)
print(res)  # [8 8], тоже как и ожидалось

Когда оба операнда — массивы пользовательских объектов, либо массив таких объектов комбинируется с одиночным экземпляром, поэлементный результат вычисляется через магические методы класса. Это ключевой момент: поэлементная диспетчеризация вызывает пользовательский оператор для каждого объекта, лежащего в массиве.

Что на самом деле происходит

При dtype=object арифметика с участием ndarray выполняется поэлементно, используя обычные правила разрешения операторов в Python. На практике это означает, что для каждого элемента вызывается его __add__ или __mul__, а в смешанных случаях подключаются «обратные» методы вроде __radd__ и __rmul__, если метод левого операнда не может обработать операцию. Это ровно то, что видно при запуске примеров и подтверждается результатами ниже.

Смешанные операции с Python ints

Если смешать массив пользовательских объектов с массивом int, прямой вызов может привести к ValueError: метод ожидает экземпляр вашего типа. Сделать класс устойчивым можно, научив его принимать и ваш тип, и обычные числа.

class Quark:
    def __init__(self, payload):
        self.payload = payload

    def __add__(self, other):
        if isinstance(other, Quark):
            return self.payload * other.payload
        return self.payload * other

    def __mul__(self, other):
        if isinstance(other, Quark):
            return self.payload + other.payload
        return self.payload + other

С этим изменением объединение массива объектов с массивом int работает нормально.

res = box * np.array([1, 2], dtype=object)  # [6 5]

Однако при смене порядка операндов вычисление всё ещё может падать: оператор для int выполняется первым и не умеет сочетаться с вашим типом. Здесь и нужны обратные магические методы.

res = np.array([1, 2], dtype=object) * box  # ошибка без обратных методов

Реализация __radd__ и __rmul__ делегирует вызов той же логике и устраняет асимметрию.

class Quark:
    def __init__(self, payload):
        self.payload = payload

    def __add__(self, other):
        if isinstance(other, Quark):
            return self.payload * other.payload
        return self.payload * other

    def __mul__(self, other):
        if isinstance(other, Quark):
            return self.payload + other.payload
        return self.payload + other

    def __radd__(self, other):
        return self.__add__(other)

    def __rmul__(self, other):
        return self.__mul__(other)

Теперь и array * array, и array * scalar работают независимо от порядка.

res = np.array([1, 2], dtype=object) * box  # [6 5]

Та же симметрия позволяет без проблем складывать смешанные массивы из объектов и чисел.

u = np.array([x, 2], dtype=object)
v = np.array([10, y], dtype=object)

out = u + v  # [50 6]

Сохранение пользовательской логики в цепочках выражений

Есть ещё одна тонкость в цепочках операций. Если ваши операторы возвращают обычные числа Python, следующий шаг цепочки пойдёт по стандартной числовой семантике, а не по вашей. Рассмотрим выражение, где левая операция возвращает int:

res = 2 * Quark(3) * 5  # 25, но при «поменянной местами» семантике ожидается 10

Первая часть, 2 * Quark(3), по «переставленным» правилам даёт 5, но затем 5 * 5 — это уже обычное умножение int и результатом будет 25. Чтобы цепочка продолжала вызывать ваши методы, операторы должны возвращать ваш собственный тип.

class Quark:
    def __init__(self, payload):
        self.payload = payload

    def __add__(self, other):
        if isinstance(other, Quark):
            return Quark(self.payload * other.payload)
        return Quark(self.payload * other)

    def __mul__(self, other):
        if isinstance(other, Quark):
            return Quark(self.payload + other.payload)
        return Quark(self.payload + other)

    def __radd__(self, other):
        return self.__add__(other)

    def __rmul__(self, other):
        return self.__mul__(other)

С таким изменением цепочка остаётся внутри вашего типа и сохраняет задуманную семантику операторов.

res = 2 * Quark(3) * 5
print(res.payload)  # 10

Для более приятного интерактивного вывода можно показать внутреннее значение через строковое представление.

def __repr__(self):
    return str(self.payload)

После этого при выводе объект показывает содержимое напрямую.

res = 2 * Quark(3) * 5  # 10

Итоги и рабочая схема

Главное: массивы объектов выполняют операции поэлементно и вызывают операторные методы самих хранимых объектов. Если реализовать __add__, __mul__, __radd__ и __rmul__, принимающие и ваш тип, и обычные числа, а также возвращать из них ваш тип, вы получите согласованное поведение для скалярных, массивных, смешанных, перевёрнутых и цепочечных операций. Это справедливо для обычных ndarrays, и даже обёртка или переобёртка через np.asarray сохраняет поведение при условии, что остаётся dtype=object.

Зачем это нужно

Понимание этой модели диспетчеризации позволяет проектировать «скаляроподобные» классы, которые интегрируются с NumPy без замены контейнера массива и без глобальных перехватов. Это полезно, когда создание массивов или паттерны их использования не под вашим контролем, а элементы встречаются в разных контекстах: как отдельные скаляры, в массивах или вперемешку с числами Python. С правильным набором магических методов одна и та же семантика работает стабильно и предсказуемо.

Что запомнить

Если вам нужна своя арифметика внутри NumPy ndarray, храните объекты в массивах с dtype=object — так каждая операция будет вызывать методы самих объектов. Реализуйте прямые и обратные операторы, чтобы покрыть смешанные случаи и обратный порядок операндов. Чтобы цепочки выражений сохраняли ту же семантику, возвращайте из арифметических методов ваш тип, а не голые примитивы Python. Для опрятного интерактивного вывода добавьте __repr__, который показывает внутреннее значение. С этими элементами обычных ndarrays достаточно, чтобы получить нужное поведение.