Создание пользовательских контейнеров массивов#

Механизм диспетчеризации NumPy, введённый в версии numpy v1.16, является рекомендуемым подходом для написания пользовательских контейнеров N-мерных массивов, совместимых с API numpy и предоставляющих пользовательские реализации функциональности numpy. Приложения включают dask массивы, N-мерный массив, распределенный по нескольким узлам, и cupy массивы, N-мерный массив на GPU.

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

>>> import numpy as np
>>> class DiagonalArray:
...     def __init__(self, N, value):
...         self._N = N
...         self._i = value
...     def __repr__(self):
...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
...     def __array__(self, dtype=None, copy=None):
...         if copy is False:
...             raise ValueError(
...                 "`copy=False` isn't supported. A copy is always created."
...             )
...         return self._i * np.eye(self._N, dtype=dtype)

Наш пользовательский массив может быть создан следующим образом:

>>> arr = DiagonalArray(5, 1)
>>> arr
DiagonalArray(N=5, value=1)

Имя файла или файловый объект для чтения. numpy.array или numpy.asarray, который вызовет его __array__ метод для получения стандартного numpy.ndarray.

>>> np.asarray(arr)
array([[1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.]])

The __array__ метод может дополнительно принимать dtype аргумент. Если предоставлен, этот аргумент указывает желаемый тип данных для результирующего массива NumPy. Ваша реализация должна попытаться преобразовать данные к этому dtype если возможно. Если преобразование не поддерживается, обычно лучше вернуться к типу по умолчанию или вызвать TypeError или ValueError.

Вот пример, демонстрирующий его использование с dtype спецификация:

>>> np.asarray(arr, dtype=np.float32)
array([[1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.]], dtype=float32)

Если мы работаем с arr с функцией numpy, numpy снова использует __array__ интерфейс для преобразования его в массив и последующего применения функции обычным способом.

>>> np.multiply(arr, 2)
array([[2., 0., 0., 0., 0.],
       [0., 2., 0., 0., 0.],
       [0., 0., 2., 0., 0.],
       [0., 0., 0., 2., 0.],
       [0., 0., 0., 0., 2.]])

Обратите внимание, что возвращаемый тип является стандартным numpy.ndarray.

>>> type(np.multiply(arr, 2))

Как мы можем передать наш пользовательский тип массива через эту функцию? NumPy позволяет классу указать, что он хочет обрабатывать вычисления определенным пользователем способом через интерфейсы __array_ufunc__ и __array_function__. Давайте рассмотрим по одному, начиная с __array_ufunc__. Этот метод охватывает Универсальные функции (ufunc), класс функций, который включает, например, numpy.multiply и numpy.sin.

The __array_ufunc__ получает:

ufunc, функция, подобная numpy.multiply
method, строка, различающая numpy.multiply(...) и варианты, такие как numpy.multiply.outer, numpy.multiply.accumulateи так далее. Для общего случая, numpy.multiply(...), method == '__call__'.
inputs, который может быть смесью различных типов
kwargs, именованные аргументы, передаваемые в функцию

Для этого примера мы будем обрабатывать только метод __call__

>>> from numbers import Number
>>> class DiagonalArray:
...     def __init__(self, N, value):
...         self._N = N
...         self._i = value
...     def __repr__(self):
...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
...     def __array__(self, dtype=None, copy=None):
...         if copy is False:
...             raise ValueError(
...                 "`copy=False` isn't supported. A copy is always created."
...             )
...         return self._i * np.eye(self._N, dtype=dtype)
...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
...         if method == '__call__':
...             N = None
...             scalars = []
...             for input in inputs:
...                 if isinstance(input, Number):
...                     scalars.append(input)
...                 elif isinstance(input, self.__class__):
...                     scalars.append(input._i)
...                     if N is not None:
...                         if N != input._N:
...                             raise TypeError("inconsistent sizes")
...                     else:
...                         N = input._N
...                 else:
...                     return NotImplemented
...             return self.__class__(N, ufunc(*scalars, **kwargs))
...         else:
...             return NotImplemented

Теперь наш пользовательский тип массива проходит через функции numpy.

>>> arr = DiagonalArray(5, 1)
>>> np.multiply(arr, 3)
DiagonalArray(N=5, value=3)
>>> np.add(arr, 3)
DiagonalArray(N=5, value=4)
>>> np.sin(arr)
DiagonalArray(N=5, value=0.8414709848078965)

На этом этапе arr + 3 не работает.

>>> arr + 3
Traceback (most recent call last):
...
TypeError: unsupported operand type(s) for +: 'DiagonalArray' and 'int'

Для поддержки этого необходимо определить интерфейсы Python __add__, __lt__, и так далее для диспетчеризации к соответствующей универсальной функции. Мы можем достичь этого удобно, наследуя от примеси NDArrayOperatorsMixin.

>>> import numpy.lib.mixins
>>> class DiagonalArray(numpy.lib.mixins.NDArrayOperatorsMixin):
...     def __init__(self, N, value):
...         self._N = N
...         self._i = value
...     def __repr__(self):
...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
...     def __array__(self, dtype=None, copy=None):
...         if copy is False:
...             raise ValueError(
...                 "`copy=False` isn't supported. A copy is always created."
...             )
...         return self._i * np.eye(self._N, dtype=dtype)
...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
...         if method == '__call__':
...             N = None
...             scalars = []
...             for input in inputs:
...                 if isinstance(input, Number):
...                     scalars.append(input)
...                 elif isinstance(input, self.__class__):
...                     scalars.append(input._i)
...                     if N is not None:
...                         if N != input._N:
...                             raise TypeError("inconsistent sizes")
...                     else:
...                         N = input._N
...                 else:
...                     return NotImplemented
...             return self.__class__(N, ufunc(*scalars, **kwargs))
...         else:
...             return NotImplemented

