Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/numpy/lib/_shape_base_impl.py: 24%

1import functools 

2import warnings 

3 

4import numpy as np 

5import numpy._core.numeric as _nx 

6from numpy._core import atleast_3d, overrides, vstack 

7from numpy._core._multiarray_umath import _array_converter 

8from numpy._core.fromnumeric import reshape, transpose 

9from numpy._core.multiarray import normalize_axis_index 

10from numpy._core.numeric import ( 

11 array, 

12 asanyarray, 

13 asarray, 

14 normalize_axis_tuple, 

15 zeros, 

16 zeros_like, 

17) 

18from numpy._core.overrides import set_module 

19from numpy._core.shape_base import _arrays_for_stack_dispatcher 

20from numpy.lib._index_tricks_impl import ndindex 

21from numpy.matrixlib.defmatrix import matrix # this raises all the right alarm bells 

22 

23__all__ = [ 

24 'column_stack', 'row_stack', 'dstack', 'array_split', 'split', 

25 'hsplit', 'vsplit', 'dsplit', 'apply_over_axes', 'expand_dims', 

26 'apply_along_axis', 'kron', 'tile', 'take_along_axis', 

27 'put_along_axis' 

28 ] 

29 

30 

31array_function_dispatch = functools.partial( 

32 overrides.array_function_dispatch, module='numpy') 

33 

34 

35def _make_along_axis_idx(arr_shape, indices, axis): 

36 # compute dimensions to iterate over 

37 if not _nx.issubdtype(indices.dtype, _nx.integer): 

38 raise IndexError('`indices` must be an integer array') 

39 if len(arr_shape) != indices.ndim: 

40 raise ValueError( 

41 "`indices` and `arr` must have the same number of dimensions") 

42 shape_ones = (1,) * indices.ndim 

43 dest_dims = list(range(axis)) + [None] + list(range(axis + 1, indices.ndim)) 

44 

45 # build a fancy index, consisting of orthogonal aranges, with the 

46 # requested index inserted at the right location 

47 fancy_index = [] 

48 for dim, n in zip(dest_dims, arr_shape): 

49 if dim is None: 

50 fancy_index.append(indices) 

51 else: 

52 ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1:] 

53 fancy_index.append(_nx.arange(n).reshape(ind_shape)) 

54 

55 return tuple(fancy_index) 

56 

57 

58def _take_along_axis_dispatcher(arr, indices, axis=None): 

59 return (arr, indices) 

60 

61 

62@array_function_dispatch(_take_along_axis_dispatcher) 

63def take_along_axis(arr, indices, axis=-1): 

64 """ 

65 Take values from the input array by matching 1d index and data slices. 

66 

67 This iterates over matching 1d slices oriented along the specified axis in 

68 the index and data arrays, and uses the former to look up values in the 

69 latter. These slices can be different lengths. 

70 

71 Functions returning an index along an axis, like `argsort` and 

72 `argpartition`, produce suitable indices for this function. 

73 

74 Parameters 

75 ---------- 

76 arr : ndarray (Ni..., M, Nk...) 

77 Source array 

78 indices : ndarray (Ni..., J, Nk...) 

79 Indices to take along each 1d slice of ``arr``. This must match the 

80 dimension of ``arr``, but dimensions Ni and Nj only need to broadcast 

81 against ``arr``. 

82 axis : int or None, optional 

83 The axis to take 1d slices along. If axis is None, the input array is 

84 treated as if it had first been flattened to 1d, for consistency with 

85 `sort` and `argsort`. 

86 

87 .. versionchanged:: 2.3 

88 The default value is now ``-1``. 

89 

90 Returns 

91 ------- 

92 out: ndarray (Ni..., J, Nk...) 

93 The indexed result. 

94 

95 Notes 

96 ----- 

97 This is equivalent to (but faster than) the following use of `ndindex` and 

98 `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices:: 

99 

100 Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:] 

101 J = indices.shape[axis] # Need not equal M 

102 out = np.empty(Ni + (J,) + Nk) 

103 

104 for ii in ndindex(Ni): 

105 for kk in ndindex(Nk): 

106 a_1d = a [ii + s_[:,] + kk] 

107 indices_1d = indices[ii + s_[:,] + kk] 

108 out_1d = out [ii + s_[:,] + kk] 

109 for j in range(J): 

110 out_1d[j] = a_1d[indices_1d[j]] 

111 

112 Equivalently, eliminating the inner loop, the last two lines would be:: 

113 

114 out_1d[:] = a_1d[indices_1d] 

115 

116 See Also 

117 -------- 

118 take : Take along an axis, using the same indices for every 1d slice 

119 put_along_axis : 

120 Put values into the destination array by matching 1d index and data slices 

121 

122 Examples 

123 -------- 

124 >>> import numpy as np 

125 

126 For this sample array 

127 

128 >>> a = np.array([[10, 30, 20], [60, 40, 50]]) 

129 

130 We can sort either by using sort directly, or argsort and this function 

131 

132 >>> np.sort(a, axis=1) 

133 array([[10, 20, 30], 

134 [40, 50, 60]]) 

135 >>> ai = np.argsort(a, axis=1) 

136 >>> ai 

137 array([[0, 2, 1], 

138 [1, 2, 0]]) 

139 >>> np.take_along_axis(a, ai, axis=1) 

140 array([[10, 20, 30], 

141 [40, 50, 60]]) 

142 

143 The same works for max and min, if you maintain the trivial dimension 

144 with ``keepdims``: 

145 

146 >>> np.max(a, axis=1, keepdims=True) 

147 array([[30], 

148 [60]]) 

149 >>> ai = np.argmax(a, axis=1, keepdims=True) 

150 >>> ai 

151 array([[1], 

152 [0]]) 

153 >>> np.take_along_axis(a, ai, axis=1) 

154 array([[30], 

155 [60]]) 

156 

157 If we want to get the max and min at the same time, we can stack the 

158 indices first 

159 

160 >>> ai_min = np.argmin(a, axis=1, keepdims=True) 

161 >>> ai_max = np.argmax(a, axis=1, keepdims=True) 

162 >>> ai = np.concatenate([ai_min, ai_max], axis=1) 

163 >>> ai 

164 array([[0, 1], 

165 [1, 0]]) 

166 >>> np.take_along_axis(a, ai, axis=1) 

167 array([[10, 30], 

168 [40, 60]]) 

169 """ 

