1 import functools

2

3 import numpy.core.numeric as _nx

4 from numpy.core.numeric import (

5 asarray, zeros, outer, concatenate, array, asanyarray

6 )

7 from numpy.core.fromnumeric import reshape, transpose

8 from numpy.core.multiarray import normalize_axis_index

9 from numpy.core import overrides

10 from numpy.core import vstack, atleast_3d

11 from numpy.core.numeric import normalize_axis_tuple

12 from numpy.core.shape_base import _arrays_for_stack_dispatcher

13 from numpy.lib.index_tricks import ndindex

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

15

16

17__all__ = [

18 'column_stack', 'row_stack', 'dstack', 'array_split', 'split',

19 'hsplit', 'vsplit', 'dsplit', 'apply_over_axes', 'expand_dims',

20 'apply_along_axis', 'kron', 'tile', 'get_array_wrap', 'take_along_axis',

21 'put_along_axis'

22 ]

23

24

25 array_function_dispatch = functools.partial(

26 overrides.array_function_dispatch, module='numpy')

27

28

29 def _make_along_axis_idx(arr_shape, indices, axis):

30 # compute dimensions to iterate over

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

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

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

34 raise ValueError(

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

36 shape_ones = (1,) * indices.ndim

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

38

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

40 # requested index inserted at the right location

41 fancy_index = []

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

43 if dim is None:

44 fancy_index.append(indices)

45 else:

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

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

48

49 return tuple(fancy_index)

50

51

52 def _take_along_axis_dispatcher(arr, indices, axis):

53 return (arr, indices)

54

55

56@array_function_dispatch(_take_along_axis_dispatcher)

57 def take_along_axis(arr, indices, axis):

58 """

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

60

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

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

63 latter. These slices can be different lengths.

64

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

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

67

68 .. versionadded:: 1.15.0

69

70 Parameters

71 ----------

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

73 Source array

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

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

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

77 against `arr`.

78 axis : int

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

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

81 `sort` and `argsort`.

82

83 Returns

84 -------

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

86 The indexed result.

87

88 Notes

89 -----

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

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

92

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

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

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

96

97 for ii in ndindex(Ni):

98 for kk in ndindex(Nk):

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

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

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

102 for j in range(J):

103 out_1d[j] = a_1d[indices_1d[j]]

104

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

106

107 out_1d[:] = a_1d[indices_1d]

108

109 See Also

110 --------

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

112 put_along_axis :

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

114

115 Examples

116 --------

117

118 For this sample array

119

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

121

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

123

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

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

126 [40, 50, 60]])

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

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

129 [1, 2, 0]])

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

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

132 [40, 50, 60]])

133

134 The same works for max and min, if you expand the dimensions:

135

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

137 array([[30],

138 [60]])

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

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.expand_dims(np.argmin(a, axis=1), axis=1)

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

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

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

174 return (arr, indices, values)

175

176

177@array_function_dispatch(_put_along_axis_dispatcher)

178 def 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.expand_dims(np.argmax(a, axis=1), axis=1)

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

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

264 return (arr,)

265

266

267@array_function_dispatch(_apply_along_axis_dispatcher)

268 def 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 as e:

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

417 def _apply_over_axes_dispatcher(func, a, axes):

418 return (a,)

419

420

421@array_function_dispatch(_apply_over_axes_dispatcher)

422 def 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

508 def _expand_dims_dispatcher(a, axis):

509 return (a,)

510

511

512@array_function_dispatch(_expand_dims_dispatcher)

513 def 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

605 row_stack = vstack

606

607

608 def _column_stack_dispatcher(tup):

609 return _arrays_for_stack_dispatcher(tup)

610

611

612@array_function_dispatch(_column_stack_dispatcher)

613 def 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

659 def _dstack_dispatcher(tup):

660 return _arrays_for_stack_dispatcher(tup)

661

662

663@array_function_dispatch(_dstack_dispatcher)

664 def 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

726 def _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

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

736 return (ary, indices_or_sections)

737

738

739@array_function_dispatch(_array_split_dispatcher)

740 def 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

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

796 return (ary, indices_or_sections)

797

798

799@array_function_dispatch(_split_dispatcher)

800 def 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

877 def _hvdsplit_dispatcher(ary, indices_or_sections):

878 return (ary, indices_or_sections)

879

880

881@array_function_dispatch(_hvdsplit_dispatcher)

882 def 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)

952 def 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)

1001 def 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 def get_array_prepare(*args):

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

1046

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

1048 """

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

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

1051 if hasattr(x, '__array_prepare__'))

