Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/numpy/lib/shape_base.py: 23%

1import functools 

2 

3import numpy.core.numeric as _nx 

4from numpy.core.numeric import asarray, zeros, array, asanyarray 

5from numpy.core.fromnumeric import reshape, transpose 

6from numpy.core.multiarray import normalize_axis_index 

7from numpy.core import overrides 

8from numpy.core import vstack, atleast_3d 

9from numpy.core.numeric import normalize_axis_tuple 

10from numpy.core.shape_base import _arrays_for_stack_dispatcher 

11from numpy.lib.index_tricks import ndindex 

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

13 

14 

15__all__ = [ 

16 'column_stack', 'row_stack', 'dstack', 'array_split', 'split', 

17 'hsplit', 'vsplit', 'dsplit', 'apply_over_axes', 'expand_dims', 

18 'apply_along_axis', 'kron', 'tile', 'get_array_wrap', 'take_along_axis', 

19 'put_along_axis' 

20 ] 

21 

22 

23array_function_dispatch = functools.partial( 

24 overrides.array_function_dispatch, module='numpy') 

25 

26 

27def _make_along_axis_idx(arr_shape, indices, axis): 

28 # compute dimensions to iterate over 

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

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

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

32 raise ValueError( 

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

34 shape_ones = (1,) * indices.ndim 

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

36 

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

38 # requested index inserted at the right location 

39 fancy_index = [] 

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

41 if dim is None: 

42 fancy_index.append(indices) 

43 else: 

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

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

46 

47 return tuple(fancy_index) 

48 

49 

50def _take_along_axis_dispatcher(arr, indices, axis): 

51 return (arr, indices) 

52 

53 

54@array_function_dispatch(_take_along_axis_dispatcher) 

55def take_along_axis(arr, indices, axis): 

56 """ 

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

58 

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

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

61 latter. These slices can be different lengths. 

62 

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

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

65 

66 .. versionadded:: 1.15.0 

67 

68 Parameters 

69 ---------- 

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

71 Source array 

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

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

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

75 against `arr`. 

76 axis : int 

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

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

79 `sort` and `argsort`. 

80 

81 Returns 

82 ------- 

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

84 The indexed result. 

85 

86 Notes 

87 ----- 

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

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

90 

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

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

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

94 

95 for ii in ndindex(Ni): 

96 for kk in ndindex(Nk): 

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

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

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

100 for j in range(J): 

101 out_1d[j] = a_1d[indices_1d[j]] 

102 

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

104 

105 out_1d[:] = a_1d[indices_1d] 

106 

107 See Also 

108 -------- 

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

110 put_along_axis : 

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

112 

113 Examples 

114 -------- 

115 

116 For this sample array 

117 

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

119 

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

121 

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

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

124 [40, 50, 60]]) 

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

126 >>> ai 

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

128 [1, 2, 0]]) 

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

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

131 [40, 50, 60]]) 

132 

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

134 with ``keepdims``: 

135 

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

137 array([[30], 

138 [60]]) 

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

140 >>> ai 

141 array([[1], 

142 [0]]) 

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

144 array([[30], 

145 [60]]) 

146 

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

148 indices first 

149 

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

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

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

153 >>> ai 

154 array([[0, 1], 

155 [1, 0]]) 

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

157 array([[10, 30], 

158 [40, 60]]) 

159 """ 

160 # normalize inputs 

161 if axis is None: 

162 arr = arr.flat 

163 arr_shape = (len(arr),) # flatiter has no .shape 

164 axis = 0 

165 else: 

166 axis = normalize_axis_index(axis, arr.ndim) 

167 arr_shape = arr.shape 

168 

169 # use the fancy index 

170 return arr[_make_along_axis_idx(arr_shape, indices, axis)] 

171 

172 

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

174 return (arr, indices, values) 

175 

176 

177@array_function_dispatch(_put_along_axis_dispatcher) 

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

179 """ 

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

181 

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

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

184 latter. These slices can be different lengths. 

185 

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

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

188 

189 .. versionadded:: 1.15.0 

190 

191 Parameters 

192 ---------- 

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

194 Destination array. 

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

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

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

198 against `arr`. 

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

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

201 broadcast to match that of `indices`. 

202 axis : int 

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

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

205 

206 Notes 

207 ----- 

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

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

210 

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

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

213 

214 for ii in ndindex(Ni): 

215 for kk in ndindex(Nk): 

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

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

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

219 for j in range(J): 

220 a_1d[indices_1d[j]] = values_1d[j] 

221 

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

223 

224 a_1d[indices_1d] = values_1d 

225 

226 See Also 

227 -------- 

228 take_along_axis : 

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

230 

231 Examples 

232 -------- 

233 

234 For this sample array 

235 

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

237 

238 We can replace the maximum values with: 

239 

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

241 >>> ai 

242 array([[1], 

243 [0]]) 

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

245 >>> a 

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

247 [99, 40, 50]]) 

248 

249 """ 

250 # normalize inputs 

251 if axis is None: 

252 arr = arr.flat 

253 axis = 0 

254 arr_shape = (len(arr),) # flatiter has no .shape 

255 else: 

256 axis = normalize_axis_index(axis, arr.ndim) 

257 arr_shape = arr.shape 

258 

259 # use the fancy index 

260 arr[_make_along_axis_idx(arr_shape, indices, axis)] = values 

261 

262 

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

264 return (arr,) 

265 

266 

267@array_function_dispatch(_apply_along_axis_dispatcher) 

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

269 """ 

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

271 

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

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

274 

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

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

277 

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

279 for ii in ndindex(Ni): 

280 for kk in ndindex(Nk): 

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

282 Nj = f.shape 

283 for jj in ndindex(Nj): 

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

285 

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

287 

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

289 for ii in ndindex(Ni): 

290 for kk in ndindex(Nk): 

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

292 

293 Parameters 

294 ---------- 

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

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

297 slices of `arr` along the specified axis. 

298 axis : integer 

299 Axis along which `arr` is sliced. 

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

301 Input array. 

302 args : any 

303 Additional arguments to `func1d`. 

304 kwargs : any 

305 Additional named arguments to `func1d`. 

306 

307 .. versionadded:: 1.9.0 

308 

309 

310 Returns 

311 ------- 

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

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

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

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

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

317 fewer dimensions than `arr`. 

318 

319 See Also 

320 -------- 

321 apply_over_axes : Apply a function repeatedly over multiple axes. 

322 

323 Examples 

324 -------- 

325 >>> def my_func(a): 

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

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

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

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

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

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

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

333 

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

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

336 

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

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

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

340 [3, 4, 9], 

341 [2, 5, 6]]) 

342 

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

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

345 

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

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

348 array([[[1, 0, 0], 

349 [0, 2, 0], 

350 [0, 0, 3]], 

351 [[4, 0, 0], 

352 [0, 5, 0], 

353 [0, 0, 6]], 

354 [[7, 0, 0], 

355 [0, 8, 0], 

356 [0, 0, 9]]]) 

357 """ 

358 # handle negative axes 

359 arr = asanyarray(arr) 

360 nd = arr.ndim 

361 axis = normalize_axis_index(axis, nd) 

362 

363 # arr, with the iteration axis at the end 

364 in_dims = list(range(nd)) 

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

366 

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

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

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

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

371 

372 # invoke the function on the first item 

373 try: 

374 ind0 = next(inds) 

375 except StopIteration: 

376 raise ValueError( 

377 'Cannot apply_along_axis when any iteration dimensions are 0' 

378 ) from None 

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

380 

381 # build a buffer for storing evaluations of func1d. 

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

383 # laid out so that each write is contiguous. 

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

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

386 

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

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

389 buff_permute = ( 

390 buff_dims[0 : axis] + 

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

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

393 ) 

394 

395 # matrices have a nasty __array_prepare__ and __array_wrap__ 

396 if not isinstance(res, matrix): 

397 buff = res.__array_prepare__(buff) 

398 

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

400 buff[ind0] = res 

401 for ind in inds: 

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

403 

404 if not isinstance(res, matrix): 

405 # wrap the array, to preserve subclasses 

406 buff = res.__array_wrap__(buff) 

407 

408 # finally, rotate the inserted axes back to where they belong 

409 return transpose(buff, buff_permute) 

410 

411 else: 

412 # matrices have to be transposed first, because they collapse dimensions! 

413 out_arr = transpose(buff, buff_permute) 

414 return res.__array_wrap__(out_arr) 

415 

416 

417def _apply_over_axes_dispatcher(func, a, axes): 

418 return (a,) 

419 

420 

421@array_function_dispatch(_apply_over_axes_dispatcher) 

422def apply_over_axes(func, a, axes): 

423 """ 

424 Apply a function repeatedly over multiple axes. 

425 

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

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

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

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

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

431 with `res` as the first argument. 

432 

433 Parameters 

434 ---------- 

435 func : function 

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

437 a : array_like 

438 Input array. 

439 axes : array_like 

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

441 

442 Returns 

443 ------- 

444 apply_over_axis : ndarray 

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

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

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

448 

449 See Also 

450 -------- 

451 apply_along_axis : 

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

453 

454 Notes 

455 ----- 

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

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

458 version 1.7.0. 

459 

460 Examples 

461 -------- 

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

463 >>> a 

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

465 [ 4, 5, 6, 7], 

466 [ 8, 9, 10, 11]], 

467 [[12, 13, 14, 15], 

468 [16, 17, 18, 19], 

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

470 

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

472 as the original array: 

473 

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

475 array([[[ 60], 

476 [ 92], 

477 [124]]]) 

478 

479 Tuple axis arguments to ufuncs are equivalent: 

480 

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

482 array([[[ 60], 

483 [ 92], 

484 [124]]]) 

485 

486 """ 

487 val = asarray(a) 

488 N = a.ndim 

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

490 axes = (axes,) 

491 for axis in axes: 

492 if axis < 0: 

493 axis = N + axis 

494 args = (val, axis) 

495 res = func(*args) 

496 if res.ndim == val.ndim: 

497 val = res 

498 else: 

499 res = expand_dims(res, axis) 

500 if res.ndim == val.ndim: 

501 val = res 

502 else: 

503 raise ValueError("function is not returning " 

504 "an array of the correct shape") 

505 return val 

506 

507 

508def _expand_dims_dispatcher(a, axis): 

509 return (a,) 

510 

511 

512@array_function_dispatch(_expand_dims_dispatcher) 

513def expand_dims(a, axis): 

514 """ 

515 Expand the shape of an array. 

516 

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

518 array shape. 

519 

520 Parameters 

521 ---------- 

522 a : array_like 

523 Input array. 

524 axis : int or tuple of ints 

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

526 

527 .. deprecated:: 1.13.0 

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

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

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

531 

532 .. versionchanged:: 1.18.0 

533 A tuple of axes is now supported. Out of range axes as 

534 described above are now forbidden and raise an `AxisError`. 

535 

536 Returns 

537 ------- 

538 result : ndarray 

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

540 

541 See Also 

542 -------- 

543 squeeze : The inverse operation, removing singleton dimensions 

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

545 doc.indexing, atleast_1d, atleast_2d, atleast_3d 

546 

547 Examples 

548 -------- 

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 type(axis) not in (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 

605row_stack = vstack 

606 

607 

608def _column_stack_dispatcher(tup): 

609 return _arrays_for_stack_dispatcher(tup) 

610 

611 

612@array_function_dispatch(_column_stack_dispatcher) 

613def column_stack(tup): 

614 """ 

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

616 

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

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

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

620 first. 

621 

622 Parameters 

623 ---------- 

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

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

626 

627 Returns 

628 ------- 

629 stacked : 2-D array 

630 The array formed by stacking the given arrays. 

631 

632 See Also 

633 -------- 

634 stack, hstack, vstack, concatenate 

635 

636 Examples 

637 -------- 

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

639 >>> b = np.array((2,3,4)) 

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

641 array([[1, 2], 

642 [2, 3], 

643 [3, 4]]) 

644 

645 """ 

646 if not overrides.ARRAY_FUNCTION_ENABLED: 

647 # raise warning if necessary 

648 _arrays_for_stack_dispatcher(tup, stacklevel=2) 

649 

650 arrays = [] 

651 for v in tup: 

652 arr = asanyarray(v) 

653 if arr.ndim < 2: 

654 arr = array(arr, copy=False, subok=True, ndmin=2).T 

655 arrays.append(arr) 

656 return _nx.concatenate(arrays, 1) 

657 

658 

659def _dstack_dispatcher(tup): 

660 return _arrays_for_stack_dispatcher(tup) 

661 

662 

663@array_function_dispatch(_dstack_dispatcher) 

664def dstack(tup): 

665 """ 

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

667 

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

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

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

671 `dsplit`. 

672 

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

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

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

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

677 

678 Parameters 

679 ---------- 

680 tup : sequence of arrays 

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

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

683 

684 Returns 

685 ------- 

686 stacked : ndarray 

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

688 

689 See Also 

690 -------- 

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

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

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

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

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

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

697 dsplit : Split array along third axis. 

698 

699 Examples 

700 -------- 

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

702 >>> b = np.array((2,3,4)) 

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

704 array([[[1, 2], 

705 [2, 3], 

706 [3, 4]]]) 

707 

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

709 >>> b = np.array([[2],[3],[4]]) 

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

711 array([[[1, 2]], 

712 [[2, 3]], 

713 [[3, 4]]]) 

714 

715 """ 

716 if not overrides.ARRAY_FUNCTION_ENABLED: 

717 # raise warning if necessary 

718 _arrays_for_stack_dispatcher(tup, stacklevel=2) 

719 

720 arrs = atleast_3d(*tup) 

721 if not isinstance(arrs, list): 

722 arrs = [arrs] 

723 return _nx.concatenate(arrs, 2) 

724 

725 

726def _replace_zero_by_x_arrays(sub_arys): 

727 for i in range(len(sub_arys)): 

728 if _nx.ndim(sub_arys[i]) == 0: 

729 sub_arys[i] = _nx.empty(0, dtype=sub_arys[i].dtype) 

730 elif _nx.sometrue(_nx.equal(_nx.shape(sub_arys[i]), 0)): 

731 sub_arys[i] = _nx.empty(0, dtype=sub_arys[i].dtype) 

732 return sub_arys 

733 

734 

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

736 return (ary, indices_or_sections) 

737 

738 

739@array_function_dispatch(_array_split_dispatcher) 

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

741 """ 

742 Split an array into multiple sub-arrays. 

743 

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

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

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

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

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

749 and the rest of size l//n. 

750 

751 See Also 

752 -------- 

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

754 

755 Examples 

756 -------- 

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

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

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

760 

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

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

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

764 

765 """ 

766 try: 

767 Ntotal = ary.shape[axis] 

768 except AttributeError: 

769 Ntotal = len(ary) 

770 try: 

771 # handle array case. 

772 Nsections = len(indices_or_sections) + 1 

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

774 except TypeError: 

775 # indices_or_sections is a scalar, not an array. 

776 Nsections = int(indices_or_sections) 

777 if Nsections <= 0: 

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

779 Neach_section, extras = divmod(Ntotal, Nsections) 

780 section_sizes = ([0] + 

781 extras * [Neach_section+1] + 

782 (Nsections-extras) * [Neach_section]) 

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

784 

785 sub_arys = [] 

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

787 for i in range(Nsections): 

788 st = div_points[i] 

789 end = div_points[i + 1] 

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

791 

792 return sub_arys 

793 

794 

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

796 return (ary, indices_or_sections) 

797 

798 

799@array_function_dispatch(_split_dispatcher) 

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

801 """ 

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

803 

804 Parameters 

805 ---------- 

806 ary : ndarray 

807 Array to be divided into sub-arrays. 

808 indices_or_sections : int or 1-D array 

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

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

811 an error is raised. 

812 

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

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

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

816 

817 - ary[:2] 

818 - ary[2:3] 

819 - ary[3:] 

820 

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

822 an empty sub-array is returned correspondingly. 

823 axis : int, optional 

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

825 

826 Returns 

827 ------- 

828 sub-arrays : list of ndarrays 

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

830 

831 Raises 

832 ------ 

833 ValueError 

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

835 a split does not result in equal division. 

836 

837 See Also 

838 -------- 

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

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

841 an equal division cannot be made. 

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

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

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

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

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

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

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

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

850 

851 Examples 

852 -------- 

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 >>> x = np.arange(16.0).reshape(4, 4) 

897 >>> x 

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

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

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

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

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

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

904 [ 4., 5.], 

905 [ 8., 9.], 

906 [12., 13.]]), 

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

