建议使用numpy版本v1.16中引入的Numpy的分派机制来编写与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):
... return self._i * np.eye(self._N)
...
我们的自定义数组可以像这样实例化:
>>> arr = DiagonalArray(5, 1)
>>> arr
DiagonalArray(N=5, value=1)
我们可以使用numpy.array
或
转换为numpy数组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.]])
如果我们使用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.]])
请注意,返回类型是standard numpy.ndarray
。
>>> type(arr)
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):
... return self._i * np.eye(self._N)
... 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 != self._N:
... raise TypeError("inconsistent sizes")
... else:
... N = self._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
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):
... return self._i * np.eye(self._N)
... 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 != self._N:
... raise TypeError("inconsistent sizes")
... else:
... N = self._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):
... return self._i * np.eye(self._N)
... 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 != self._N:
... raise TypeError("inconsistent sizes")
... else:
... N = self._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
,TypeError
则numpy将引发a ,表示不支持此操作。例如,连接两个
DiagonalArrays
不会产生另一个对角线数组,因此不支持它。
>>> np.concatenate([arr, arr])
TypeError: no implementation found for 'numpy.concatenate' on types that implement __array_function__: [<class '__main__.DiagonalArray'>]
此外,我们的实现,sum
并且mean
不接受numpy实现的可选参数。
>>> np.sum(arr, axis=0)
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.]])
有关自定义数组容器的更完整的示例,请参考dask源代码和 cupy源代码。
另请参阅NEP 18。