编写自定义数组容器#
NumPy 在 v1.16 版本中引入的调度机制是编写与 NumPy API 兼容并提供 NumPy 功能自定义实现的自定义 N 维数组容器的推荐方法。应用包括dask 数组（分布在多个节点上的 N 维数组）和cupy 数组（GPU 上的 N 维数组）。
为了了解如何编写自定义数组容器，我们将从一个简单的示例开始，该示例的实用性比较窄，但说明了所涉及的概念。
>>> 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将其转换为 NumPy 数组，这将调用其__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.]])
如果我们使用 NumPy 函数对arr进行操作，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))
<class 'numpy.ndarray'>
我们如何通过此函数传递自定义数组类型？NumPy 允许一个类通过__array_ufunc__和__array_function__接口指示它希望以自定义方式处理计算。让我们一次一个地来看，从__array_ufunc__开始。此方法涵盖了通用函数 (ufunc)，这是一类函数，例如numpy.multiply和numpy.sin。
__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__等等，以调度到相应的 ufunc。我们可以通过继承自 mixin 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
...
现在我们编写DiagonalArray的 NumPy 函数实现。为完整起见，为了支持arr.sum()的使用，添加一个调用numpy.sum(self)的方法sum，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
如果用户尝试使用HANDLED_FUNCTIONS中未包含的任何 NumPy 函数，NumPy 将引发TypeError，指示不支持此操作。例如，连接两个DiagonalArrays不会产生另一个对角数组，因此不支持。
>>> np.concatenate([arr, arr])
Traceback (most recent call last):
...
TypeError: no implementation found for 'numpy.concatenate' on types that implement __array_function__: [<class '__main__.DiagonalArray'>]
此外，我们对sum和mean的实现不接受 NumPy 实现所接受的可选参数。
>>> np.sum(arr, axis=0)
Traceback (most recent call last):
...
TypeError: sum() got an unexpected keyword argument 'axis'
用户始终可以选择使用numpy.asarray转换为普通的numpy.ndarray，然后从那里使用标准的 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函数。NumPy API 中还有许多其他函数可以包装，一个完整的自定义数组容器可以明确地支持 NumPy 提供的所有可包装函数。
NumPy 提供了一些实用程序来帮助测试实现__array_ufunc__和__array_function__协议的自定义数组容器，这些实用程序位于numpy.testing.overrides命名空间中。
要检查是否可以通过__array_ufunc__覆盖 NumPy 函数，可以使用allows_array_ufunc_override
>>> from numpy.testing.overrides import allows_array_ufunc_override
>>> allows_array_ufunc_override(np.add)
True
类似地，您可以使用allows_array_function_override检查是否可以通过__array_function__覆盖函数。
NumPy API 中每个可覆盖函数的列表也可通过get_overridable_numpy_array_functions（对于支持__array_function__协议的函数）和get_overridable_numpy_ufuncs（对于支持__array_ufunc__协议的函数）获得。这两个函数都返回 NumPy 公共 API 中存在的函数集。在这些集合中不存在用户定义的 ufunc 或其他依赖于 NumPy 的库中定义的 ufunc。
有关自定义数组容器的更完整的示例，请参阅dask 源代码和cupy 源代码。
另请参阅NEP 18。