免費AI編程助手
NumPy 編寫自定義數(shù)組容器
在 numpy v1.16 版本中引入的 Numpy 調(diào)度機制是編寫與 numpy API 兼容并提供 numpy 功能的自定義實現(xiàn)的自定義 N 維數(shù)組容器的推薦方法。應用程序包括dask數(shù)組（分布在多個節(jié)點上的 N 維數(shù)組）和cupy數(shù)組（GPU 上的 N 維數(shù)組）。
為了感受如何編寫自定義數(shù)組容器，我們將從一個簡單的示例開始，該示例具有相當狹窄的實用性，但說明了所涉及的概念。
>>> 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):
... return self._i * np.eye(self._N, dtype=dtype)我們的自定義數(shù)組可以像這樣實例化：
>>> arr = DiagonalArray(5, 1)
>>> arr
DiagonalArray(N=5, value=1)我們可以使用numpy.arrayor?轉換成一個 numpy 數(shù)組numpy.asarray，它會調(diào)用它的__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.]])如果我們使用arrnumpy 函數(shù)進行操作，numpy 將再次使用該?__array__接口將其轉換為數(shù)組，然后以通常的方式應用該函數(shù)。
>>> 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.ndarray我們?nèi)绾瓮ㄟ^這個函數(shù)傳遞我們的自定義數(shù)組類型？Numpy 允許一個類通過接口__array_ufunc__和__array_function__.?讓我們一次一個，從_array_ufunc__.?此方法涵蓋?通用函數(shù) (ufunc)，這是一類函數(shù)，例如包括?numpy.multiply和numpy.sin。
該__array_ufunc__接收：
ufunc，函數(shù)如?numpy.multiplymethod, 一個字符串，區(qū)分numpy.multiply(...)和 變體，如numpy.multiply.outer、numpy.multiply.accumulate等。對于常見情況，numpy.multiply(...),?。method?==?'__call__'inputs，這可能是不同類型的混合kwargs, 傳遞給函數(shù)的關鍵字參數(shù)
對于這個例子，我們將只處理方法?__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):
... 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 != self._N:
... raise TypeError("inconsistent sizes")
... else:
... N = self._N
... else:
... return NotImplemented
... return self.__class__(N, ufunc(*scalars, **kwargs))
... else:
... return NotImplemented現(xiàn)在我們的自定義數(shù)組類型通過 numpy 函數(shù)。
>>> 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 繼承來方便地實現(xiàn)這一點?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):
... 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 != 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)現(xiàn)在讓我們解決__array_function__.?我們將創(chuàng)建 dict 將 numpy 函數(shù)映射到我們的自定義變體。
>>> 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):
... 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 != 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
...現(xiàn)在我們編寫 numpy 函數(shù)的實現(xiàn)DiagonalArray。為了完整起見，為了支持用法，arr.sum()添加一個sum調(diào)用的方法，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 函數(shù)?HANDLED_FUNCTIONS，TypeError則 numpy 將引發(fā)a?，表示不支持此操作。例如，連接兩個?DiagonalArrays不會產(chǎn)生另一個對角數(shù)組，因此不支持。
>>> np.concatenate([arr, arr])
TypeError: no implementation found for 'numpy.concatenate' on types that implement __array_function__: [<class '__main__.DiagonalArray'>]此外，我們的sum和mean實現(xiàn)不接受 numpy 的實現(xiàn)所做的可選參數(shù)。
>>> 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.]])有關自定義數(shù)組容器的更完整示例，請參閱dask 源代碼和?cupy 源代碼。 另見NEP 18。
更多建議：