1052 if wrappers:

1053 return wrappers[-1][-1]

1054 return None

1055

1056 def get_array_wrap(*args):

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

1058

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

1060 """

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

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

1063 if hasattr(x, '__array_wrap__'))

1064 if wrappers:

1065 return wrappers[-1][-1]

1066 return None

1067

1068

1069 def _kron_dispatcher(a, b):

1070 return (a, b)

1071

1072

1073@array_function_dispatch(_kron_dispatcher)

1074 def kron(a, b):

1075 """

1076 Kronecker product of two arrays.

1077

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

1079 second array scaled by the first.

1080

1081 Parameters

1082 ----------

1083 a, b : array_like

1084

1085 Returns

1086 -------

1087 out : ndarray

1088

1089 See Also

1090 --------

1091 outer : The outer product

1092

1093 Notes

1094 -----

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

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

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

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

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

1100 explicitly by::

1101

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

1103

1104 where::

1105

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

1107

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

1109

1110 [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ],

1111 [ ... ... ],

1112 [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]]

1113

1114

1115 Examples

1116 --------

1117 >>> np.kron([1,10,100], [5,6,7])

1118 array([ 5, 6, 7, ..., 500, 600, 700])

1119 >>> np.kron([5,6,7], [1,10,100])

1120 array([ 5, 50, 500, ..., 7, 70, 700])

1121

1122 >>> np.kron(np.eye(2), np.ones((2,2)))

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

1124 [1., 1., 0., 0.],

1125 [0., 0., 1., 1.],

1126 [0., 0., 1., 1.]])

1127

1128 >>> a = np.arange(100).reshape((2,5,2,5))

1129 >>> b = np.arange(24).reshape((2,3,4))

1130 >>> c = np.kron(a,b)

1131 >>> c.shape

1132 (2, 10, 6, 20)

1133 >>> I = (1,3,0,2)

1134 >>> J = (0,2,1)

1135 >>> J1 = (0,) + J # extend to ndim=4

1136 >>> S1 = (1,) + b.shape

1137 >>> K = tuple(np.array(I) * np.array(S1) + np.array(J1))

1138 >>> c[K] == a[I]*b[J]

1139 True

1140

1141 """

1142 # Working:

1143 # 1. Equalise the shapes by prepending smaller array with 1s

1144 # 2. Expand shapes of both the arrays by adding new axes at

1145 # odd positions for 1st array and even positions for 2nd

1146 # 3. Compute the product of the modified array

1147 # 4. The inner most array elements now contain the rows of

1148 # the Kronecker product

1149 # 5. Reshape the result to kron's shape, which is same as

1150 # product of shapes of the two arrays.

1151 b = asanyarray(b)

1152 a = array(a, copy=False, subok=True, ndmin=b.ndim)

1153 is_any_mat = isinstance(a, matrix) or isinstance(b, matrix)

1154 ndb, nda = b.ndim, a.ndim

1155 nd = max(ndb, nda)

1156

1157 if (nda == 0 or ndb == 0):

1158 return _nx.multiply(a, b)

1159

1160 as_ = a.shape

1161 bs = b.shape

1162 if not a.flags.contiguous:

1163 a = reshape(a, as_)

1164 if not b.flags.contiguous:

1165 b = reshape(b, bs)

1166

1167 # Equalise the shapes by prepending smaller one with 1s

1168 as_ = (1,)*max(0, ndb-nda) + as_

1169 bs = (1,)*max(0, nda-ndb) + bs

1170

1171 # Insert empty dimensions

1172 a_arr = expand_dims(a, axis=tuple(range(ndb-nda)))

1173 b_arr = expand_dims(b, axis=tuple(range(nda-ndb)))

1174

1175 # Compute the product

1176 a_arr = expand_dims(a_arr, axis=tuple(range(1, nd*2, 2)))

1177 b_arr = expand_dims(b_arr, axis=tuple(range(0, nd*2, 2)))

1178 # In case of `mat`, convert result to `array`

1179 result = _nx.multiply(a_arr, b_arr, subok=(not is_any_mat))

1180

1181 # Reshape back

1182 result = result.reshape(_nx.multiply(as_, bs))

1183

1184 return result if not is_any_mat else matrix(result, copy=False)

1185

1186

1187 def _tile_dispatcher(A, reps):

1188 return (A, reps)

1189

1190

1191@array_function_dispatch(_tile_dispatcher)

1192 def tile(A, reps):

1193 """