170 # normalize inputs 

171 if axis is None: 

172 if indices.ndim != 1: 

173 raise ValueError( 

174 'when axis=None, `indices` must have a single dimension.') 

175 arr = np.array(arr.flat) 

176 axis = 0 

177 else: 

178 axis = normalize_axis_index(axis, arr.ndim) 

179 

180 # use the fancy index 

181 return arr[_make_along_axis_idx(arr.shape, indices, axis)] 

182 

183 

184def _put_along_axis_dispatcher(arr, indices, values, axis): 

185 return (arr, indices, values) 

186 

187 

188@array_function_dispatch(_put_along_axis_dispatcher) 

189def put_along_axis(arr, indices, values, axis): 

190 """ 

191 Put values into the destination array by matching 1d index and data slices. 

192 

193 This iterates over matching 1d slices oriented along the specified axis in 

194 the index and data arrays, and uses the former to place values into the 

195 latter. These slices can be different lengths. 

196 

197 Functions returning an index along an axis, like `argsort` and 

198 `argpartition`, produce suitable indices for this function. 

199 

200 Parameters 

201 ---------- 

202 arr : ndarray (Ni..., M, Nk...) 

203 Destination array. 

204 indices : ndarray (Ni..., J, Nk...) 

205 Indices to change along each 1d slice of `arr`. This must match the 

206 dimension of arr, but dimensions in Ni and Nj may be 1 to broadcast 

207 against `arr`. 

208 values : array_like (Ni..., J, Nk...) 

209 values to insert at those indices. Its shape and dimension are 

210 broadcast to match that of `indices`. 

211 axis : int 

212 The axis to take 1d slices along. If axis is None, the destination 

213 array is treated as if a flattened 1d view had been created of it. 

214 

215 Notes 

216 ----- 

217 This is equivalent to (but faster than) the following use of `ndindex` and 

218 `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices:: 

219 

220 Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:] 

221 J = indices.shape[axis] # Need not equal M 

222 

223 for ii in ndindex(Ni): 

224 for kk in ndindex(Nk): 

225 a_1d = a [ii + s_[:,] + kk] 

226 indices_1d = indices[ii + s_[:,] + kk] 

227 values_1d = values [ii + s_[:,] + kk] 

228 for j in range(J): 

229 a_1d[indices_1d[j]] = values_1d[j] 

230 

231 Equivalently, eliminating the inner loop, the last two lines would be:: 

232 

233 a_1d[indices_1d] = values_1d 

234 

235 See Also 

236 -------- 

237 take_along_axis : 

238 Take values from the input array by matching 1d index and data slices 

239 

240 Examples 

241 -------- 

242 >>> import numpy as np 

243 

244 For this sample array 

245 

246 >>> a = np.array([[10, 30, 20], [60, 40, 50]]) 

247 

248 We can replace the maximum values with: 

249 

250 >>> ai = np.argmax(a, axis=1, keepdims=True) 

251 >>> ai 

252 array([[1], 

253 [0]]) 

254 >>> np.put_along_axis(a, ai, 99, axis=1) 

255 >>> a 

256 array([[10, 99, 20], 

257 [99, 40, 50]]) 

258 

259 """ 

260 # normalize inputs 

261 if axis is None: 

262 if indices.ndim != 1: 

263 raise ValueError( 

264 'when axis=None, `indices` must have a single dimension.') 

265 arr = np.array(arr.flat) 

266 axis = 0 

267 else: 

268 axis = normalize_axis_index(axis, arr.ndim) 

269 

270 # use the fancy index 

271 arr[_make_along_axis_idx(arr.shape, indices, axis)] = values 

272 

273 

274def _apply_along_axis_dispatcher(func1d, axis, arr, *args, **kwargs): 

275 return (arr,) 

276 

277 

278@array_function_dispatch(_apply_along_axis_dispatcher) 

