stride_tricks.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268
  1. """
  2. Utilities that manipulate strides to achieve desirable effects.
  3. An explanation of strides can be found in the "ndarray.rst" file in the
  4. NumPy reference guide.
  5. """
  6. from __future__ import division, absolute_import, print_function
  7. import numpy as np
  8. from numpy.core.overrides import array_function_dispatch
  9. __all__ = ['broadcast_to', 'broadcast_arrays']
  10. class DummyArray(object):
  11. """Dummy object that just exists to hang __array_interface__ dictionaries
  12. and possibly keep alive a reference to a base array.
  13. """
  14. def __init__(self, interface, base=None):
  15. self.__array_interface__ = interface
  16. self.base = base
  17. def _maybe_view_as_subclass(original_array, new_array):
  18. if type(original_array) is not type(new_array):
  19. # if input was an ndarray subclass and subclasses were OK,
  20. # then view the result as that subclass.
  21. new_array = new_array.view(type=type(original_array))
  22. # Since we have done something akin to a view from original_array, we
  23. # should let the subclass finalize (if it has it implemented, i.e., is
  24. # not None).
  25. if new_array.__array_finalize__:
  26. new_array.__array_finalize__(original_array)
  27. return new_array
  28. def as_strided(x, shape=None, strides=None, subok=False, writeable=True):
  29. """
  30. Create a view into the array with the given shape and strides.
  31. .. warning:: This function has to be used with extreme care, see notes.
  32. Parameters
  33. ----------
  34. x : ndarray
  35. Array to create a new.
  36. shape : sequence of int, optional
  37. The shape of the new array. Defaults to ``x.shape``.
  38. strides : sequence of int, optional
  39. The strides of the new array. Defaults to ``x.strides``.
  40. subok : bool, optional
  41. .. versionadded:: 1.10
  42. If True, subclasses are preserved.
  43. writeable : bool, optional
  44. .. versionadded:: 1.12
  45. If set to False, the returned array will always be readonly.
  46. Otherwise it will be writable if the original array was. It
  47. is advisable to set this to False if possible (see Notes).
  48. Returns
  49. -------
  50. view : ndarray
  51. See also
  52. --------
  53. broadcast_to: broadcast an array to a given shape.
  54. reshape : reshape an array.
  55. Notes
  56. -----
  57. ``as_strided`` creates a view into the array given the exact strides
  58. and shape. This means it manipulates the internal data structure of
  59. ndarray and, if done incorrectly, the array elements can point to
  60. invalid memory and can corrupt results or crash your program.
  61. It is advisable to always use the original ``x.strides`` when
  62. calculating new strides to avoid reliance on a contiguous memory
  63. layout.
  64. Furthermore, arrays created with this function often contain self
  65. overlapping memory, so that two elements are identical.
  66. Vectorized write operations on such arrays will typically be
  67. unpredictable. They may even give different results for small, large,
  68. or transposed arrays.
  69. Since writing to these arrays has to be tested and done with great
  70. care, you may want to use ``writeable=False`` to avoid accidental write
  71. operations.
  72. For these reasons it is advisable to avoid ``as_strided`` when
  73. possible.
  74. """
  75. # first convert input to array, possibly keeping subclass
  76. x = np.array(x, copy=False, subok=subok)
  77. interface = dict(x.__array_interface__)
  78. if shape is not None:
  79. interface['shape'] = tuple(shape)
  80. if strides is not None:
  81. interface['strides'] = tuple(strides)
  82. array = np.asarray(DummyArray(interface, base=x))
  83. # The route via `__interface__` does not preserve structured
  84. # dtypes. Since dtype should remain unchanged, we set it explicitly.
  85. array.dtype = x.dtype
  86. view = _maybe_view_as_subclass(x, array)
  87. if view.flags.writeable and not writeable:
  88. view.flags.writeable = False
  89. return view
  90. def _broadcast_to(array, shape, subok, readonly):
  91. shape = tuple(shape) if np.iterable(shape) else (shape,)
  92. array = np.array(array, copy=False, subok=subok)
  93. if not shape and array.shape:
  94. raise ValueError('cannot broadcast a non-scalar to a scalar array')
  95. if any(size < 0 for size in shape):
  96. raise ValueError('all elements of broadcast shape must be non-'
  97. 'negative')
  98. needs_writeable = not readonly and array.flags.writeable
  99. extras = ['reduce_ok'] if needs_writeable else []
  100. op_flag = 'readwrite' if needs_writeable else 'readonly'
  101. it = np.nditer(
  102. (array,), flags=['multi_index', 'refs_ok', 'zerosize_ok'] + extras,
  103. op_flags=[op_flag], itershape=shape, order='C')
  104. with it:
  105. # never really has writebackifcopy semantics
  106. broadcast = it.itviews[0]
  107. result = _maybe_view_as_subclass(array, broadcast)
  108. if needs_writeable and not result.flags.writeable:
  109. result.flags.writeable = True
  110. return result
  111. def _broadcast_to_dispatcher(array, shape, subok=None):
  112. return (array,)
  113. @array_function_dispatch(_broadcast_to_dispatcher, module='numpy')
  114. def broadcast_to(array, shape, subok=False):
  115. """Broadcast an array to a new shape.
  116. Parameters
  117. ----------
  118. array : array_like
  119. The array to broadcast.
  120. shape : tuple
  121. The shape of the desired array.
  122. subok : bool, optional
  123. If True, then sub-classes will be passed-through, otherwise
  124. the returned array will be forced to be a base-class array (default).
  125. Returns
  126. -------
  127. broadcast : array
  128. A readonly view on the original array with the given shape. It is
  129. typically not contiguous. Furthermore, more than one element of a
  130. broadcasted array may refer to a single memory location.
  131. Raises
  132. ------
  133. ValueError
  134. If the array is not compatible with the new shape according to NumPy's
  135. broadcasting rules.
  136. Notes
  137. -----
  138. .. versionadded:: 1.10.0
  139. Examples
  140. --------
  141. >>> x = np.array([1, 2, 3])
  142. >>> np.broadcast_to(x, (3, 3))
  143. array([[1, 2, 3],
  144. [1, 2, 3],
  145. [1, 2, 3]])
  146. """
  147. return _broadcast_to(array, shape, subok=subok, readonly=True)
  148. def _broadcast_shape(*args):
  149. """Returns the shape of the arrays that would result from broadcasting the
  150. supplied arrays against each other.
  151. """
  152. if not args:
  153. return ()
  154. # use the old-iterator because np.nditer does not handle size 0 arrays
  155. # consistently
  156. b = np.broadcast(*args[:32])
  157. # unfortunately, it cannot handle 32 or more arguments directly
  158. for pos in range(32, len(args), 31):
  159. # ironically, np.broadcast does not properly handle np.broadcast
  160. # objects (it treats them as scalars)
  161. # use broadcasting to avoid allocating the full array
  162. b = broadcast_to(0, b.shape)
  163. b = np.broadcast(b, *args[pos:(pos + 31)])
  164. return b.shape
  165. def _broadcast_arrays_dispatcher(*args, **kwargs):
  166. return args
  167. @array_function_dispatch(_broadcast_arrays_dispatcher, module='numpy')
  168. def broadcast_arrays(*args, **kwargs):
  169. """
  170. Broadcast any number of arrays against each other.
  171. Parameters
  172. ----------
  173. `*args` : array_likes
  174. The arrays to broadcast.
  175. subok : bool, optional
  176. If True, then sub-classes will be passed-through, otherwise
  177. the returned arrays will be forced to be a base-class array (default).
  178. Returns
  179. -------
  180. broadcasted : list of arrays
  181. These arrays are views on the original arrays. They are typically
  182. not contiguous. Furthermore, more than one element of a
  183. broadcasted array may refer to a single memory location. If you
  184. need to write to the arrays, make copies first.
  185. Examples
  186. --------
  187. >>> x = np.array([[1,2,3]])
  188. >>> y = np.array([[4],[5]])
  189. >>> np.broadcast_arrays(x, y)
  190. [array([[1, 2, 3],
  191. [1, 2, 3]]), array([[4, 4, 4],
  192. [5, 5, 5]])]
  193. Here is a useful idiom for getting contiguous copies instead of
  194. non-contiguous views.
  195. >>> [np.array(a) for a in np.broadcast_arrays(x, y)]
  196. [array([[1, 2, 3],
  197. [1, 2, 3]]), array([[4, 4, 4],
  198. [5, 5, 5]])]
  199. """
  200. # nditer is not used here to avoid the limit of 32 arrays.
  201. # Otherwise, something like the following one-liner would suffice:
  202. # return np.nditer(args, flags=['multi_index', 'zerosize_ok'],
  203. # order='C').itviews
  204. subok = kwargs.pop('subok', False)
  205. if kwargs:
  206. raise TypeError('broadcast_arrays() got an unexpected keyword '
  207. 'argument {!r}'.format(list(kwargs.keys())[0]))
  208. args = [np.array(_m, copy=False, subok=subok) for _m in args]
  209. shape = _broadcast_shape(*args)
  210. if all(array.shape == shape for array in args):
  211. # Common case where nothing needs to be broadcasted.
  212. return args
  213. # TODO: consider making the results of broadcast_arrays readonly to match
  214. # broadcast_to. This will require a deprecation cycle.
  215. return [_broadcast_to(array, shape, subok=subok, readonly=False)
  216. for array in args]