图片解析应用
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

176 lines
6.9 KiB

  1. """Mixin classes for custom array types that don't inherit from ndarray."""
  2. from numpy.core import umath as um
  3. __all__ = ['NDArrayOperatorsMixin']
  4. def _disables_array_ufunc(obj):
  5. """True when __array_ufunc__ is set to None."""
  6. try:
  7. return obj.__array_ufunc__ is None
  8. except AttributeError:
  9. return False
  10. def _binary_method(ufunc, name):
  11. """Implement a forward binary method with a ufunc, e.g., __add__."""
  12. def func(self, other):
  13. if _disables_array_ufunc(other):
  14. return NotImplemented
  15. return ufunc(self, other)
  16. func.__name__ = '__{}__'.format(name)
  17. return func
  18. def _reflected_binary_method(ufunc, name):
  19. """Implement a reflected binary method with a ufunc, e.g., __radd__."""
  20. def func(self, other):
  21. if _disables_array_ufunc(other):
  22. return NotImplemented
  23. return ufunc(other, self)
  24. func.__name__ = '__r{}__'.format(name)
  25. return func
  26. def _inplace_binary_method(ufunc, name):
  27. """Implement an in-place binary method with a ufunc, e.g., __iadd__."""
  28. def func(self, other):
  29. return ufunc(self, other, out=(self,))
  30. func.__name__ = '__i{}__'.format(name)
  31. return func
  32. def _numeric_methods(ufunc, name):
  33. """Implement forward, reflected and inplace binary methods with a ufunc."""
  34. return (_binary_method(ufunc, name),
  35. _reflected_binary_method(ufunc, name),
  36. _inplace_binary_method(ufunc, name))
  37. def _unary_method(ufunc, name):
  38. """Implement a unary special method with a ufunc."""
  39. def func(self):
  40. return ufunc(self)
  41. func.__name__ = '__{}__'.format(name)
  42. return func
  43. class NDArrayOperatorsMixin:
  44. """Mixin defining all operator special methods using __array_ufunc__.
  45. This class implements the special methods for almost all of Python's
  46. builtin operators defined in the `operator` module, including comparisons
  47. (``==``, ``>``, etc.) and arithmetic (``+``, ``*``, ``-``, etc.), by
  48. deferring to the ``__array_ufunc__`` method, which subclasses must
  49. implement.
  50. It is useful for writing classes that do not inherit from `numpy.ndarray`,
  51. but that should support arithmetic and numpy universal functions like
  52. arrays as described in `A Mechanism for Overriding Ufuncs
  53. <https://numpy.org/neps/nep-0013-ufunc-overrides.html>`_.
  54. As an trivial example, consider this implementation of an ``ArrayLike``
  55. class that simply wraps a NumPy array and ensures that the result of any
  56. arithmetic operation is also an ``ArrayLike`` object::
  57. class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin):
  58. def __init__(self, value):
  59. self.value = np.asarray(value)
  60. # One might also consider adding the built-in list type to this
  61. # list, to support operations like np.add(array_like, list)
  62. _HANDLED_TYPES = (np.ndarray, numbers.Number)
  63. def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
  64. out = kwargs.get('out', ())
  65. for x in inputs + out:
  66. # Only support operations with instances of _HANDLED_TYPES.
  67. # Use ArrayLike instead of type(self) for isinstance to
  68. # allow subclasses that don't override __array_ufunc__ to
  69. # handle ArrayLike objects.
  70. if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)):
  71. return NotImplemented
  72. # Defer to the implementation of the ufunc on unwrapped values.
  73. inputs = tuple(x.value if isinstance(x, ArrayLike) else x
  74. for x in inputs)
  75. if out:
  76. kwargs['out'] = tuple(
  77. x.value if isinstance(x, ArrayLike) else x
  78. for x in out)
  79. result = getattr(ufunc, method)(*inputs, **kwargs)
  80. if type(result) is tuple:
  81. # multiple return values
  82. return tuple(type(self)(x) for x in result)
  83. elif method == 'at':
  84. # no return value
  85. return None
  86. else:
  87. # one return value
  88. return type(self)(result)
  89. def __repr__(self):
  90. return '%s(%r)' % (type(self).__name__, self.value)
  91. In interactions between ``ArrayLike`` objects and numbers or numpy arrays,
  92. the result is always another ``ArrayLike``:
  93. >>> x = ArrayLike([1, 2, 3])
  94. >>> x - 1
  95. ArrayLike(array([0, 1, 2]))
  96. >>> 1 - x
  97. ArrayLike(array([ 0, -1, -2]))
  98. >>> np.arange(3) - x
  99. ArrayLike(array([-1, -1, -1]))
  100. >>> x - np.arange(3)
  101. ArrayLike(array([1, 1, 1]))
  102. Note that unlike ``numpy.ndarray``, ``ArrayLike`` does not allow operations
  103. with arbitrary, unrecognized types. This ensures that interactions with
  104. ArrayLike preserve a well-defined casting hierarchy.
  105. .. versionadded:: 1.13
  106. """
  107. # Like np.ndarray, this mixin class implements "Option 1" from the ufunc
  108. # overrides NEP.
  109. # comparisons don't have reflected and in-place versions
  110. __lt__ = _binary_method(um.less, 'lt')
  111. __le__ = _binary_method(um.less_equal, 'le')
  112. __eq__ = _binary_method(um.equal, 'eq')
  113. __ne__ = _binary_method(um.not_equal, 'ne')
  114. __gt__ = _binary_method(um.greater, 'gt')
  115. __ge__ = _binary_method(um.greater_equal, 'ge')
  116. # numeric methods
  117. __add__, __radd__, __iadd__ = _numeric_methods(um.add, 'add')
  118. __sub__, __rsub__, __isub__ = _numeric_methods(um.subtract, 'sub')
  119. __mul__, __rmul__, __imul__ = _numeric_methods(um.multiply, 'mul')
  120. __matmul__, __rmatmul__, __imatmul__ = _numeric_methods(
  121. um.matmul, 'matmul')
  122. # Python 3 does not use __div__, __rdiv__, or __idiv__
  123. __truediv__, __rtruediv__, __itruediv__ = _numeric_methods(
  124. um.true_divide, 'truediv')
  125. __floordiv__, __rfloordiv__, __ifloordiv__ = _numeric_methods(
  126. um.floor_divide, 'floordiv')
  127. __mod__, __rmod__, __imod__ = _numeric_methods(um.remainder, 'mod')
  128. __divmod__ = _binary_method(um.divmod, 'divmod')
  129. __rdivmod__ = _reflected_binary_method(um.divmod, 'divmod')
  130. # __idivmod__ does not exist
  131. # TODO: handle the optional third argument for __pow__?
  132. __pow__, __rpow__, __ipow__ = _numeric_methods(um.power, 'pow')
  133. __lshift__, __rlshift__, __ilshift__ = _numeric_methods(
  134. um.left_shift, 'lshift')
  135. __rshift__, __rrshift__, __irshift__ = _numeric_methods(
  136. um.right_shift, 'rshift')
  137. __and__, __rand__, __iand__ = _numeric_methods(um.bitwise_and, 'and')
  138. __xor__, __rxor__, __ixor__ = _numeric_methods(um.bitwise_xor, 'xor')
  139. __or__, __ror__, __ior__ = _numeric_methods(um.bitwise_or, 'or')
  140. # unary methods
  141. __neg__ = _unary_method(um.negative, 'neg')
  142. __pos__ = _unary_method(um.positive, 'pos')
  143. __abs__ = _unary_method(um.absolute, 'abs')
  144. __invert__ = _unary_method(um.invert, 'invert')