279def apply_along_axis(func1d, axis, arr, *args, **kwargs): 

280 """ 

281 Apply a function to 1-D slices along the given axis. 

282 

283 Execute `func1d(a, *args, **kwargs)` where `func1d` operates on 1-D arrays 

284 and `a` is a 1-D slice of `arr` along `axis`. 

285 

286 This is equivalent to (but faster than) the following use of `ndindex` and 

287 `s_`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of indices:: 

288 

289 Ni, Nk = a.shape[:axis], a.shape[axis+1:] 

290 for ii in ndindex(Ni): 

291 for kk in ndindex(Nk): 

292 f = func1d(arr[ii + s_[:,] + kk]) 

293 Nj = f.shape 

294 for jj in ndindex(Nj): 

295 out[ii + jj + kk] = f[jj] 

296 

297 Equivalently, eliminating the inner loop, this can be expressed as:: 

298 

299 Ni, Nk = a.shape[:axis], a.shape[axis+1:] 

300 for ii in ndindex(Ni): 

301 for kk in ndindex(Nk): 

302 out[ii + s_[...,] + kk] = func1d(arr[ii + s_[:,] + kk]) 

303 

304 Parameters 

305 ---------- 

306 func1d : function (M,) -> (Nj...) 

307 This function should accept 1-D arrays. It is applied to 1-D 

308 slices of `arr` along the specified axis. 

309 axis : integer 

310 Axis along which `arr` is sliced. 

311 arr : ndarray (Ni..., M, Nk...) 

312 Input array. 

313 args : any 

314 Additional arguments to `func1d`. 

315 kwargs : any 

316 Additional named arguments to `func1d`. 

317 

318 Returns 

319 ------- 

320 out : ndarray (Ni..., Nj..., Nk...) 

321 The output array. The shape of `out` is identical to the shape of 

322 `arr`, except along the `axis` dimension. This axis is removed, and 

323 replaced with new dimensions equal to the shape of the return value 

324 of `func1d`. So if `func1d` returns a scalar `out` will have one 

325 fewer dimensions than `arr`. 

326 

327 See Also 

328 -------- 

329 apply_over_axes : Apply a function repeatedly over multiple axes. 

330 

331 Examples 

332 -------- 

333 >>> import numpy as np 

334 >>> def my_func(a): 

335 ... \"\"\"Average first and last element of a 1-D array\"\"\" 

336 ... return (a[0] + a[-1]) * 0.5 

337 >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]]) 

338 >>> np.apply_along_axis(my_func, 0, b) 

339 array([4., 5., 6.]) 

340 >>> np.apply_along_axis(my_func, 1, b) 

341 array([2., 5., 8.]) 

342 

343 For a function that returns a 1D array, the number of dimensions in 

344 `outarr` is the same as `arr`. 

345 

346 >>> b = np.array([[8,1,7], [4,3,9], [5,2,6]]) 

347 >>> np.apply_along_axis(sorted, 1, b) 

348 array([[1, 7, 8], 

349 [3, 4, 9], 

350 [2, 5, 6]]) 

351 

352 For a function that returns a higher dimensional array, those dimensions 

353 are inserted in place of the `axis` dimension. 

354 

355 >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]]) 

356 >>> np.apply_along_axis(np.diag, -1, b) 

357 array([[[1, 0, 0], 

358 [0, 2, 0], 

359 [0, 0, 3]], 

360 [[4, 0, 0], 

361 [0, 5, 0], 

362 [0, 0, 6]], 

363 [[7, 0, 0], 

364 [0, 8, 0], 

365 [0, 0, 9]]]) 

366 """ 

367 # handle negative axes 

368 conv = _array_converter(arr) 

369 arr = conv[0] 

370 

371 nd = arr.ndim 

372 axis = normalize_axis_index(axis, nd) 

373 

374 # arr, with the iteration axis at the end 

375 in_dims = list(range(nd)) 

376 inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis + 1:] + [axis]) 

377 

378 # compute indices for the iteration axes, and append a trailing ellipsis to 

379 # prevent 0d arrays decaying to scalars, which fixes gh-8642 

380 inds = ndindex(inarr_view.shape[:-1]) 

381 inds = (ind + (Ellipsis,) for ind in inds) 

382 

383 # invoke the function on the first item 

384 try: 

385 ind0 = next(inds) 

386 except StopIteration: 

387 raise ValueError( 

388 'Cannot apply_along_axis when any iteration dimensions are 0' 

389 ) from None 

390 res = asanyarray(func1d(inarr_view[ind0], *args, **kwargs)) 

391 

392 # build a buffer for storing evaluations of func1d. 

393 # remove the requested axis, and add the new ones on the end. 

394 # laid out so that each write is contiguous. 

395 # for a tuple index inds, buff[inds] = func1d(inarr_view[inds]) 

396 if not isinstance(res, matrix): 

397 buff = zeros_like(res, shape=inarr_view.shape[:-1] + res.shape) 

398 else: 

399 # Matrices are nasty with reshaping, so do not preserve them here. 

