stride_tricks.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545
  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. import numpy as np
  7. from numpy.core.numeric import normalize_axis_tuple
  8. from numpy.core.overrides import array_function_dispatch, set_module
  9. __all__ = ['broadcast_to', 'broadcast_arrays', 'broadcast_shapes']
  10. class DummyArray:
  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. lib.stride_tricks.sliding_window_view :
  56. userfriendly and safe function for the creation of sliding window views.
  57. Notes
  58. -----
  59. ``as_strided`` creates a view into the array given the exact strides
  60. and shape. This means it manipulates the internal data structure of
  61. ndarray and, if done incorrectly, the array elements can point to
  62. invalid memory and can corrupt results or crash your program.
  63. It is advisable to always use the original ``x.strides`` when
  64. calculating new strides to avoid reliance on a contiguous memory
  65. layout.
  66. Furthermore, arrays created with this function often contain self
  67. overlapping memory, so that two elements are identical.
  68. Vectorized write operations on such arrays will typically be
  69. unpredictable. They may even give different results for small, large,
  70. or transposed arrays.
  71. Since writing to these arrays has to be tested and done with great
  72. care, you may want to use ``writeable=False`` to avoid accidental write
  73. operations.
  74. For these reasons it is advisable to avoid ``as_strided`` when
  75. possible.
  76. """
  77. # first convert input to array, possibly keeping subclass
  78. x = np.array(x, copy=False, subok=subok)
  79. interface = dict(x.__array_interface__)
  80. if shape is not None:
  81. interface['shape'] = tuple(shape)
  82. if strides is not None:
  83. interface['strides'] = tuple(strides)
  84. array = np.asarray(DummyArray(interface, base=x))
  85. # The route via `__interface__` does not preserve structured
  86. # dtypes. Since dtype should remain unchanged, we set it explicitly.
  87. array.dtype = x.dtype
  88. view = _maybe_view_as_subclass(x, array)
  89. if view.flags.writeable and not writeable:
  90. view.flags.writeable = False
  91. return view
  92. def _sliding_window_view_dispatcher(x, window_shape, axis=None, *,
  93. subok=None, writeable=None):
  94. return (x,)
  95. @array_function_dispatch(_sliding_window_view_dispatcher)
  96. def sliding_window_view(x, window_shape, axis=None, *,
  97. subok=False, writeable=False):
  98. """
  99. Create a sliding window view into the array with the given window shape.
  100. Also known as rolling or moving window, the window slides across all
  101. dimensions of the array and extracts subsets of the array at all window
  102. positions.
  103. .. versionadded:: 1.20.0
  104. Parameters
  105. ----------
  106. x : array_like
  107. Array to create the sliding window view from.
  108. window_shape : int or tuple of int
  109. Size of window over each axis that takes part in the sliding window.
  110. If `axis` is not present, must have same length as the number of input
  111. array dimensions. Single integers `i` are treated as if they were the
  112. tuple `(i,)`.
  113. axis : int or tuple of int, optional
  114. Axis or axes along which the sliding window is applied.
  115. By default, the sliding window is applied to all axes and
  116. `window_shape[i]` will refer to axis `i` of `x`.
  117. If `axis` is given as a `tuple of int`, `window_shape[i]` will refer to
  118. the axis `axis[i]` of `x`.
  119. Single integers `i` are treated as if they were the tuple `(i,)`.
  120. subok : bool, optional
  121. If True, sub-classes will be passed-through, otherwise the returned
  122. array will be forced to be a base-class array (default).
  123. writeable : bool, optional
  124. When true, allow writing to the returned view. The default is false,
  125. as this should be used with caution: the returned view contains the
  126. same memory location multiple times, so writing to one location will
  127. cause others to change.
  128. Returns
  129. -------
  130. view : ndarray
  131. Sliding window view of the array. The sliding window dimensions are
  132. inserted at the end, and the original dimensions are trimmed as
  133. required by the size of the sliding window.
  134. That is, ``view.shape = x_shape_trimmed + window_shape``, where
  135. ``x_shape_trimmed`` is ``x.shape`` with every entry reduced by one less
  136. than the corresponding window size.
  137. See Also
  138. --------
  139. lib.stride_tricks.as_strided: A lower-level and less safe routine for
  140. creating arbitrary views from custom shape and strides.
  141. broadcast_to: broadcast an array to a given shape.
  142. Notes
  143. -----
  144. For many applications using a sliding window view can be convenient, but
  145. potentially very slow. Often specialized solutions exist, for example:
  146. - `scipy.signal.fftconvolve`
  147. - filtering functions in `scipy.ndimage`
  148. - moving window functions provided by
  149. `bottleneck <https://github.com/pydata/bottleneck>`_.
  150. As a rough estimate, a sliding window approach with an input size of `N`
  151. and a window size of `W` will scale as `O(N*W)` where frequently a special
  152. algorithm can achieve `O(N)`. That means that the sliding window variant
  153. for a window size of 100 can be a 100 times slower than a more specialized
  154. version.
  155. Nevertheless, for small window sizes, when no custom algorithm exists, or
  156. as a prototyping and developing tool, this function can be a good solution.
  157. Examples
  158. --------
  159. >>> x = np.arange(6)
  160. >>> x.shape
  161. (6,)
  162. >>> v = sliding_window_view(x, 3)
  163. >>> v.shape
  164. (4, 3)
  165. >>> v
  166. array([[0, 1, 2],
  167. [1, 2, 3],
  168. [2, 3, 4],
  169. [3, 4, 5]])
  170. This also works in more dimensions, e.g.
  171. >>> i, j = np.ogrid[:3, :4]
  172. >>> x = 10*i + j
  173. >>> x.shape
  174. (3, 4)
  175. >>> x
  176. array([[ 0, 1, 2, 3],
  177. [10, 11, 12, 13],
  178. [20, 21, 22, 23]])
  179. >>> shape = (2,2)
  180. >>> v = sliding_window_view(x, shape)
  181. >>> v.shape
  182. (2, 3, 2, 2)
  183. >>> v
  184. array([[[[ 0, 1],
  185. [10, 11]],
  186. [[ 1, 2],
  187. [11, 12]],
  188. [[ 2, 3],
  189. [12, 13]]],
  190. [[[10, 11],
  191. [20, 21]],
  192. [[11, 12],
  193. [21, 22]],
  194. [[12, 13],
  195. [22, 23]]]])
  196. The axis can be specified explicitly:
  197. >>> v = sliding_window_view(x, 3, 0)
  198. >>> v.shape
  199. (1, 4, 3)
  200. >>> v
  201. array([[[ 0, 10, 20],
  202. [ 1, 11, 21],
  203. [ 2, 12, 22],
  204. [ 3, 13, 23]]])
  205. The same axis can be used several times. In that case, every use reduces
  206. the corresponding original dimension:
  207. >>> v = sliding_window_view(x, (2, 3), (1, 1))
  208. >>> v.shape
  209. (3, 1, 2, 3)
  210. >>> v
  211. array([[[[ 0, 1, 2],
  212. [ 1, 2, 3]]],
  213. [[[10, 11, 12],
  214. [11, 12, 13]]],
  215. [[[20, 21, 22],
  216. [21, 22, 23]]]])
  217. Combining with stepped slicing (`::step`), this can be used to take sliding
  218. views which skip elements:
  219. >>> x = np.arange(7)
  220. >>> sliding_window_view(x, 5)[:, ::2]
  221. array([[0, 2, 4],
  222. [1, 3, 5],
  223. [2, 4, 6]])
  224. or views which move by multiple elements
  225. >>> x = np.arange(7)
  226. >>> sliding_window_view(x, 3)[::2, :]
  227. array([[0, 1, 2],
  228. [2, 3, 4],
  229. [4, 5, 6]])
  230. A common application of `sliding_window_view` is the calculation of running
  231. statistics. The simplest example is the
  232. `moving average <https://en.wikipedia.org/wiki/Moving_average>`_:
  233. >>> x = np.arange(6)
  234. >>> x.shape
  235. (6,)
  236. >>> v = sliding_window_view(x, 3)
  237. >>> v.shape
  238. (4, 3)
  239. >>> v
  240. array([[0, 1, 2],
  241. [1, 2, 3],
  242. [2, 3, 4],
  243. [3, 4, 5]])
  244. >>> moving_average = v.mean(axis=-1)
  245. >>> moving_average
  246. array([1., 2., 3., 4.])
  247. Note that a sliding window approach is often **not** optimal (see Notes).
  248. """
  249. window_shape = (tuple(window_shape)
  250. if np.iterable(window_shape)
  251. else (window_shape,))
  252. # first convert input to array, possibly keeping subclass
  253. x = np.array(x, copy=False, subok=subok)
  254. window_shape_array = np.array(window_shape)
  255. if np.any(window_shape_array < 0):
  256. raise ValueError('`window_shape` cannot contain negative values')
  257. if axis is None:
  258. axis = tuple(range(x.ndim))
  259. if len(window_shape) != len(axis):
  260. raise ValueError(f'Since axis is `None`, must provide '
  261. f'window_shape for all dimensions of `x`; '
  262. f'got {len(window_shape)} window_shape elements '
  263. f'and `x.ndim` is {x.ndim}.')
  264. else:
  265. axis = normalize_axis_tuple(axis, x.ndim, allow_duplicate=True)
  266. if len(window_shape) != len(axis):
  267. raise ValueError(f'Must provide matching length window_shape and '
  268. f'axis; got {len(window_shape)} window_shape '
  269. f'elements and {len(axis)} axes elements.')
  270. out_strides = x.strides + tuple(x.strides[ax] for ax in axis)
  271. # note: same axis can be windowed repeatedly
  272. x_shape_trimmed = list(x.shape)
  273. for ax, dim in zip(axis, window_shape):
  274. if x_shape_trimmed[ax] < dim:
  275. raise ValueError(
  276. 'window shape cannot be larger than input array shape')
  277. x_shape_trimmed[ax] -= dim - 1
  278. out_shape = tuple(x_shape_trimmed) + window_shape
  279. return as_strided(x, strides=out_strides, shape=out_shape,
  280. subok=subok, writeable=writeable)
  281. def _broadcast_to(array, shape, subok, readonly):
  282. shape = tuple(shape) if np.iterable(shape) else (shape,)
  283. array = np.array(array, copy=False, subok=subok)
  284. if not shape and array.shape:
  285. raise ValueError('cannot broadcast a non-scalar to a scalar array')
  286. if any(size < 0 for size in shape):
  287. raise ValueError('all elements of broadcast shape must be non-'
  288. 'negative')
  289. extras = []
  290. it = np.nditer(
  291. (array,), flags=['multi_index', 'refs_ok', 'zerosize_ok'] + extras,
  292. op_flags=['readonly'], itershape=shape, order='C')
  293. with it:
  294. # never really has writebackifcopy semantics
  295. broadcast = it.itviews[0]
  296. result = _maybe_view_as_subclass(array, broadcast)
  297. # In a future version this will go away
  298. if not readonly and array.flags._writeable_no_warn:
  299. result.flags.writeable = True
  300. result.flags._warn_on_write = True
  301. return result
  302. def _broadcast_to_dispatcher(array, shape, subok=None):
  303. return (array,)
  304. @array_function_dispatch(_broadcast_to_dispatcher, module='numpy')
  305. def broadcast_to(array, shape, subok=False):
  306. """Broadcast an array to a new shape.
  307. Parameters
  308. ----------
  309. array : array_like
  310. The array to broadcast.
  311. shape : tuple
  312. The shape of the desired array.
  313. subok : bool, optional
  314. If True, then sub-classes will be passed-through, otherwise
  315. the returned array will be forced to be a base-class array (default).
  316. Returns
  317. -------
  318. broadcast : array
  319. A readonly view on the original array with the given shape. It is
  320. typically not contiguous. Furthermore, more than one element of a
  321. broadcasted array may refer to a single memory location.
  322. Raises
  323. ------
  324. ValueError
  325. If the array is not compatible with the new shape according to NumPy's
  326. broadcasting rules.
  327. See Also
  328. --------
  329. broadcast
  330. broadcast_arrays
  331. broadcast_shapes
  332. Notes
  333. -----
  334. .. versionadded:: 1.10.0
  335. Examples
  336. --------
  337. >>> x = np.array([1, 2, 3])
  338. >>> np.broadcast_to(x, (3, 3))
  339. array([[1, 2, 3],
  340. [1, 2, 3],
  341. [1, 2, 3]])
  342. """
  343. return _broadcast_to(array, shape, subok=subok, readonly=True)
  344. def _broadcast_shape(*args):
  345. """Returns the shape of the arrays that would result from broadcasting the
  346. supplied arrays against each other.
  347. """
  348. # use the old-iterator because np.nditer does not handle size 0 arrays
  349. # consistently
  350. b = np.broadcast(*args[:32])
  351. # unfortunately, it cannot handle 32 or more arguments directly
  352. for pos in range(32, len(args), 31):
  353. # ironically, np.broadcast does not properly handle np.broadcast
  354. # objects (it treats them as scalars)
  355. # use broadcasting to avoid allocating the full array
  356. b = broadcast_to(0, b.shape)
  357. b = np.broadcast(b, *args[pos:(pos + 31)])
  358. return b.shape
  359. @set_module('numpy')
  360. def broadcast_shapes(*args):
  361. """
  362. Broadcast the input shapes into a single shape.
  363. :ref:`Learn more about broadcasting here <basics.broadcasting>`.
  364. .. versionadded:: 1.20.0
  365. Parameters
  366. ----------
  367. `*args` : tuples of ints, or ints
  368. The shapes to be broadcast against each other.
  369. Returns
  370. -------
  371. tuple
  372. Broadcasted shape.
  373. Raises
  374. ------
  375. ValueError
  376. If the shapes are not compatible and cannot be broadcast according
  377. to NumPy's broadcasting rules.
  378. See Also
  379. --------
  380. broadcast
  381. broadcast_arrays
  382. broadcast_to
  383. Examples
  384. --------
  385. >>> np.broadcast_shapes((1, 2), (3, 1), (3, 2))
  386. (3, 2)
  387. >>> np.broadcast_shapes((6, 7), (5, 6, 1), (7,), (5, 1, 7))
  388. (5, 6, 7)
  389. """
  390. arrays = [np.empty(x, dtype=[]) for x in args]
  391. return _broadcast_shape(*arrays)
  392. def _broadcast_arrays_dispatcher(*args, subok=None):
  393. return args
  394. @array_function_dispatch(_broadcast_arrays_dispatcher, module='numpy')
  395. def broadcast_arrays(*args, subok=False):
  396. """
  397. Broadcast any number of arrays against each other.
  398. Parameters
  399. ----------
  400. `*args` : array_likes
  401. The arrays to broadcast.
  402. subok : bool, optional
  403. If True, then sub-classes will be passed-through, otherwise
  404. the returned arrays will be forced to be a base-class array (default).
  405. Returns
  406. -------
  407. broadcasted : list of arrays
  408. These arrays are views on the original arrays. They are typically
  409. not contiguous. Furthermore, more than one element of a
  410. broadcasted array may refer to a single memory location. If you need
  411. to write to the arrays, make copies first. While you can set the
  412. ``writable`` flag True, writing to a single output value may end up
  413. changing more than one location in the output array.
  414. .. deprecated:: 1.17
  415. The output is currently marked so that if written to, a deprecation
  416. warning will be emitted. A future version will set the
  417. ``writable`` flag False so writing to it will raise an error.
  418. See Also
  419. --------
  420. broadcast
  421. broadcast_to
  422. broadcast_shapes
  423. Examples
  424. --------
  425. >>> x = np.array([[1,2,3]])
  426. >>> y = np.array([[4],[5]])
  427. >>> np.broadcast_arrays(x, y)
  428. [array([[1, 2, 3],
  429. [1, 2, 3]]), array([[4, 4, 4],
  430. [5, 5, 5]])]
  431. Here is a useful idiom for getting contiguous copies instead of
  432. non-contiguous views.
  433. >>> [np.array(a) for a in np.broadcast_arrays(x, y)]
  434. [array([[1, 2, 3],
  435. [1, 2, 3]]), array([[4, 4, 4],
  436. [5, 5, 5]])]
  437. """
  438. # nditer is not used here to avoid the limit of 32 arrays.
  439. # Otherwise, something like the following one-liner would suffice:
  440. # return np.nditer(args, flags=['multi_index', 'zerosize_ok'],
  441. # order='C').itviews
  442. args = [np.array(_m, copy=False, subok=subok) for _m in args]
  443. shape = _broadcast_shape(*args)
  444. if all(array.shape == shape for array in args):
  445. # Common case where nothing needs to be broadcasted.
  446. return args
  447. return [_broadcast_to(array, shape, subok=subok, readonly=False)
  448. for array in args]