908 [ 6., 7.], 

909 [10., 11.], 

910 [14., 15.]])] 

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

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

913 [ 4., 5., 6.], 

914 [ 8., 9., 10.], 

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

916 array([[ 3.], 

917 [ 7.], 

918 [11.], 

919 [15.]]), 

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

921 

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

923 

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

925 >>> x 

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

927 [2., 3.]], 

928 [[4., 5.], 

929 [6., 7.]]]) 

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

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

932 [[4., 5.]]]), 

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

934 [[6., 7.]]])] 

935 

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

937 

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

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

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

941 

942 """ 

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

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

945 if ary.ndim > 1: 

946 return split(ary, indices_or_sections, 1) 

947 else: 

948 return split(ary, indices_or_sections, 0) 

949 

950 

951@array_function_dispatch(_hvdsplit_dispatcher) 

952def vsplit(ary, indices_or_sections): 

953 """ 

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

955 

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

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

958 first axis regardless of the array dimension. 

959 

960 See Also 

961 -------- 

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

963 

964 Examples 

965 -------- 

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

967 >>> x 

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

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

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

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

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

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

974 [4., 5., 6., 7.]]), array([[ 8., 9., 10., 11.], 

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

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

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

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

979 [ 8., 9., 10., 11.]]), array([[12., 13., 14., 15.]]), array([], shape=(0, 4), dtype=float64)] 