>>> arr = DiagonalArray(5, 1)
>>> arr + 3
DiagonalArray(N=5, value=4)
>>> arr > 0
DiagonalArray(N=5, value=True)

Теперь давайте разберёмся с __array_function__. Мы создадим словарь, который сопоставляет функции numpy с нашими пользовательскими вариантами.

>>> HANDLED_FUNCTIONS = {}
>>> class DiagonalArray(numpy.lib.mixins.NDArrayOperatorsMixin):
...     def __init__(self, N, value):
...         self._N = N
...         self._i = value
...     def __repr__(self):
...         return f"{self.__class__.__name__}(N={self._N}, value={self._i})"
...     def __array__(self, dtype=None, copy=None):
...         if copy is False:
...             raise ValueError(
...                 "`copy=False` isn't supported. A copy is always created."
...             )
...         return self._i * np.eye(self._N, dtype=dtype)
...     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
...         if method == '__call__':
...             N = None
...             scalars = []
...             for input in inputs:
...                 # In this case we accept only scalar numbers or DiagonalArrays.
...                 if isinstance(input, Number):
...                     scalars.append(input)
...                 elif isinstance(input, self.__class__):
...                     scalars.append(input._i)
...                     if N is not None:
...                         if N != input._N:
...                             raise TypeError("inconsistent sizes")
...                     else:
...                         N = input._N
...                 else:
...                     return NotImplemented
...             return self.__class__(N, ufunc(*scalars, **kwargs))
...         else:
...             return NotImplemented
...     def __array_function__(self, func, types, args, kwargs):
...         if func not in HANDLED_FUNCTIONS:
...             return NotImplemented
...         # Note: this allows subclasses that don't override
...         # __array_function__ to handle DiagonalArray objects.
...         if not all(issubclass(t, self.__class__) for t in types):
...             return NotImplemented
...         return HANDLED_FUNCTIONS[func](*args, **kwargs)
...

Удобный шаблон — определить декоратор implements который можно использовать для добавления функций в HANDLED_FUNCTIONS.

>>> def implements(np_function):
...    "Register an __array_function__ implementation for DiagonalArray objects."
...    def decorator(func):
...        HANDLED_FUNCTIONS[np_function] = func
...        return func
...    return decorator
...

Теперь мы пишем реализации функций NumPy для DiagonalArray. Для полноты, чтобы поддержать использование arr.sum() добавить метод sum который вызывает numpy.sum(self), и то же самое для mean.

>>> @implements(np.sum)
... def sum(arr):
...     "Implementation of np.sum for DiagonalArray objects"
...     return arr._i * arr._N
...
>>> @implements(np.mean)
... def mean(arr):
...     "Implementation of np.mean for DiagonalArray objects"
...     return arr._i / arr._N
...
>>> arr = DiagonalArray(5, 1)
>>> np.sum(arr)
5
>>> np.mean(arr)
0.2

Если пользователь попытается использовать любые функции numpy, не включенные в HANDLED_FUNCTIONS, a TypeError будет вызвано исключение в numpy, указывающее, что эта операция не поддерживается. Например, конкатенация двух DiagonalArrays не создаёт другой диагональный массив, поэтому не поддерживается.

>>> np.concatenate([arr, arr])
Traceback (most recent call last):
...
TypeError: no implementation found for 'numpy.concatenate' on types that implement __array_function__: []

Кроме того, наши реализации sum и mean не принимают необязательные аргументы, которые принимает реализация numpy.

>>> np.sum(arr, axis=0)
Traceback (most recent call last):
...
TypeError: sum() got an unexpected keyword argument 'axis'

Пользователь всегда имеет возможность преобразовать в обычный numpy.ndarray с numpy.asarray и использование стандартного numpy оттуда.

>>> np.concatenate([np.asarray(arr), np.asarray(arr)])
array([[1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.],
       [1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.]])

Реализация DiagonalArray в этом примере обрабатывает только np.sum и np.mean функции для краткости. Многие другие функции в API Numpy также доступны для обертывания, и полноценный пользовательский контейнер массива может явно поддерживать все функции, которые Numpy делает доступными для обертывания.

NumPy предоставляет некоторые утилиты для помощи в тестировании пользовательских контейнеров массивов, которые реализуют __array_ufunc__ и __array_function__ протоколы в numpy.testing.overrides пространство имён.

Чтобы проверить, можно ли переопределить функцию Numpy через __array_ufunc__, вы можете использовать allows_array_ufunc_override:

>>> from numpy.testing.overrides import allows_array_ufunc_override
>>> allows_array_ufunc_override(np.add)
True

Аналогично, вы можете проверить, можно ли переопределить функцию через __array_function__ используя allows_array_function_override.

Списки каждой переопределяемой функции в API Numpy также доступны через get_overridable_numpy_array_functions для функций, поддерживающих __array_function__ протокол и get_overridable_numpy_ufuncs для функций, которые поддерживают __array_ufunc__ протокол. Обе функции возвращают наборы функций, присутствующих в публичном API Numpy. Пользовательские ufunc или ufunc, определённые в других библиотеках, зависящих от Numpy, не входят в эти наборы.

См. исходный код dask и исходный код cupy для более полных примеров пользовательских контейнеров массивов.

Смотрите также NEP 18.