400 buff = zeros(inarr_view.shape[:-1] + res.shape, dtype=res.dtype) 

401 

402 # permutation of axes such that out = buff.transpose(buff_permute) 

403 buff_dims = list(range(buff.ndim)) 

404 buff_permute = ( 

405 buff_dims[0 : axis] + 

406 buff_dims[buff.ndim - res.ndim : buff.ndim] + 

407 buff_dims[axis : buff.ndim - res.ndim] 

408 ) 

409 

410 # save the first result, then compute and save all remaining results 

411 buff[ind0] = res 

412 for ind in inds: 

413 buff[ind] = asanyarray(func1d(inarr_view[ind], *args, **kwargs)) 

414 

415 res = transpose(buff, buff_permute) 

416 return conv.wrap(res) 

417 

418 

419def _apply_over_axes_dispatcher(func, a, axes): 

420 return (a,) 

421 

422 

423@array_function_dispatch(_apply_over_axes_dispatcher) 

424def apply_over_axes(func, a, axes): 

425 """ 

426 Apply a function repeatedly over multiple axes. 

427 

428 `func` is called as `res = func(a, axis)`, where `axis` is the first 

429 element of `axes`. The result `res` of the function call must have 

430 either the same dimensions as `a` or one less dimension. If `res` 

431 has one less dimension than `a`, a dimension is inserted before 

432 `axis`. The call to `func` is then repeated for each axis in `axes`, 

433 with `res` as the first argument. 

434 

435 Parameters 

436 ---------- 

437 func : function 

438 This function must take two arguments, `func(a, axis)`. 

439 a : array_like 

440 Input array. 

441 axes : array_like 

442 Axes over which `func` is applied; the elements must be integers. 

443 

444 Returns 

445 ------- 

446 apply_over_axis : ndarray 

447 The output array. The number of dimensions is the same as `a`, 

448 but the shape can be different. This depends on whether `func` 

449 changes the shape of its output with respect to its input. 

450 

451 See Also 

452 -------- 

453 apply_along_axis : 

454 Apply a function to 1-D slices of an array along the given axis. 

455 

456 Notes 

457 ----- 

458 This function is equivalent to tuple axis arguments to reorderable ufuncs 

459 with keepdims=True. Tuple axis arguments to ufuncs have been available since 

460 version 1.7.0. 

461 

462 Examples 

463 -------- 

464 >>> import numpy as np 

465 >>> a = np.arange(24).reshape(2,3,4) 

466 >>> a 

467 array([[[ 0, 1, 2, 3], 

468 [ 4, 5, 6, 7], 

469 [ 8, 9, 10, 11]], 

470 [[12, 13, 14, 15], 

471 [16, 17, 18, 19], 

472 [20, 21, 22, 23]]]) 

473 

474 Sum over axes 0 and 2. The result has same number of dimensions 

475 as the original array: 

476 

477 >>> np.apply_over_axes(np.sum, a, [0,2]) 

478 array([[[ 60], 

479 [ 92], 

480 [124]]]) 

481 

482 Tuple axis arguments to ufuncs are equivalent: 

483 

484 >>> np.sum(a, axis=(0,2), keepdims=True) 

485 array([[[ 60], 

486 [ 92], 

487 [124]]]) 

488 

489 """ 

490 val = asarray(a) 

491 N = a.ndim 

492 if array(axes).ndim == 0: 

493 axes = (axes,) 

494 for axis in axes: 

495 if axis < 0: 

496 axis = N + axis 

497 args = (val, axis) 

498 res = func(*args) 

499 if res.ndim == val.ndim: 

500 val = res 

501 else: 

502 res = expand_dims(res, axis) 

503 if res.ndim == val.ndim: 

504 val = res 

505 else: 

506 raise ValueError("function is not returning " 

507 "an array of the correct shape") 

508 return val 

509 

510 

511def _expand_dims_dispatcher(a, axis): 

512 return (a,) 

513 

514 

515@array_function_dispatch(_expand_dims_dispatcher) 

516def expand_dims(a, axis): 

517 """ 

518 Expand the shape of an array. 

519 

520 Insert a new axis that will appear at the `axis` position in the expanded 

521 array shape. 

522 

523 Parameters 

524 ---------- 

525 a : array_like 

526 Input array. 

527 axis : int or tuple of ints 

528 Position in the expanded axes where the new axis (or axes) is placed. 

529 

530 .. deprecated:: 1.13.0 

531 Passing an axis where ``axis > a.ndim`` will be treated as 

532 ``axis == a.ndim``, and passing ``axis < -a.ndim - 1`` will 

533 be treated as ``axis == 0``. This behavior is deprecated. 

534 

535 Returns 

536 ------- 

537 result : ndarray 

538 View of `a` with the number of dimensions increased. 

539 

540 See Also 

541 -------- 

542 squeeze : The inverse operation, removing singleton dimensions 

543 reshape : Insert, remove, and combine dimensions, and resize existing ones 

544 atleast_1d, atleast_2d, atleast_3d 

545 

546 Examples 

547 -------- 

548 >>> import numpy as np 

549 >>> x = np.array([1, 2]) 

550 >>> x.shape 

551 (2,) 

552 

553 The following is equivalent to ``x[np.newaxis, :]`` or ``x[np.newaxis]``: 

554 

555 >>> y = np.expand_dims(x, axis=0) 

556 >>> y 

557 array([[1, 2]]) 

558 >>> y.shape 

559 (1, 2) 

560 

561 The following is equivalent to ``x[:, np.newaxis]``: 

562 

563 >>> y = np.expand_dims(x, axis=1) 

564 >>> y 

565 array([[1], 

566 [2]]) 

567 >>> y.shape 

568 (2, 1) 

569 

570 ``axis`` may also be a tuple: 

571 

572 >>> y = np.expand_dims(x, axis=(0, 1)) 

573 >>> y 

574 array([[[1, 2]]]) 

575 

576 >>> y = np.expand_dims(x, axis=(2, 0)) 

577 >>> y 

578 array([[[1], 

579 [2]]]) 

580 

581 Note that some examples may use ``None`` instead of ``np.newaxis``. These 

582 are the same objects: 

583 

584 >>> np.newaxis is None 

585 True 

586 

587 """ 