1194 Construct an array by repeating A the number of times given by reps.

1195

1196 If `reps` has length ``d``, the result will have dimension of

1197 ``max(d, A.ndim)``.

1198

1199 If ``A.ndim < d``, `A` is promoted to be d-dimensional by prepending new

1200 axes. So a shape (3,) array is promoted to (1, 3) for 2-D replication,

1201 or shape (1, 1, 3) for 3-D replication. If this is not the desired

1202 behavior, promote `A` to d-dimensions manually before calling this

1203 function.

1204

1205 If ``A.ndim > d``, `reps` is promoted to `A`.ndim by pre-pending 1's to it.

1206 Thus for an `A` of shape (2, 3, 4, 5), a `reps` of (2, 2) is treated as

1207 (1, 1, 2, 2).

1208

1209 Note : Although tile may be used for broadcasting, it is strongly

1210 recommended to use numpy's broadcasting operations and functions.

1211

1212 Parameters

1213 ----------

1214 A : array_like

1215 The input array.

1216 reps : array_like

1217 The number of repetitions of `A` along each axis.

1218

1219 Returns

1220 -------

1221 c : ndarray

1222 The tiled output array.

1223

1224 See Also

1225 --------

1226 repeat : Repeat elements of an array.

1227 broadcast_to : Broadcast an array to a new shape

1228

1229 Examples

1230 --------

1231 >>> a = np.array([0, 1, 2])

1232 >>> np.tile(a, 2)

1233 array([0, 1, 2, 0, 1, 2])

1234 >>> np.tile(a, (2, 2))

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

1236 [0, 1, 2, 0, 1, 2]])

1237 >>> np.tile(a, (2, 1, 2))

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

1239 [[0, 1, 2, 0, 1, 2]]])

1240

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

1242 >>> np.tile(b, 2)

1243 array([[1, 2, 1, 2],

1244 [3, 4, 3, 4]])

1245 >>> np.tile(b, (2, 1))

1246 array([[1, 2],

1247 [3, 4],

1248 [1, 2],

1249 [3, 4]])

1250

1251 >>> c = np.array([1,2,3,4])

1252 >>> np.tile(c,(4,1))

1253 array([[1, 2, 3, 4],

1254 [1, 2, 3, 4],

1255 [1, 2, 3, 4],

1256 [1, 2, 3, 4]])

1257 """

1258 try:

1259 tup = tuple(reps)

1260 except TypeError:

1261 tup = (reps,)

1262 d = len(tup)

1263 if all(x == 1 for x in tup) and isinstance(A, _nx.ndarray):

1264 # Fixes the problem that the function does not make a copy if A is a

1265 # numpy array and the repetitions are 1 in all dimensions

1266 return _nx.array(A, copy=True, subok=True, ndmin=d)

1267 else:

1268 # Note that no copy of zero-sized arrays is made. However since they

1269 # have no data there is no risk of an inadvertent overwrite.

1270 c = _nx.array(A, copy=False, subok=True, ndmin=d)

1271 if (d < c.ndim):

1272 tup = (1,)*(c.ndim-d) + tup

1273 shape_out = tuple(s*t for s, t in zip(c.shape, tup))

1274 n = c.size

1275 if n > 0:

1276 for dim_in, nrep in zip(c.shape, tup):

1277 if nrep != 1:

1278 c = c.reshape(-1, n).repeat(nrep, 0)

1279 n //= dim_in

1280 return c.reshape(shape_out)