codekingpro/portable-devtools
115k
1__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
2 'stack', 'unstack', 'vstack']
3
4import functools
5import itertools
6import operator
7
8from . import fromnumeric as _from_nx, numeric as _nx, overrides
9from .multiarray import array, asanyarray, normalize_axis_index
10
11array_function_dispatch = functools.partial(
12 overrides.array_function_dispatch, module='numpy')
13
14
15def _atleast_1d_dispatcher(*arys):
16 return arys
17
18
19@array_function_dispatch(_atleast_1d_dispatcher)
20def atleast_1d(*arys):
21 """
22 Convert inputs to arrays with at least one dimension.
23
24 Scalar inputs are converted to 1-dimensional arrays, whilst
25 higher-dimensional inputs are preserved.
26
27 Parameters
28 ----------
29 arys1, arys2, ... : array_like
30 One or more input arrays.
31
32 Returns
33 -------
34 ret : ndarray
35 An array, or tuple of arrays, each with ``a.ndim >= 1``.
36 Copies are made only if necessary.
37
38 See Also
39 --------
40 atleast_2d, atleast_3d
41
42 Examples
43 --------
44 >>> import numpy as np
45 >>> np.atleast_1d(1.0)
46 array([1.])
47
48 >>> x = np.arange(9.0).reshape(3,3)
49 >>> np.atleast_1d(x)
50 array([[0., 1., 2.],
51 [3., 4., 5.],
52 [6., 7., 8.]])
53 >>> np.atleast_1d(x) is x
54 True
55
56 >>> np.atleast_1d(1, [3, 4])
57 (array([1]), array([3, 4]))
58
59 """
60 if len(arys) == 1:
61 result = asanyarray(arys[0])
62 if result.ndim == 0:
63 result = result.reshape(1)
64 return result
65 res = []
66 for ary in arys:
67 result = asanyarray(ary)
68 if result.ndim == 0:
69 result = result.reshape(1)
70 res.append(result)
71 return tuple(res)
72
73
74def _atleast_2d_dispatcher(*arys):
75 return arys
76
77
78@array_function_dispatch(_atleast_2d_dispatcher)
79def atleast_2d(*arys):
80 """
81 View inputs as arrays with at least two dimensions.
82
83 Parameters
84 ----------
85 arys1, arys2, ... : array_like
86 One or more array-like sequences. Non-array inputs are converted
87 to arrays. Arrays that already have two or more dimensions are
88 preserved.
89
90 Returns
91 -------
92 res, res2, ... : ndarray
93 An array, or tuple of arrays, each with ``a.ndim >= 2``.
94 Copies are avoided where possible, and views with two or more
95 dimensions are returned.
96
97 See Also
98 --------
99 atleast_1d, atleast_3d
100
101 Examples
102 --------
103 >>> import numpy as np
104 >>> np.atleast_2d(3.0)
105 array([[3.]])
106
107 >>> x = np.arange(3.0)
108 >>> np.atleast_2d(x)
109 array([[0., 1., 2.]])
110 >>> np.atleast_2d(x).base is x
111 True
112
113 >>> np.atleast_2d(1, [1, 2], [[1, 2]])
114 (array([[1]]), array([[1, 2]]), array([[1, 2]]))
115
116 """
117 res = []
118 for ary in arys:
119 ary = asanyarray(ary)
120 if ary.ndim == 0:
121 result = ary.reshape(1, 1)
122 elif ary.ndim == 1:
123 result = ary[_nx.newaxis, :]
124 else:
125 result = ary
126 res.append(result)
127 if len(res) == 1:
128 return res[0]
129 else:
130 return tuple(res)
131
132
133def _atleast_3d_dispatcher(*arys):
134 return arys
135
136
137@array_function_dispatch(_atleast_3d_dispatcher)
138def atleast_3d(*arys):
139 """
140 View inputs as arrays with at least three dimensions.
141
142 Parameters
143 ----------
144 arys1, arys2, ... : array_like
145 One or more array-like sequences. Non-array inputs are converted to
146 arrays. Arrays that already have three or more dimensions are
147 preserved.
148
149 Returns
150 -------
151 res1, res2, ... : ndarray
152 An array, or tuple of arrays, each with ``a.ndim >= 3``. Copies are
153 avoided where possible, and views with three or more dimensions are
154 returned. For example, a 1-D array of shape ``(N,)`` becomes a view
155 of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a
156 view of shape ``(M, N, 1)``.
157
158 See Also
159 --------
160 atleast_1d, atleast_2d
161
162 Examples
163 --------
164 >>> import numpy as np
165 >>> np.atleast_3d(3.0)
166 array([[[3.]]])
167
168 >>> x = np.arange(3.0)
169 >>> np.atleast_3d(x).shape
170 (1, 3, 1)
171
172 >>> x = np.arange(12.0).reshape(4,3)
173 >>> np.atleast_3d(x).shape
174 (4, 3, 1)
175 >>> np.atleast_3d(x).base is x.base # x is a reshape, so not base itself
176 True
177
178 >>> for arr in np.atleast_3d([1, 2], [[1, 2]], [[[1, 2]]]):
179 ... print(arr, arr.shape) # doctest: +SKIP
180 ...
181 [[[1]
182 [2]]] (1, 2, 1)
183 [[[1]
184 [2]]] (1, 2, 1)
185 [[[1 2]]] (1, 1, 2)
186
187 """
188 res = []
189 for ary in arys:
190 ary = asanyarray(ary)
191 if ary.ndim == 0:
192 result = ary.reshape(1, 1, 1)
193 elif ary.ndim == 1:
194 result = ary[_nx.newaxis, :, _nx.newaxis]
195 elif ary.ndim == 2:
196 result = ary[:, :, _nx.newaxis]
197 else:
198 result = ary
199 res.append(result)
200 if len(res) == 1:
201 return res[0]
202 else:
203 return tuple(res)
204
205
206def _arrays_for_stack_dispatcher(arrays):
207 if not hasattr(arrays, "__getitem__"):
208 raise TypeError('arrays to stack must be passed as a "sequence" type '
209 'such as list or tuple.')
210
211 return tuple(arrays)
212
213
214def _vhstack_dispatcher(tup, *, dtype=None, casting=None):
215 return _arrays_for_stack_dispatcher(tup)
216
217
218@array_function_dispatch(_vhstack_dispatcher)
219def vstack(tup, *, dtype=None, casting="same_kind"):
220 """
221 Stack arrays in sequence vertically (row wise).
222
223 This is equivalent to concatenation along the first axis after 1-D arrays
224 of shape `(N,)` have been reshaped to `(1,N)`. Rebuilds arrays divided by
225 `vsplit`.
226
227 This function makes most sense for arrays with up to 3 dimensions. For
228 instance, for pixel-data with a height (first axis), width (second axis),
229 and r/g/b channels (third axis). The functions `concatenate`, `stack` and
230 `block` provide more general stacking and concatenation operations.
231
232 Parameters
233 ----------
234 tup : sequence of ndarrays
235 The arrays must have the same shape along all but the first axis.
236 1-D arrays must have the same length. In the case of a single
237 array_like input, it will be treated as a sequence of arrays; i.e.,
238 each element along the zeroth axis is treated as a separate array.
239
240 dtype : str or dtype
241 If provided, the destination array will have this dtype. Cannot be
242 provided together with `out`.
243
244 .. versionadded:: 1.24
245
246 casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
247 Controls what kind of data casting may occur. Defaults to 'same_kind'.
248
249 .. versionadded:: 1.24
250
251 Returns
252 -------
253 stacked : ndarray
254 The array formed by stacking the given arrays, will be at least 2-D.
255
256 See Also
257 --------
258 concatenate : Join a sequence of arrays along an existing axis.
259 stack : Join a sequence of arrays along a new axis.
260 block : Assemble an nd-array from nested lists of blocks.
261 hstack : Stack arrays in sequence horizontally (column wise).
262 dstack : Stack arrays in sequence depth wise (along third axis).
263 column_stack : Stack 1-D arrays as columns into a 2-D array.
264 vsplit : Split an array into multiple sub-arrays vertically (row-wise).
265 unstack : Split an array into a tuple of sub-arrays along an axis.
266
267 Examples
268 --------
269 >>> import numpy as np
270 >>> a = np.array([1, 2, 3])
271 >>> b = np.array([4, 5, 6])
272 >>> np.vstack((a,b))
273 array([[1, 2, 3],
274 [4, 5, 6]])
275
276 >>> a = np.array([[1], [2], [3]])
277 >>> b = np.array([[4], [5], [6]])
278 >>> np.vstack((a,b))
279 array([[1],
280 [2],
281 [3],
282 [4],
283 [5],
284 [6]])
285
286 """
287 arrs = atleast_2d(*tup)
288 if not isinstance(arrs, tuple):
289 arrs = (arrs,)
290 return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting)
291
292
293@array_function_dispatch(_vhstack_dispatcher)
294def hstack(tup, *, dtype=None, casting="same_kind"):
295 """
296 Stack arrays in sequence horizontally (column wise).
297
298 This is equivalent to concatenation along the second axis, except for 1-D
299 arrays where it concatenates along the first axis. Rebuilds arrays divided
300 by `hsplit`.
301
302 This function makes most sense for arrays with up to 3 dimensions. For
303 instance, for pixel-data with a height (first axis), width (second axis),
304 and r/g/b channels (third axis). The functions `concatenate`, `stack` and
305 `block` provide more general stacking and concatenation operations.
306
307 Parameters
308 ----------
309 tup : sequence of ndarrays
310 The arrays must have the same shape along all but the second axis,
311 except 1-D arrays which can be any length. In the case of a single
312 array_like input, it will be treated as a sequence of arrays; i.e.,
313 each element along the zeroth axis is treated as a separate array.
314
315 dtype : str or dtype
316 If provided, the destination array will have this dtype. Cannot be
317 provided together with `out`.
318
319 .. versionadded:: 1.24
320
321 casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
322 Controls what kind of data casting may occur. Defaults to 'same_kind'.
323
324 .. versionadded:: 1.24
325
326 Returns
327 -------
328 stacked : ndarray
329 The array formed by stacking the given arrays.
330
331 See Also
332 --------
333 concatenate : Join a sequence of arrays along an existing axis.
334 stack : Join a sequence of arrays along a new axis.
335 block : Assemble an nd-array from nested lists of blocks.
336 vstack : Stack arrays in sequence vertically (row wise).
337 dstack : Stack arrays in sequence depth wise (along third axis).
338 column_stack : Stack 1-D arrays as columns into a 2-D array.
339 hsplit : Split an array into multiple sub-arrays
340 horizontally (column-wise).
341 unstack : Split an array into a tuple of sub-arrays along an axis.
342
343 Examples
344 --------
345 >>> import numpy as np
346 >>> a = np.array((1,2,3))
347 >>> b = np.array((4,5,6))
348 >>> np.hstack((a,b))
349 array([1, 2, 3, 4, 5, 6])
350 >>> a = np.array([[1],[2],[3]])
351 >>> b = np.array([[4],[5],[6]])
352 >>> np.hstack((a,b))
353 array([[1, 4],
354 [2, 5],
355 [3, 6]])
356
357 """
358 arrs = atleast_1d(*tup)
359 if not isinstance(arrs, tuple):
360 arrs = (arrs,)
361 # As a special case, dimension 0 of 1-dimensional arrays is "horizontal"
362 if arrs and arrs[0].ndim == 1:
363 return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting)
364 else:
365 return _nx.concatenate(arrs, 1, dtype=dtype, casting=casting)
366
367
368def _stack_dispatcher(arrays, axis=None, out=None, *,
369 dtype=None, casting=None):
370 arrays = _arrays_for_stack_dispatcher(arrays)
371 if out is not None:
372 # optimize for the typical case where only arrays is provided
373 arrays = list(arrays)
374 arrays.append(out)
375 return arrays
376
377
378@array_function_dispatch(_stack_dispatcher)
379def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"):
380 """
381 Join a sequence of arrays along a new axis.
382
383 The ``axis`` parameter specifies the index of the new axis in the
384 dimensions of the result. For example, if ``axis=0`` it will be the first
385 dimension and if ``axis=-1`` it will be the last dimension.
386
387 Parameters
388 ----------
389 arrays : sequence of ndarrays
390 Each array must have the same shape. In the case of a single ndarray
391 array_like input, it will be treated as a sequence of arrays; i.e.,
392 each element along the zeroth axis is treated as a separate array.
393
394 axis : int, optional
395 The axis in the result array along which the input arrays are stacked.
396
397 out : ndarray, optional
398 If provided, the destination to place the result. The shape must be
399 correct, matching that of what stack would have returned if no
400 out argument were specified.
401
402 dtype : str or dtype
403 If provided, the destination array will have this dtype. Cannot be
404 provided together with `out`.
405
406 .. versionadded:: 1.24
407
408 casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
409 Controls what kind of data casting may occur. Defaults to 'same_kind'.
410
411 .. versionadded:: 1.24
412
413
414 Returns
415 -------
416 stacked : ndarray
417 The stacked array has one more dimension than the input arrays.
418
419 See Also
420 --------
421 concatenate : Join a sequence of arrays along an existing axis.
422 block : Assemble an nd-array from nested lists of blocks.
423 split : Split array into a list of multiple sub-arrays of equal size.
424 unstack : Split an array into a tuple of sub-arrays along an axis.
425
426 Examples
427 --------
428 >>> import numpy as np
429 >>> rng = np.random.default_rng()
430 >>> arrays = [rng.normal(size=(3,4)) for _ in range(10)]
431 >>> np.stack(arrays, axis=0).shape
432 (10, 3, 4)
433
434 >>> np.stack(arrays, axis=1).shape
435 (3, 10, 4)
436
437 >>> np.stack(arrays, axis=2).shape
438 (3, 4, 10)
439
440 >>> a = np.array([1, 2, 3])
441 >>> b = np.array([4, 5, 6])
442 >>> np.stack((a, b))
443 array([[1, 2, 3],
444 [4, 5, 6]])
445
446 >>> np.stack((a, b), axis=-1)
447 array([[1, 4],
448 [2, 5],
449 [3, 6]])
450
451 """
452 arrays = [asanyarray(arr) for arr in arrays]
453 if not arrays:
454 raise ValueError('need at least one array to stack')
455
456 shapes = {arr.shape for arr in arrays}
457 if len(shapes) != 1:
458 raise ValueError('all input arrays must have the same shape')
459
460 result_ndim = arrays[0].ndim + 1
461 axis = normalize_axis_index(axis, result_ndim)
462
463 sl = (slice(None),) * axis + (_nx.newaxis,)
464 expanded_arrays = [arr[sl] for arr in arrays]
465 return _nx.concatenate(expanded_arrays, axis=axis, out=out,
466 dtype=dtype, casting=casting)
467
468def _unstack_dispatcher(x, /, *, axis=None):
469 return (x,)
470
471@array_function_dispatch(_unstack_dispatcher)
472def unstack(x, /, *, axis=0):
473 """
474 Split an array into a sequence of arrays along the given axis.
475
476 The ``axis`` parameter specifies the dimension along which the array will
477 be split. For example, if ``axis=0`` (the default) it will be the first
478 dimension and if ``axis=-1`` it will be the last dimension.
479
480 The result is a tuple of arrays split along ``axis``.
481
482 .. versionadded:: 2.1.0
483
484 Parameters
485 ----------
486 x : ndarray
487 The array to be unstacked.
488 axis : int, optional
489 Axis along which the array will be split. Default: ``0``.
490
491 Returns
492 -------
493 unstacked : tuple of ndarrays
494 The unstacked arrays.
495
496 See Also
497 --------
498 stack : Join a sequence of arrays along a new axis.
499 concatenate : Join a sequence of arrays along an existing axis.
500 block : Assemble an nd-array from nested lists of blocks.
501 split : Split array into a list of multiple sub-arrays of equal size.
502
503 Notes
504 -----
505 ``unstack`` serves as the reverse operation of :py:func:`stack`, i.e.,
506 ``stack(unstack(x, axis=axis), axis=axis) == x``.
507
508 This function is equivalent to ``tuple(np.moveaxis(x, axis, 0))``, since
509 iterating on an array iterates along the first axis.
510
511 Examples
512 --------
513 >>> arr = np.arange(24).reshape((2, 3, 4))
514 >>> np.unstack(arr)
515 (array([[ 0, 1, 2, 3],
516 [ 4, 5, 6, 7],
517 [ 8, 9, 10, 11]]),
518 array([[12, 13, 14, 15],
519 [16, 17, 18, 19],
520 [20, 21, 22, 23]]))
521 >>> np.unstack(arr, axis=1)
522 (array([[ 0, 1, 2, 3],
523 [12, 13, 14, 15]]),
524 array([[ 4, 5, 6, 7],
525 [16, 17, 18, 19]]),
526 array([[ 8, 9, 10, 11],
527 [20, 21, 22, 23]]))
528 >>> arr2 = np.stack(np.unstack(arr, axis=1), axis=1)
529 >>> arr2.shape
530 (2, 3, 4)
531 >>> np.all(arr == arr2)
532 np.True_
533
534 """
535 if x.ndim == 0:
536 raise ValueError("Input array must be at least 1-d.")
537 return tuple(_nx.moveaxis(x, axis, 0))
538
539
540# Internal functions to eliminate the overhead of repeated dispatch in one of
541# the two possible paths inside np.block.
542# Use getattr to protect against __array_function__ being disabled.
543_size = getattr(_from_nx.size, '__wrapped__', _from_nx.size)
544_ndim = getattr(_from_nx.ndim, '__wrapped__', _from_nx.ndim)
545_concatenate = getattr(_from_nx.concatenate,
546 '__wrapped__', _from_nx.concatenate)
547
548
549def _block_format_index(index):
550 """
551 Convert a list of indices ``[0, 1, 2]`` into ``"arrays[0][1][2]"``.
552 """
553 idx_str = ''.join(f'[{i}]' for i in index if i is not None)
554 return 'arrays' + idx_str
555
556
557def _block_check_depths_match(arrays, parent_index=[]):
558 """
559 Recursive function checking that the depths of nested lists in `arrays`
560 all match. Mismatch raises a ValueError as described in the block
561 docstring below.
562
563 The entire index (rather than just the depth) needs to be calculated
564 for each innermost list, in case an error needs to be raised, so that
565 the index of the offending list can be printed as part of the error.
566
567 Parameters
568 ----------
569 arrays : nested list of arrays
570 The arrays to check
571 parent_index : list of int
572 The full index of `arrays` within the nested lists passed to
573 `_block_check_depths_match` at the top of the recursion.
574
575 Returns
576 -------
577 first_index : list of int
578 The full index of an element from the bottom of the nesting in
579 `arrays`. If any element at the bottom is an empty list, this will
580 refer to it, and the last index along the empty axis will be None.
581 max_arr_ndim : int
582 The maximum of the ndims of the arrays nested in `arrays`.
583 final_size: int
584 The number of elements in the final array. This is used the motivate
585 the choice of algorithm used using benchmarking wisdom.
586
587 """
588 if isinstance(arrays, tuple):
589 # not strictly necessary, but saves us from:
590 # - more than one way to do things - no point treating tuples like
591 # lists
592 # - horribly confusing behaviour that results when tuples are
593 # treated like ndarray
594 raise TypeError(
595 f'{_block_format_index(parent_index)} is a tuple. '
596 'Only lists can be used to arrange blocks, and np.block does '
597 'not allow implicit conversion from tuple to ndarray.'
598 )
599 elif isinstance(arrays, list) and len(arrays) > 0:
600 idxs_ndims = (_block_check_depths_match(arr, parent_index + [i])
601 for i, arr in enumerate(arrays))
602
603 first_index, max_arr_ndim, final_size = next(idxs_ndims)
604 for index, ndim, size in idxs_ndims:
605 final_size += size
606 if ndim > max_arr_ndim:
607 max_arr_ndim = ndim
608 if len(index) != len(first_index):
609 raise ValueError(
610 "List depths are mismatched. First element was at "
611 f"depth {len(first_index)}, but there is an element at "
612 f"depth {len(index)} ({_block_format_index(index)})"
613 )
614 # propagate our flag that indicates an empty list at the bottom
615 if index[-1] is None:
616 first_index = index
617
618 return first_index, max_arr_ndim, final_size
619 elif isinstance(arrays, list) and len(arrays) == 0:
620 # We've 'bottomed out' on an empty list
621 return parent_index + [None], 0, 0
622 else:
623 # We've 'bottomed out' - arrays is either a scalar or an array
624 size = _size(arrays)
625 return parent_index, _ndim(arrays), size
626
627
628def _atleast_nd(a, ndim):
629 # Ensures `a` has at least `ndim` dimensions by prepending
630 # ones to `a.shape` as necessary
631 return array(a, ndmin=ndim, copy=None, subok=True)
632
633
634def _accumulate(values):
635 return list(itertools.accumulate(values))
636
637
638def _concatenate_shapes(shapes, axis):
639 """Given array shapes, return the resulting shape and slices prefixes.
640
641 These help in nested concatenation.
642
643 Returns
644 -------
645 shape: tuple of int
646 This tuple satisfies::
647
648 shape, _ = _concatenate_shapes([arr.shape for shape in arrs], axis)
649 shape == concatenate(arrs, axis).shape
650
651 slice_prefixes: tuple of (slice(start, end), )
652 For a list of arrays being concatenated, this returns the slice
653 in the larger array at axis that needs to be sliced into.
654
655 For example, the following holds::
656
657 ret = concatenate([a, b, c], axis)
658 _, (sl_a, sl_b, sl_c) = concatenate_slices([a, b, c], axis)
659
660 ret[(slice(None),) * axis + sl_a] == a
661 ret[(slice(None),) * axis + sl_b] == b
662 ret[(slice(None),) * axis + sl_c] == c
663
664 These are called slice prefixes since they are used in the recursive
665 blocking algorithm to compute the left-most slices during the
666 recursion. Therefore, they must be prepended to rest of the slice
667 that was computed deeper in the recursion.
668
669 These are returned as tuples to ensure that they can quickly be added
670 to existing slice tuple without creating a new tuple every time.
671
672 """
673 # Cache a result that will be reused.
674 shape_at_axis = [shape[axis] for shape in shapes]
675
676 # Take a shape, any shape
677 first_shape = shapes[0]
678 first_shape_pre = first_shape[:axis]
679 first_shape_post = first_shape[axis + 1:]
680
681 if any(shape[:axis] != first_shape_pre or
682 shape[axis + 1:] != first_shape_post for shape in shapes):
683 raise ValueError(
684 f'Mismatched array shapes in block along axis {axis}.')
685
686 shape = (first_shape_pre + (sum(shape_at_axis),) + first_shape[axis + 1:])
687
688 offsets_at_axis = _accumulate(shape_at_axis)
689 slice_prefixes = [(slice(start, end),)
690 for start, end in zip([0] + offsets_at_axis,
691 offsets_at_axis)]
692 return shape, slice_prefixes
693
694
695def _block_info_recursion(arrays, max_depth, result_ndim, depth=0):
696 """
697 Returns the shape of the final array, along with a list
698 of slices and a list of arrays that can be used for assignment inside the
699 new array
700
701 Parameters
702 ----------
703 arrays : nested list of arrays
704 The arrays to check
705 max_depth : list of int
706 The number of nested lists
707 result_ndim : int
708 The number of dimensions in thefinal array.
709
710 Returns
711 -------
712 shape : tuple of int
713 The shape that the final array will take on.
714 slices: list of tuple of slices
715 The slices into the full array required for assignment. These are
716 required to be prepended with ``(Ellipsis, )`` to obtain to correct
717 final index.
718 arrays: list of ndarray
719 The data to assign to each slice of the full array
720
721 """
722 if depth < max_depth:
723 shapes, slices, arrays = zip(
724 *[_block_info_recursion(arr, max_depth, result_ndim, depth + 1)
725 for arr in arrays])
726
727 axis = result_ndim - max_depth + depth
728 shape, slice_prefixes = _concatenate_shapes(shapes, axis)
729
730 # Prepend the slice prefix and flatten the slices
731 slices = [slice_prefix + the_slice
732 for slice_prefix, inner_slices in zip(slice_prefixes, slices)
733 for the_slice in inner_slices]
734
735 # Flatten the array list
736 arrays = functools.reduce(operator.add, arrays)
737
738 return shape, slices, arrays
739 else:
740 # We've 'bottomed out' - arrays is either a scalar or an array
741 # type(arrays) is not list
742 # Return the slice and the array inside a list to be consistent with
743 # the recursive case.
744 arr = _atleast_nd(arrays, result_ndim)
745 return arr.shape, [()], [arr]
746
747
748def _block(arrays, max_depth, result_ndim, depth=0):
749 """
750 Internal implementation of block based on repeated concatenation.
751 `arrays` is the argument passed to
752 block. `max_depth` is the depth of nested lists within `arrays` and
753 `result_ndim` is the greatest of the dimensions of the arrays in
754 `arrays` and the depth of the lists in `arrays` (see block docstring
755 for details).
756 """
757 if depth < max_depth:
758 arrs = [_block(arr, max_depth, result_ndim, depth + 1)
759 for arr in arrays]
760 return _concatenate(arrs, axis=-(max_depth - depth))
761 else:
762 # We've 'bottomed out' - arrays is either a scalar or an array
763 # type(arrays) is not list
764 return _atleast_nd(arrays, result_ndim)
765
766
767def _block_dispatcher(arrays):
768 # Use type(...) is list to match the behavior of np.block(), which special
769 # cases list specifically rather than allowing for generic iterables or
770 # tuple. Also, we know that list.__array_function__ will never exist.
771 if isinstance(arrays, list):
772 for subarrays in arrays:
773 yield from _block_dispatcher(subarrays)
774 else:
775 yield arrays
776
777
778@array_function_dispatch(_block_dispatcher)
779def block(arrays):
780 """
781 Assemble an nd-array from nested lists of blocks.
782
783 Blocks in the innermost lists are concatenated (see `concatenate`) along
784 the last dimension (-1), then these are concatenated along the
785 second-last dimension (-2), and so on until the outermost list is reached.
786
787 Blocks can be of any dimension, but will not be broadcasted using
788 the normal rules. Instead, leading axes of size 1 are inserted,
789 to make ``block.ndim`` the same for all blocks. This is primarily useful
790 for working with scalars, and means that code like ``np.block([v, 1])``
791 is valid, where ``v.ndim == 1``.
792
793 When the nested list is two levels deep, this allows block matrices to be
794 constructed from their components.
795
796 Parameters
797 ----------
798 arrays : nested list of array_like or scalars (but not tuples)
799 If passed a single ndarray or scalar (a nested list of depth 0), this
800 is returned unmodified (and not copied).
801
802 Elements shapes must match along the appropriate axes (without
803 broadcasting), but leading 1s will be prepended to the shape as
804 necessary to make the dimensions match.
805
806 Returns
807 -------
808 block_array : ndarray
809 The array assembled from the given blocks.
810
811 The dimensionality of the output is equal to the greatest of:
812
813 * the dimensionality of all the inputs
814 * the depth to which the input list is nested
815
816 Raises
817 ------
818 ValueError
819 * If list depths are mismatched - for instance, ``[[a, b], c]`` is
820 illegal, and should be spelt ``[[a, b], [c]]``
821 * If lists are empty - for instance, ``[[a, b], []]``
822
823 See Also
824 --------
825 concatenate : Join a sequence of arrays along an existing axis.
826 stack : Join a sequence of arrays along a new axis.
827 vstack : Stack arrays in sequence vertically (row wise).
828 hstack : Stack arrays in sequence horizontally (column wise).
829 dstack : Stack arrays in sequence depth wise (along third axis).
830 column_stack : Stack 1-D arrays as columns into a 2-D array.
831 vsplit : Split an array into multiple sub-arrays vertically (row-wise).
832 unstack : Split an array into a tuple of sub-arrays along an axis.
833
834 Notes
835 -----
836 When called with only scalars, ``np.block`` is equivalent to an ndarray
837 call. So ``np.block([[1, 2], [3, 4]])`` is equivalent to
838 ``np.array([[1, 2], [3, 4]])``.
839
840 This function does not enforce that the blocks lie on a fixed grid.
841 ``np.block([[a, b], [c, d]])`` is not restricted to arrays of the form::
842
843 AAAbb
844 AAAbb
845 cccDD
846
847 But is also allowed to produce, for some ``a, b, c, d``::
848
849 AAAbb
850 AAAbb
851 cDDDD
852
853 Since concatenation happens along the last axis first, `block` is *not*
854 capable of producing the following directly::
855
856 AAAbb
857 cccbb
858 cccDD
859
860 Matlab's "square bracket stacking", ``[A, B, ...; p, q, ...]``, is
861 equivalent to ``np.block([[A, B, ...], [p, q, ...]])``.
862
863 Examples
864 --------
865 The most common use of this function is to build a block matrix:
866
867 >>> import numpy as np
868 >>> A = np.eye(2) * 2
869 >>> B = np.eye(3) * 3
870 >>> np.block([
871 ... [A, np.zeros((2, 3))],
872 ... [np.ones((3, 2)), B ]
873 ... ])
874 array([[2., 0., 0., 0., 0.],
875 [0., 2., 0., 0., 0.],
876 [1., 1., 3., 0., 0.],
877 [1., 1., 0., 3., 0.],
878 [1., 1., 0., 0., 3.]])
879
880 With a list of depth 1, `block` can be used as `hstack`:
881
882 >>> np.block([1, 2, 3]) # hstack([1, 2, 3])
883 array([1, 2, 3])
884
885 >>> a = np.array([1, 2, 3])
886 >>> b = np.array([4, 5, 6])
887 >>> np.block([a, b, 10]) # hstack([a, b, 10])
888 array([ 1, 2, 3, 4, 5, 6, 10])
889
890 >>> A = np.ones((2, 2), int)
891 >>> B = 2 * A
892 >>> np.block([A, B]) # hstack([A, B])
893 array([[1, 1, 2, 2],
894 [1, 1, 2, 2]])
895
896 With a list of depth 2, `block` can be used in place of `vstack`:
897
898 >>> a = np.array([1, 2, 3])
899 >>> b = np.array([4, 5, 6])
900 >>> np.block([[a], [b]]) # vstack([a, b])
901 array([[1, 2, 3],
902 [4, 5, 6]])
903
904 >>> A = np.ones((2, 2), int)
905 >>> B = 2 * A
906 >>> np.block([[A], [B]]) # vstack([A, B])
907 array([[1, 1],
908 [1, 1],
909 [2, 2],
910 [2, 2]])
911
912 It can also be used in place of `atleast_1d` and `atleast_2d`:
913
914 >>> a = np.array(0)
915 >>> b = np.array([1])
916 >>> np.block([a]) # atleast_1d(a)
917 array([0])
918 >>> np.block([b]) # atleast_1d(b)
919 array([1])
920
921 >>> np.block([[a]]) # atleast_2d(a)
922 array([[0]])
923 >>> np.block([[b]]) # atleast_2d(b)
924 array([[1]])
925
926
927 """
928 arrays, list_ndim, result_ndim, final_size = _block_setup(arrays)
929
930 # It was found through benchmarking that making an array of final size
931 # around 256x256 was faster by straight concatenation on a
932 # i7-7700HQ processor and dual channel ram 2400MHz.
933 # It didn't seem to matter heavily on the dtype used.
934 #
935 # A 2D array using repeated concatenation requires 2 copies of the array.
936 #
937 # The fastest algorithm will depend on the ratio of CPU power to memory
938 # speed.
939 # One can monitor the results of the benchmark
940 # https://pv.github.io/numpy-bench/#bench_shape_base.Block2D.time_block2d
941 # to tune this parameter until a C version of the `_block_info_recursion`
942 # algorithm is implemented which would likely be faster than the python
943 # version.
944 if list_ndim * final_size > (2 * 512 * 512):
945 return _block_slicing(arrays, list_ndim, result_ndim)
946 else:
947 return _block_concatenate(arrays, list_ndim, result_ndim)
948
949
950# These helper functions are mostly used for testing.
951# They allow us to write tests that directly call `_block_slicing`
952# or `_block_concatenate` without blocking large arrays to force the wisdom
953# to trigger the desired path.
954def _block_setup(arrays):
955 """
956 Returns
957 (`arrays`, list_ndim, result_ndim, final_size)
958 """
959 bottom_index, arr_ndim, final_size = _block_check_depths_match(arrays)
960 list_ndim = len(bottom_index)
961 if bottom_index and bottom_index[-1] is None:
962 raise ValueError(
963 f'List at {_block_format_index(bottom_index)} cannot be empty'
964 )
965 result_ndim = max(arr_ndim, list_ndim)
966 return arrays, list_ndim, result_ndim, final_size
967
968
969def _block_slicing(arrays, list_ndim, result_ndim):
970 shape, slices, arrays = _block_info_recursion(
971 arrays, list_ndim, result_ndim)
972 dtype = _nx.result_type(*[arr.dtype for arr in arrays])
973
974 # Test preferring F only in the case that all input arrays are F
975 F_order = all(arr.flags['F_CONTIGUOUS'] for arr in arrays)
976 C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays)
977 order = 'F' if F_order and not C_order else 'C'
978 result = _nx.empty(shape=shape, dtype=dtype, order=order)
979 # Note: In a c implementation, the function
980 # PyArray_CreateMultiSortedStridePerm could be used for more advanced
981 # guessing of the desired order.
982
983 for the_slice, arr in zip(slices, arrays):
984 result[(Ellipsis,) + the_slice] = arr
985 return result
986
987
988def _block_concatenate(arrays, list_ndim, result_ndim):
989 result = _block(arrays, list_ndim, result_ndim)
990 if list_ndim == 0:
991 # Catch an edge case where _block returns a view because
992 # `arrays` is a single numpy array and not a list of numpy arrays.
993 # This might copy scalars or lists twice, but this isn't a likely
994 # usecase for those interested in performance
995 result = result.copy()
996 return result
997