588 if isinstance(a, matrix): 

589 a = asarray(a) 

590 else: 

591 a = asanyarray(a) 

592 

593 if not isinstance(axis, (tuple, list)): 

594 axis = (axis,) 

595 

596 out_ndim = len(axis) + a.ndim 

597 axis = normalize_axis_tuple(axis, out_ndim) 

598 

599 shape_it = iter(a.shape) 

600 shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)] 

601 

602 return a.reshape(shape) 

603 

604 

605# NOTE: Remove once deprecation period passes 

606@set_module("numpy") 

607def row_stack(tup, *, dtype=None, casting="same_kind"): 

608 # Deprecated in NumPy 2.0, 2023-08-18 

609 warnings.warn( 

610 "`row_stack` alias is deprecated. " 

611 "Use `np.vstack` directly.", 

612 DeprecationWarning, 

613 stacklevel=2 

614 ) 

615 return vstack(tup, dtype=dtype, casting=casting) 

616 

617 

618row_stack.__doc__ = vstack.__doc__ 

619 

620 

621def _column_stack_dispatcher(tup): 

622 return _arrays_for_stack_dispatcher(tup) 

623 

624 

625@array_function_dispatch(_column_stack_dispatcher) 

626def column_stack(tup): 

627 """ 

628 Stack 1-D arrays as columns into a 2-D array. 

629 

630 Take a sequence of 1-D arrays and stack them as columns 

631 to make a single 2-D array. 2-D arrays are stacked as-is, 

632 just like with `hstack`. 1-D arrays are turned into 2-D columns 

633 first. 

634 

635 Parameters 

636 ---------- 

637 tup : sequence of 1-D or 2-D arrays. 

638 Arrays to stack. All of them must have the same first dimension. 

639 

640 Returns 

641 ------- 

642 stacked : 2-D array 

643 The array formed by stacking the given arrays. 

644 

645 See Also 

646 -------- 

647 stack, hstack, vstack, concatenate 

648 

649 Examples 

650 -------- 

651 >>> import numpy as np 

652 >>> a = np.array((1,2,3)) 

653 >>> b = np.array((4,5,6)) 

654 >>> np.column_stack((a,b)) 

655 array([[1, 4], 

656 [2, 5], 

657 [3, 6]]) 

658 

659 """ 

660 arrays = [] 

661 for v in tup: 

662 arr = asanyarray(v) 

663 if arr.ndim < 2: 

664 arr = array(arr, copy=None, subok=True, ndmin=2).T 

665 arrays.append(arr) 

666 return _nx.concatenate(arrays, 1) 

667 

668 

669def _dstack_dispatcher(tup): 

670 return _arrays_for_stack_dispatcher(tup) 

671 

672 

673@array_function_dispatch(_dstack_dispatcher) 

674def dstack(tup): 

675 """ 

676 Stack arrays in sequence depth wise (along third axis). 

677 

678 This is equivalent to concatenation along the third axis after 2-D arrays 

679 of shape `(M,N)` have been reshaped to `(M,N,1)` and 1-D arrays of shape 

680 `(N,)` have been reshaped to `(1,N,1)`. Rebuilds arrays divided by 

681 `dsplit`. 

682 

683 This function makes most sense for arrays with up to 3 dimensions. For 

684 instance, for pixel-data with a height (first axis), width (second axis), 

685 and r/g/b channels (third axis). The functions `concatenate`, `stack` and 

686 `block` provide more general stacking and concatenation operations. 

687 

688 Parameters 

689 ---------- 

690 tup : sequence of arrays 

691 The arrays must have the same shape along all but the third axis. 

692 1-D or 2-D arrays must have the same shape. 

693 

694 Returns 

695 ------- 

696 stacked : ndarray 

697 The array formed by stacking the given arrays, will be at least 3-D. 

698 

699 See Also 

700 -------- 

701 concatenate : Join a sequence of arrays along an existing axis. 

702 stack : Join a sequence of arrays along a new axis. 

703 block : Assemble an nd-array from nested lists of blocks. 

704 vstack : Stack arrays in sequence vertically (row wise). 

705 hstack : Stack arrays in sequence horizontally (column wise). 

706 column_stack : Stack 1-D arrays as columns into a 2-D array. 

707 dsplit : Split array along third axis. 

708 

709 Examples 

710 -------- 

711 >>> import numpy as np 

712 >>> a = np.array((1,2,3)) 

713 >>> b = np.array((4,5,6)) 

714 >>> np.dstack((a,b)) 

715 array([[[1, 4], 

716 [2, 5], 

717 [3, 6]]]) 

718 

719 >>> a = np.array([[1],[2],[3]]) 

720 >>> b = np.array([[4],[5],[6]]) 

721 >>> np.dstack((a,b)) 

722 array([[[1, 4]], 

723 [[2, 5]], 

724 [[3, 6]]]) 

725 

726 """ 