980 

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

982 

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

984 >>> x 

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

986 [2., 3.]], 

987 [[4., 5.], 

988 [6., 7.]]]) 

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

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

991 [2., 3.]]]), array([[[4., 5.], 

992 [6., 7.]]])] 

993 

994 """ 

995 if _nx.ndim(ary) < 2: 

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

997 return split(ary, indices_or_sections, 0) 

998 

999 

1000@array_function_dispatch(_hvdsplit_dispatcher) 

1001def dsplit(ary, indices_or_sections): 

1002 """ 

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

1004 

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

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

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

1008 

1009 See Also 

1010 -------- 

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

1012 

1013 Examples 

1014 -------- 

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

1016 >>> x 

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

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

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

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

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

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

1023 [ 4., 5.]], 

1024 [[ 8., 9.], 

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

1026 [ 6., 7.]], 

1027 [[10., 11.], 

1028 [14., 15.]]])] 

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

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

1031 [ 4., 5., 6.]], 

1032 [[ 8., 9., 10.], 

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

1034 array([[[ 3.], 

1035 [ 7.]], 

1036 [[11.], 

1037 [15.]]]), 

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

1039 """ 

1040 if _nx.ndim(ary) < 3: 

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

1042 return split(ary, indices_or_sections, 2) 

1043 

1044 

1045def get_array_prepare(*args): 

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

1047 

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

1049 """ 

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

1051 x.__array_prepare__) for i, x in enumerate(args) 

1052 if hasattr(x, '__array_prepare__')) 

1053 if wrappers: 

1054 return wrappers[-1][-1] 

1055 return None 

1056 

1057 

1058def get_array_wrap(*args): 

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

1060 

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

1062 """ 

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

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

1065 if hasattr(x, '__array_wrap__')) 

1066 if wrappers: 

1067 return wrappers[-1][-1] 

1068 return None 

1069 

1070 

1071def _kron_dispatcher(a, b): 

1072 return (a, b) 

1073 

1074 

1075@array_function_dispatch(_kron_dispatcher) 

1076def kron(a, b): 

1077 """ 

1078 Kronecker product of two arrays. 

1079 

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

1081 second array scaled by the first. 

1082 

1083 Parameters 

1084 ---------- 

1085 a, b : array_like 

1086 

1087 Returns 

1088 ------- 

1089 out : ndarray 

1090 

1091 See Also 

1092 -------- 

1093 outer : The outer product 

1094 

1095 Notes 

1096 ----- 

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

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

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