Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-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 arrays = [] 

647 for v in tup: 

648 arr = asanyarray(v) 

649 if arr.ndim < 2: 

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

651 arrays.append(arr) 

652 return _nx.concatenate(arrays, 1) 

653 

654 

655def _dstack_dispatcher(tup): 

656 return _arrays_for_stack_dispatcher(tup) 

657 

658 

659@array_function_dispatch(_dstack_dispatcher) 

660def dstack(tup): 

661 """ 

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

663 

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

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

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

667 `dsplit`. 

668 

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

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

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

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

673 

674 Parameters 

675 ---------- 

676 tup : sequence of arrays 

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

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

679 

680 Returns 

681 ------- 

682 stacked : ndarray 

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

684 

685 See Also 

686 -------- 

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

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

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

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

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

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

693 dsplit : Split array along third axis. 

694 

695 Examples 

696 -------- 

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

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

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

700 array([[[1, 2], 

701 [2, 3], 

702 [3, 4]]]) 

703 

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

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

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

707 array([[[1, 2]], 

708 [[2, 3]], 

709 [[3, 4]]]) 

710 

711 """ 

712 arrs = atleast_3d(*tup) 

713 if not isinstance(arrs, list): 

714 arrs = [arrs] 

715 return _nx.concatenate(arrs, 2) 

716 

717 

718def _replace_zero_by_x_arrays(sub_arys): 

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

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

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

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

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

724 return sub_arys 

725 

726 

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

728 return (ary, indices_or_sections) 

729 

730 

731@array_function_dispatch(_array_split_dispatcher) 

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

733 """ 

734 Split an array into multiple sub-arrays. 

735 

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

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

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

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

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

741 and the rest of size l//n. 

742 

743 See Also 

744 -------- 

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

746 

747 Examples 

748 -------- 

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

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

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

752 

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

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

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

756 

757 """ 

758 try: 

759 Ntotal = ary.shape[axis] 

760 except AttributeError: 

761 Ntotal = len(ary) 

762 try: 

763 # handle array case. 

764 Nsections = len(indices_or_sections) + 1 

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

766 except TypeError: 

767 # indices_or_sections is a scalar, not an array. 

768 Nsections = int(indices_or_sections) 

769 if Nsections <= 0: 

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

771 Neach_section, extras = divmod(Ntotal, Nsections) 

772 section_sizes = ([0] + 

773 extras * [Neach_section+1] + 

774 (Nsections-extras) * [Neach_section]) 

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

776 

777 sub_arys = [] 

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

779 for i in range(Nsections): 

780 st = div_points[i] 

781 end = div_points[i + 1] 

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

783 

784 return sub_arys 

785 

786 

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

788 return (ary, indices_or_sections) 

789 

790 

791@array_function_dispatch(_split_dispatcher) 

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

793 """ 

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

795 

796 Parameters 

797 ---------- 

798 ary : ndarray 

799 Array to be divided into sub-arrays. 

800 indices_or_sections : int or 1-D array 

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

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

803 an error is raised. 

804 

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

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

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

808 

809 - ary[:2] 

810 - ary[2:3] 

811 - ary[3:] 

812 

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

814 an empty sub-array is returned correspondingly. 

815 axis : int, optional 

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

817 

818 Returns 

819 ------- 

820 sub-arrays : list of ndarrays 

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

822 

823 Raises 

824 ------ 

825 ValueError 

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

827 a split does not result in equal division. 

828 

829 See Also 

830 -------- 

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

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

833 an equal division cannot be made. 

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

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

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

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

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

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

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

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

842 

843 Examples 

844 -------- 

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

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

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

848 

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

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

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

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

853 array([5.]), 

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

855 array([], dtype=float64)] 

856 

857 """ 

858 try: 

859 len(indices_or_sections) 

860 except TypeError: 

861 sections = indices_or_sections 

862 N = ary.shape[axis] 

863 if N % sections: 