727 arrs = atleast_3d(*tup) 

728 if not isinstance(arrs, tuple): 

729 arrs = (arrs,) 

730 return _nx.concatenate(arrs, 2) 

731 

732 

733def _array_split_dispatcher(ary, indices_or_sections, axis=None): 

734 return (ary, indices_or_sections) 

735 

736 

737@array_function_dispatch(_array_split_dispatcher) 

738def array_split(ary, indices_or_sections, axis=0): 

739 """ 

740 Split an array into multiple sub-arrays. 

741 

742 Please refer to the ``split`` documentation. The only difference 

743 between these functions is that ``array_split`` allows 

744 `indices_or_sections` to be an integer that does *not* equally 

745 divide the axis. For an array of length l that should be split 

746 into n sections, it returns l % n sub-arrays of size l//n + 1 

747 and the rest of size l//n. 

748 

749 See Also 

750 -------- 

751 split : Split array into multiple sub-arrays of equal size. 

752 

753 Examples 

754 -------- 

755 >>> import numpy as np 

756 >>> x = np.arange(8.0) 

757 >>> np.array_split(x, 3) 

758 [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])] 

759 

760 >>> x = np.arange(9) 

761 >>> np.array_split(x, 4) 

762 [array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8])] 

763 

764 """ 

765 try: 

766 Ntotal = ary.shape[axis] 

767 except AttributeError: 

768 Ntotal = len(ary) 

769 try: 

770 # handle array case. 

771 Nsections = len(indices_or_sections) + 1 

772 div_points = [0] + list(indices_or_sections) + [Ntotal] 

773 except TypeError: 

774 # indices_or_sections is a scalar, not an array. 

775 Nsections = int(indices_or_sections) 

776 if Nsections <= 0: 

777 raise ValueError('number sections must be larger than 0.') from None 

778 Neach_section, extras = divmod(Ntotal, Nsections) 

779 section_sizes = ([0] + 

780 extras * [Neach_section + 1] + 

781 (Nsections - extras) * [Neach_section]) 

782 div_points = _nx.array(section_sizes, dtype=_nx.intp).cumsum() 

783 

784 sub_arys = [] 

785 sary = _nx.swapaxes(ary, axis, 0) 

786 for i in range(Nsections): 

787 st = div_points[i] 

788 end = div_points[i + 1] 

789 sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0)) 

790 

791 return sub_arys 

792 

793 

794def _split_dispatcher(ary, indices_or_sections, axis=None): 

795 return (ary, indices_or_sections) 

796 

797 

798@array_function_dispatch(_split_dispatcher) 

799def split(ary, indices_or_sections, axis=0): 

800 """ 

801 Split an array into multiple sub-arrays as views into `ary`. 

802 

803 Parameters 

804 ---------- 

805 ary : ndarray 

806 Array to be divided into sub-arrays. 

807 indices_or_sections : int or 1-D array 

808 If `indices_or_sections` is an integer, N, the array will be divided 

809 into N equal arrays along `axis`. If such a split is not possible, 

810 an error is raised. 

811 

812 If `indices_or_sections` is a 1-D array of sorted integers, the entries 

813 indicate where along `axis` the array is split. For example, 

814 ``[2, 3]`` would, for ``axis=0``, result in 

815 

816 - ary[:2] 

817 - ary[2:3] 

818 - ary[3:] 

819 

820 If an index exceeds the dimension of the array along `axis`, 

821 an empty sub-array is returned correspondingly. 

822 axis : int, optional 

823 The axis along which to split, default is 0. 

824 

825 Returns 

826 ------- 

827 sub-arrays : list of ndarrays 

828 A list of sub-arrays as views into `ary`. 

829 

830 Raises 

831 ------ 

832 ValueError 

833 If `indices_or_sections` is given as an integer, but 

834 a split does not result in equal division. 

835 

836 See Also 

837 -------- 

838 array_split : Split an array into multiple sub-arrays of equal or 

839 near-equal size. Does not raise an exception if 

840 an equal division cannot be made. 

841 hsplit : Split array into multiple sub-arrays horizontally (column-wise). 

842 vsplit : Split array into multiple sub-arrays vertically (row wise). 

843 dsplit : Split array into multiple sub-arrays along the 3rd axis (depth). 

844 concatenate : Join a sequence of arrays along an existing axis. 

845 stack : Join a sequence of arrays along a new axis. 

846 hstack : Stack arrays in sequence horizontally (column wise). 

847 vstack : Stack arrays in sequence vertically (row wise). 

848 dstack : Stack arrays in sequence depth wise (along third dimension). 

849 

850 Examples 

851 -------- 

852 >>> import numpy as np 

853 >>> x = np.arange(9.0) 

854 >>> np.split(x, 3) 

855 [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])] 

856 

857 >>> x = np.arange(8.0) 

858 >>> np.split(x, [3, 5, 6, 10]) 

859 [array([0., 1., 2.]), 

860 array([3., 4.]), 

861 array([5.]), 

862 array([6., 7.]), 

863 array([], dtype=float64)] 

864 

865 """ 

866 try: 

867 len(indices_or_sections) 

868 except TypeError: 

869 sections = indices_or_sections 

870 N = ary.shape[axis] 

871 if N % sections: 

872 raise ValueError( 

873 'array split does not result in an equal division') from None 

874 return array_split(ary, indices_or_sections, axis) 

875 

876 

877def _hvdsplit_dispatcher(ary, indices_or_sections): 

878 return (ary, indices_or_sections) 

879 

880 

881@array_function_dispatch(_hvdsplit_dispatcher) 

882def hsplit(ary, indices_or_sections): 

883 """ 

884 Split an array into multiple sub-arrays horizontally (column-wise). 

885 

886 Please refer to the `split` documentation. `hsplit` is equivalent 

887 to `split` with ``axis=1``, the array is always split along the second 

888 axis except for 1-D arrays, where it is split at ``axis=0``. 

889 

890 See Also 

891 -------- 

892 split : Split an array into multiple sub-arrays of equal size. 

893 

894 Examples 

895 -------- 

896 >>> import numpy as np 

897 >>> x = np.arange(16.0).reshape(4, 4) 

898 >>> x 

899 array([[ 0., 1., 2., 3.], 

900 [ 4., 5., 6., 7.], 

901 [ 8., 9., 10., 11.], 

902 [12., 13., 14., 15.]]) 

903 >>> np.hsplit(x, 2) 

904 [array([[ 0., 1.], 

905 [ 4., 5.], 

906 [ 8., 9.], 

907 [12., 13.]]), 

908 array([[ 2., 3.], 

909 [ 6., 7.], 

910 [10., 11.], 

911 [14., 15.]])] 

912 >>> np.hsplit(x, np.array([3, 6])) 

913 [array([[ 0., 1., 2.], 

914 [ 4., 5., 6.], 

915 [ 8., 9., 10.], 

916 [12., 13., 14.]]), 

917 array([[ 3.], 

918 [ 7.], 

919 [11.], 

920 [15.]]), 

921 array([], shape=(4, 0), dtype=float64)] 

922 

923 With a higher dimensional array the split is still along the second axis. 

924 

925 >>> x = np.arange(8.0).reshape(2, 2, 2) 

926 >>> x 

927 array([[[0., 1.], 

928 [2., 3.]], 

929 [[4., 5.], 

930 [6., 7.]]]) 

931 >>> np.hsplit(x, 2) 

932 [array([[[0., 1.]], 

933 [[4., 5.]]]), 

934 array([[[2., 3.]], 

935 [[6., 7.]]])] 

936 

937 With a 1-D array, the split is along axis 0. 

938 

939 >>> x = np.array([0, 1, 2, 3, 4, 5]) 

940 >>> np.hsplit(x, 2) 

941 [array([0, 1, 2]), array([3, 4, 5])] 

942 

943 """ 

944 if _nx.ndim(ary) == 0: 

945 raise ValueError('hsplit only works on arrays of 1 or more dimensions') 

946 if ary.ndim > 1: 

947 return split(ary, indices_or_sections, 1) 

948 else: 

949 return split(ary, indices_or_sections, 0) 

950 

951 

952@array_function_dispatch(_hvdsplit_dispatcher) 

953def vsplit(ary, indices_or_sections): 

954 """ 

955 Split an array into multiple sub-arrays vertically (row-wise). 

956 

957 Please refer to the ``split`` documentation. ``vsplit`` is equivalent 

958 to ``split`` with `axis=0` (default), the array is always split along the 

959 first axis regardless of the array dimension. 

960 

961 See Also 

962 -------- 

963 split : Split an array into multiple sub-arrays of equal size. 

964 

965 Examples 

966 -------- 

967 >>> import numpy as np 

968 >>> x = np.arange(16.0).reshape(4, 4) 

969 >>> x 

970 array([[ 0., 1., 2., 3.], 

971 [ 4., 5., 6., 7.], 

972 [ 8., 9., 10., 11.], 

973 [12., 13., 14., 15.]]) 

974 >>> np.vsplit(x, 2) 

975 [array([[0., 1., 2., 3.], 

976 [4., 5., 6., 7.]]), 

977 array([[ 8., 9., 10., 11.], 

978 [12., 13., 14., 15.]])] 

979 >>> np.vsplit(x, np.array([3, 6])) 

980 [array([[ 0., 1., 2., 3.], 

981 [ 4., 5., 6., 7.], 

982 [ 8., 9., 10., 11.]]), 

983 array([[12., 13., 14., 15.]]), 

984 array([], shape=(0, 4), dtype=float64)] 

985 

986 With a higher dimensional array the split is still along the first axis. 

987 

988 >>> x = np.arange(8.0).reshape(2, 2, 2) 

989 >>> x 

990 array([[[0., 1.], 

991 [2., 3.]], 

992 [[4., 5.], 

993 [6., 7.]]]) 

994 >>> np.vsplit(x, 2) 

995 [array([[[0., 1.], 

996 [2., 3.]]]), 

997 array([[[4., 5.], 

998 [6., 7.]]])] 

999 

1000 """ 