864 raise ValueError( 

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

866 return array_split(ary, indices_or_sections, axis) 

867 

868 

869def _hvdsplit_dispatcher(ary, indices_or_sections): 

870 return (ary, indices_or_sections) 

871 

872 

873@array_function_dispatch(_hvdsplit_dispatcher) 

874def hsplit(ary, indices_or_sections): 

875 """ 

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

877 

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

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

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

881 

882 See Also 

883 -------- 

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

885 

886 Examples 

887 -------- 

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

889 >>> x 

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

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

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

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

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

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

896 [ 4., 5.], 

897 [ 8., 9.], 

898 [12., 13.]]), 

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

900 [ 6., 7.], 

901 [10., 11.], 

902 [14., 15.]])] 

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

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

905 [ 4., 5., 6.], 

906 [ 8., 9., 10.], 

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

908 array([[ 3.], 

909 [ 7.], 

910 [11.], 

911 [15.]]), 

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

913 

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

915 

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

917 >>> x 

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

919 [2., 3.]], 

920 [[4., 5.], 

921 [6., 7.]]]) 

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

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

924 [[4., 5.]]]), 

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

926 [[6., 7.]]])] 

927 

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

929 

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

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

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

933 

934 """ 

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

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

937 if ary.ndim > 1: 

938 return split(ary, indices_or_sections, 1) 

939 else: 

940 return split(ary, indices_or_sections, 0) 

941 

942 

943@array_function_dispatch(_hvdsplit_dispatcher) 

944def vsplit(ary, indices_or_sections): 

945 """ 

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

947 

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

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

950 first axis regardless of the array dimension. 

951 

952 See Also 

953 -------- 

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

955 

956 Examples 

957 -------- 

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

959 >>> x 

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

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

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

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

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

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

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

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

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

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

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

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

972 

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

974 

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

976 >>> x 

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

978 [2., 3.]], 

979 [[4., 5.], 

980 [6., 7.]]]) 

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

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

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

984 [6., 7.]]])] 

985 

986 """ 

987 if _nx.ndim(ary) < 2: 

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

989 return split(ary, indices_or_sections, 0) 

990 

991 

992@array_function_dispatch(_hvdsplit_dispatcher) 

993def dsplit(ary, indices_or_sections): 

994 """ 

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

996 

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

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

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

1000 

1001 See Also 

1002 -------- 

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

1004 

1005 Examples 

1006 -------- 

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

1008 >>> x 

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

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

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

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

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

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

1015 [ 4., 5.]], 

1016 [[ 8., 9.], 

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

1018 [ 6., 7.]], 

1019 [[10., 11.], 

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

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

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

1023 [ 4., 5., 6.]], 

1024 [[ 8., 9., 10.], 

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

1026 array([[[ 3.], 

1027 [ 7.]], 

1028 [[11.], 

1029 [15.]]]), 

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

1031 """ 

1032 if _nx.ndim(ary) < 3: 

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

1034 return split(ary, indices_or_sections, 2) 

1035 

1036 

1037def get_array_prepare(*args): 

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

1039 

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

1041 """ 

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

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

1044 if hasattr(x, '__array_prepare__')) 

1045 if wrappers: 

1046 return wrappers[-1][-1] 

1047 return None 

1048 

1049 

1050def get_array_wrap(*args): 

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

1052 

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

1054 """ 

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

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

1057 if hasattr(x, '__array_wrap__')) 

1058 if wrappers: 

1059 return wrappers[-1][-1] 

1060 return None 

1061 

1062 

1063def _kron_dispatcher(a, b): 

1064 return (a, b) 

1065 

1066 

1067@array_function_dispatch(_kron_dispatcher) 

1068def kron(a, b): 

1069 """ 

1070 Kronecker product of two arrays. 

1071 

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

1073 second array scaled by the first. 

1074 

1075 Parameters 

1076 ---------- 

1077 a, b : array_like 

1078 

1079 Returns 

1080 ------- 

1081 out : ndarray 

1082 

1083 See Also 

1084 -------- 

1085 outer : The outer product 

1086 

1087 Notes 

1088 ----- 

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

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

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

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

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

1094 explicitly by:: 

1095 

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

1097 

1098 where:: 

1099 

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

1101 

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

1103