1001 if _nx.ndim(ary) < 2: 

1002 raise ValueError('vsplit only works on arrays of 2 or more dimensions') 

1003 return split(ary, indices_or_sections, 0) 

1004 

1005 

1006@array_function_dispatch(_hvdsplit_dispatcher) 

1007def dsplit(ary, indices_or_sections): 

1008 """ 

1009 Split array into multiple sub-arrays along the 3rd axis (depth). 

1010 

1011 Please refer to the `split` documentation. `dsplit` is equivalent 

1012 to `split` with ``axis=2``, the array is always split along the third 

1013 axis provided the array dimension is greater than or equal to 3. 

1014 

1015 See Also 

1016 -------- 

1017 split : Split an array into multiple sub-arrays of equal size. 

1018 

1019 Examples 

1020 -------- 

1021 >>> import numpy as np 

1022 >>> x = np.arange(16.0).reshape(2, 2, 4) 

1023 >>> x 

1024 array([[[ 0., 1., 2., 3.], 

1025 [ 4., 5., 6., 7.]], 

1026 [[ 8., 9., 10., 11.], 

1027 [12., 13., 14., 15.]]]) 

1028 >>> np.dsplit(x, 2) 

1029 [array([[[ 0., 1.], 

1030 [ 4., 5.]], 

1031 [[ 8., 9.], 

1032 [12., 13.]]]), array([[[ 2., 3.], 

1033 [ 6., 7.]], 

1034 [[10., 11.], 

1035 [14., 15.]]])] 

1036 >>> np.dsplit(x, np.array([3, 6])) 

1037 [array([[[ 0., 1., 2.], 

1038 [ 4., 5., 6.]], 

1039 [[ 8., 9., 10.], 

1040 [12., 13., 14.]]]), 

1041 array([[[ 3.], 

1042 [ 7.]], 

1043 [[11.], 

1044 [15.]]]), 

1045 array([], shape=(2, 2, 0), dtype=float64)] 

1046 """ 

1047 if _nx.ndim(ary) < 3: 

1048 raise ValueError('dsplit only works on arrays of 3 or more dimensions') 

1049 return split(ary, indices_or_sections, 2) 

1050 

1051 

1052def get_array_wrap(*args): 

1053 """Find the wrapper for the array with the highest priority. 

1054 

1055 In case of ties, leftmost wins. If no wrapper is found, return None. 

1056 

1057 .. deprecated:: 2.0 

1058 """ 

1059 

1060 # Deprecated in NumPy 2.0, 2023-07-11 

1061 warnings.warn( 

1062 "`get_array_wrap` is deprecated. " 

1063 "(deprecated in NumPy 2.0)", 

1064 DeprecationWarning, 

1065 stacklevel=2 

1066 ) 

1067 

1068 wrappers = sorted((getattr(x, '__array_priority__', 0), -i, 

1069 x.__array_wrap__) for i, x in enumerate(args) 

1070 if hasattr(x, '__array_wrap__')) 

1071 if wrappers: 

1072 return wrappers[-1][-1] 

1073 return None 

1074 

1075 

1076def _kron_dispatcher(a, b): 

1077 return (a, b) 

1078 

1079 

1080@array_function_dispatch(_kron_dispatcher) 

1081def kron(a, b): 

1082 """ 

1083 Kronecker product of two arrays. 

1084 

1085 Computes the Kronecker product, a composite array made of blocks of the 

1086 second array scaled by the first. 

1087 

1088 Parameters 

1089 ---------- 

1090 a, b : array_like 

1091 

1092 Returns 

1093 ------- 

1094 out : ndarray 

1095 

1096 See Also 

1097 -------- 

1098 outer : The outer product 

1099 

1100 Notes 

1101 ----- 

1102 The function assumes that the number of dimensions of `a` and `b` 

1103 are the same, if necessary prepending the smallest with ones. 

1104 If ``a.shape = (r0,r1,...,rN)`` and ``b.shape = (s0,s1,...,sN)``, 

1105 the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``. 

1106 The elements are products of elements from `a` and `b`, organized 

1107 explicitly by:: 

1108 

1109 kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN] 

1110 

1111 where:: 

1112 

1113 kt = it * st + jt, t = 0,...,N 

1114 

1115 In the common 2-D case (N=1), the block structure can be visualized::