Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/numpy/lib/_index_tricks

1import functools

2import sys

3import math

4import warnings

6import numpy as np

7from .._utils import set_module

8import numpy._core.numeric as _nx

9from numpy._core.numeric import ScalarType, array

10from numpy._core.numerictypes import issubdtype

12import numpy.matrixlib as matrixlib

13from numpy._core.multiarray import ravel_multi_index, unravel_index

14from numpy._core import overrides, linspace

15from numpy.lib.stride_tricks import as_strided

16from numpy.lib._function_base_impl import diff

19array_function_dispatch = functools.partial(

20 overrides.array_function_dispatch, module='numpy')

23__all__ = [

24 'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_',

25 's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal',

26 'diag_indices', 'diag_indices_from'

27]

30def _ix__dispatcher(*args):

31 return args

34@array_function_dispatch(_ix__dispatcher)

35def ix_(*args):

36 """

37 Construct an open mesh from multiple sequences.

39 This function takes N 1-D sequences and returns N outputs with N

40 dimensions each, such that the shape is 1 in all but one dimension

41 and the dimension with the non-unit shape value cycles through all

42 N dimensions.

44 Using `ix_` one can quickly construct index arrays that will index

45 the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array

46 ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``.

48 Parameters

49 ----------

50 args : 1-D sequences

51 Each sequence should be of integer or boolean type.

52 Boolean sequences will be interpreted as boolean masks for the

53 corresponding dimension (equivalent to passing in

54 ``np.nonzero(boolean_sequence)``).

56 Returns

57 -------

58 out : tuple of ndarrays

59 N arrays with N dimensions each, with N the number of input

60 sequences. Together these arrays form an open mesh.

62 See Also

63 --------

64 ogrid, mgrid, meshgrid

66 Examples

67 --------

68 >>> a = np.arange(10).reshape(2, 5)

69 >>> a

70 array([[0, 1, 2, 3, 4],

71 [5, 6, 7, 8, 9]])

72 >>> ixgrid = np.ix_([0, 1], [2, 4])

73 >>> ixgrid

74 (array([[0],

75 [1]]), array([[2, 4]]))

76 >>> ixgrid[0].shape, ixgrid[1].shape

77 ((2, 1), (1, 2))

78 >>> a[ixgrid]

79 array([[2, 4],

80 [7, 9]])

82 >>> ixgrid = np.ix_([True, True], [2, 4])

83 >>> a[ixgrid]

84 array([[2, 4],

85 [7, 9]])

86 >>> ixgrid = np.ix_([True, True], [False, False, True, False, True])

87 >>> a[ixgrid]

88 array([[2, 4],

89 [7, 9]])

91 """

92 out = []

93 nd = len(args)

94 for k, new in enumerate(args):

95 if not isinstance(new, _nx.ndarray):

96 new = np.asarray(new)

97 if new.size == 0:

98 # Explicitly type empty arrays to avoid float default

99 new = new.astype(_nx.intp)

100 if new.ndim != 1:

101 raise ValueError("Cross index must be 1 dimensional")

102 if issubdtype(new.dtype, _nx.bool):

103 new, = new.nonzero()

104 new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1))

105 out.append(new)

106 return tuple(out)

107

108

109class nd_grid:

110 """

111 Construct a multi-dimensional "meshgrid".

112

113 ``grid = nd_grid()`` creates an instance which will return a mesh-grid

114 when indexed. The dimension and number of the output arrays are equal

115 to the number of indexing dimensions. If the step length is not a

116 complex number, then the stop is not inclusive.

117

118 However, if the step length is a **complex number** (e.g. 5j), then the

119 integer part of its magnitude is interpreted as specifying the

120 number of points to create between the start and stop values, where

121 the stop value **is inclusive**.

122

123 If instantiated with an argument of ``sparse=True``, the mesh-grid is

124 open (or not fleshed out) so that only one-dimension of each returned

125 argument is greater than 1.

126

127 Parameters

128 ----------

129 sparse : bool, optional

130 Whether the grid is sparse or not. Default is False.

131

132 Notes

133 -----

134 Two instances of `nd_grid` are made available in the NumPy namespace,

135 `mgrid` and `ogrid`, approximately defined as::

136

137 mgrid = nd_grid(sparse=False)

138 ogrid = nd_grid(sparse=True)

139

140 Users should use these pre-defined instances instead of using `nd_grid`

141 directly.

142 """

143

144 def __init__(self, sparse=False):

145 self.sparse = sparse

146

147 def __getitem__(self, key):

148 try:

149 size = []

150 # Mimic the behavior of `np.arange` and use a data type

151 # which is at least as large as `np.int_`

152 num_list = [0]

153 for k in range(len(key)):

154 step = key[k].step

155 start = key[k].start

156 stop = key[k].stop

157 if start is None:

158 start = 0

159 if step is None:

160 step = 1

161 if isinstance(step, (_nx.complexfloating, complex)):

162 step = abs(step)

163 size.append(int(step))

164 else:

165 size.append(

166 int(math.ceil((stop - start) / (step*1.0))))

167 num_list += [start, stop, step]

168 typ = _nx.result_type(*num_list)

169 if self.sparse:

170 nn = [_nx.arange(_x, dtype=_t)

171 for _x, _t in zip(size, (typ,)*len(size))]

172 else:

173 nn = _nx.indices(size, typ)

174 for k, kk in enumerate(key):

175 step = kk.step

176 start = kk.start

177 if start is None:

178 start = 0

179 if step is None:

180 step = 1

181 if isinstance(step, (_nx.complexfloating, complex)):

182 step = int(abs(step))

183 if step != 1:

184 step = (kk.stop - start) / float(step - 1)

185 nn[k] = (nn[k]*step+start)

186 if self.sparse:

187 slobj = [_nx.newaxis]*len(size)

188 for k in range(len(size)):

189 slobj[k] = slice(None, None)

190 nn[k] = nn[k][tuple(slobj)]

191 slobj[k] = _nx.newaxis

192 return tuple(nn) # ogrid -> tuple of arrays

193 return nn # mgrid -> ndarray

194 except (IndexError, TypeError):

195 step = key.step

196 stop = key.stop

197 start = key.start

198 if start is None:

199 start = 0

200 if isinstance(step, (_nx.complexfloating, complex)):

201 # Prevent the (potential) creation of integer arrays

202 step_float = abs(step)

203 step = length = int(step_float)

204 if step != 1:

205 step = (key.stop-start)/float(step-1)

206 typ = _nx.result_type(start, stop, step_float)

207 return _nx.arange(0, length, 1, dtype=typ)*step + start

208 else:

209 return _nx.arange(start, stop, step)

210

211

212class MGridClass(nd_grid):

213 """

214 An instance which returns a dense multi-dimensional "meshgrid".

215

216 An instance which returns a dense (or fleshed out) mesh-grid

217 when indexed, so that each returned argument has the same shape.

218 The dimensions and number of the output arrays are equal to the

219 number of indexing dimensions. If the step length is not a complex

220 number, then the stop is not inclusive.

221

222 However, if the step length is a **complex number** (e.g. 5j), then

223 the integer part of its magnitude is interpreted as specifying the

224 number of points to create between the start and stop values, where

225 the stop value **is inclusive**.

226

227 Returns

228 -------

229 mesh-grid : ndarray

230 A single array, containing a set of `ndarray`\\ s all of the same

231 dimensions. stacked along the first axis.

232

233 See Also

234 --------

235 ogrid : like `mgrid` but returns open (not fleshed out) mesh grids

236 meshgrid: return coordinate matrices from coordinate vectors

237 r_ : array concatenator

238 :ref:`how-to-partition`

239

240 Examples

241 --------

242 >>> np.mgrid[0:5, 0:5]

243 array([[[0, 0, 0, 0, 0],

244 [1, 1, 1, 1, 1],

245 [2, 2, 2, 2, 2],

246 [3, 3, 3, 3, 3],

247 [4, 4, 4, 4, 4]],

248 [[0, 1, 2, 3, 4],

249 [0, 1, 2, 3, 4],

250 [0, 1, 2, 3, 4],

251 [0, 1, 2, 3, 4],

252 [0, 1, 2, 3, 4]]])

253 >>> np.mgrid[-1:1:5j]

254 array([-1. , -0.5, 0. , 0.5, 1. ])

255

256 >>> np.mgrid[0:4].shape

257 (4,)

258 >>> np.mgrid[0:4, 0:5].shape

259 (2, 4, 5)

260 >>> np.mgrid[0:4, 0:5, 0:6].shape

261 (3, 4, 5, 6)

262

263 """

264

265 def __init__(self):

266 super().__init__(sparse=False)

267

268

269mgrid = MGridClass()

270

271

272class OGridClass(nd_grid):

273 """

274 An instance which returns an open multi-dimensional "meshgrid".

275

276 An instance which returns an open (i.e. not fleshed out) mesh-grid

277 when indexed, so that only one dimension of each returned array is

278 greater than 1. The dimension and number of the output arrays are

279 equal to the number of indexing dimensions. If the step length is

280 not a complex number, then the stop is not inclusive.

281

282 However, if the step length is a **complex number** (e.g. 5j), then

283 the integer part of its magnitude is interpreted as specifying the

284 number of points to create between the start and stop values, where

285 the stop value **is inclusive**.

286

287 Returns

288 -------

289 mesh-grid : ndarray or tuple of ndarrays

290 If the input is a single slice, returns an array.

291 If the input is multiple slices, returns a tuple of arrays, with

292 only one dimension not equal to 1.

293

294 See Also

295 --------

296 mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids

297 meshgrid: return coordinate matrices from coordinate vectors

298 r_ : array concatenator

299 :ref:`how-to-partition`

300

301 Examples

302 --------

303 >>> from numpy import ogrid

304 >>> ogrid[-1:1:5j]

305 array([-1. , -0.5, 0. , 0.5, 1. ])

306 >>> ogrid[0:5, 0:5]

307 (array([[0],

308 [1],

309 [2],

310 [3],

311 [4]]),

312 array([[0, 1, 2, 3, 4]]))

313

314 """

315

316 def __init__(self):

317 super().__init__(sparse=True)

318

319

320ogrid = OGridClass()

321

322

323class AxisConcatenator:

324 """

325 Translates slice objects to concatenation along an axis.

326

327 For detailed documentation on usage, see `r_`.

328 """

329 # allow ma.mr_ to override this

330 concatenate = staticmethod(_nx.concatenate)

331 makemat = staticmethod(matrixlib.matrix)

332

333 def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):

334 self.axis = axis

335 self.matrix = matrix

336 self.trans1d = trans1d

337 self.ndmin = ndmin

338

339 def __getitem__(self, key):

340 # handle matrix builder syntax

341 if isinstance(key, str):

342 frame = sys._getframe().f_back

343 mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals)

344 return mymat

345

346 if not isinstance(key, tuple):

347 key = (key,)

348

349 # copy attributes, since they can be overridden in the first argument

350 trans1d = self.trans1d

351 ndmin = self.ndmin

352 matrix = self.matrix

353 axis = self.axis

354

355 objs = []

356 # dtypes or scalars for weak scalar handling in result_type

357 result_type_objs = []

358

359 for k, item in enumerate(key):

360 scalar = False

361 if isinstance(item, slice):

362 step = item.step

363 start = item.start

364 stop = item.stop

365 if start is None:

366 start = 0

367 if step is None:

368 step = 1

369 if isinstance(step, (_nx.complexfloating, complex)):

370 size = int(abs(step))

371 newobj = linspace(start, stop, num=size)

372 else:

373 newobj = _nx.arange(start, stop, step)

374 if ndmin > 1:

375 newobj = array(newobj, copy=None, ndmin=ndmin)

376 if trans1d != -1:

377 newobj = newobj.swapaxes(-1, trans1d)

378 elif isinstance(item, str):

379 if k != 0:

380 raise ValueError("special directives must be the "

381 "first entry.")

382 if item in ('r', 'c'):

383 matrix = True

384 col = (item == 'c')

385 continue

386 if ',' in item:

387 vec = item.split(',')

388 try:

389 axis, ndmin = [int(x) for x in vec[:2]]

390 if len(vec) == 3:

391 trans1d = int(vec[2])

392 continue

393 except Exception as e:

394 raise ValueError(

395 "unknown special directive {!r}".format(item)

396 ) from e

397 try:

398 axis = int(item)

399 continue

400 except (ValueError, TypeError) as e:

401 raise ValueError("unknown special directive") from e

402 elif type(item) in ScalarType:

403 scalar = True

404 newobj = item

405 else:

406 item_ndim = np.ndim(item)

407 newobj = array(item, copy=None, subok=True, ndmin=ndmin)

408 if trans1d != -1 and item_ndim < ndmin:

409 k2 = ndmin - item_ndim

410 k1 = trans1d

411 if k1 < 0:

412 k1 += k2 + 1

413 defaxes = list(range(ndmin))

414 axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]

415 newobj = newobj.transpose(axes)

416

417 objs.append(newobj)

418 if scalar:

419 result_type_objs.append(item)

420 else:

421 result_type_objs.append(newobj.dtype)

422

423 # Ensure that scalars won't up-cast unless warranted, for 0, drops

424 # through to error in concatenate.

425 if len(result_type_objs) != 0:

426 final_dtype = _nx.result_type(*result_type_objs)

427 # concatenate could do cast, but that can be overridden:

428 objs = [array(obj, copy=None, subok=True,

429 ndmin=ndmin, dtype=final_dtype) for obj in objs]

430

431 res = self.concatenate(tuple(objs), axis=axis)

432

433 if matrix:

434 oldndim = res.ndim

435 res = self.makemat(res)

436 if oldndim == 1 and col:

437 res = res.T

438 return res

439

440 def __len__(self):

441 return 0

442

443# separate classes are used here instead of just making r_ = concatentor(0),

444# etc. because otherwise we couldn't get the doc string to come out right

445# in help(r_)

446

447

448class RClass(AxisConcatenator):

449 """

450 Translates slice objects to concatenation along the first axis.

451

452 This is a simple way to build up arrays quickly. There are two use cases.

453

454 1. If the index expression contains comma separated arrays, then stack

455 them along their first axis.

456 2. If the index expression contains slice notation or scalars then create

457 a 1-D array with a range indicated by the slice notation.

458

459 If slice notation is used, the syntax ``start:stop:step`` is equivalent

460 to ``np.arange(start, stop, step)`` inside of the brackets. However, if

461 ``step`` is an imaginary number (i.e. 100j) then its integer portion is

462 interpreted as a number-of-points desired and the start and stop are

463 inclusive. In other words ``start:stop:stepj`` is interpreted as

464 ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.

465 After expansion of slice notation, all comma separated sequences are

466 concatenated together.

467

468 Optional character strings placed as the first element of the index

469 expression can be used to change the output. The strings 'r' or 'c' result

470 in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)

471 matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1

472 (column) matrix is produced. If the result is 2-D then both provide the

473 same matrix result.

474

475 A string integer specifies which axis to stack multiple comma separated

476 arrays along. A string of two comma-separated integers allows indication

477 of the minimum number of dimensions to force each entry into as the

478 second integer (the axis to concatenate along is still the first integer).

479

480 A string with three comma-separated integers allows specification of the

481 axis to concatenate along, the minimum number of dimensions to force the

482 entries to, and which axis should contain the start of the arrays which

483 are less than the specified number of dimensions. In other words the third

484 integer allows you to specify where the 1's should be placed in the shape

485 of the arrays that have their shapes upgraded. By default, they are placed

486 in the front of the shape tuple. The third argument allows you to specify

487 where the start of the array should be instead. Thus, a third argument of

488 '0' would place the 1's at the end of the array shape. Negative integers

489 specify where in the new shape tuple the last dimension of upgraded arrays

490 should be placed, so the default is '-1'.

491

492 Parameters

493 ----------

494 Not a function, so takes no parameters

495

496

497 Returns

498 -------

499 A concatenated ndarray or matrix.

500

501 See Also

502 --------

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

504 c_ : Translates slice objects to concatenation along the second axis.

505

506 Examples

507 --------

508 >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]

509 array([1, 2, 3, ..., 4, 5, 6])

510 >>> np.r_[-1:1:6j, [0]*3, 5, 6]

511 array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])

512

513 String integers specify the axis to concatenate along or the minimum

514 number of dimensions to force entries into.

515

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

517 >>> np.r_['-1', a, a] # concatenate along last axis

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

519 [3, 4, 5, 3, 4, 5]])

520 >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2

521 array([[1, 2, 3],

522 [4, 5, 6]])

523

524 >>> np.r_['0,2,0', [1,2,3], [4,5,6]]

525 array([[1],

526 [2],

527 [3],

528 [4],

529 [5],

530 [6]])

531 >>> np.r_['1,2,0', [1,2,3], [4,5,6]]

532 array([[1, 4],

533 [2, 5],

534 [3, 6]])

535

536 Using 'r' or 'c' as a first string argument creates a matrix.

537

538 >>> np.r_['r',[1,2,3], [4,5,6]]

539 matrix([[1, 2, 3, 4, 5, 6]])

540

541 """

542

543 def __init__(self):

544 AxisConcatenator.__init__(self, 0)

547r_ = RClass()

550class CClass(AxisConcatenator):

551 """

552 Translates slice objects to concatenation along the second axis.

553

554 This is short-hand for ``np.r_['-1,2,0', index expression]``, which is

555 useful because of its common occurrence. In particular, arrays will be

556 stacked along their last axis after being upgraded to at least 2-D with

557 1's post-pended to the shape (column vectors made out of 1-D arrays).

558

559 See Also

560 --------

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

562 r_ : For more detailed documentation.

563

564 Examples

565 --------

566 >>> np.c_[np.array([1,2,3]), np.array([4,5,6])]

567 array([[1, 4],

568 [2, 5],

569 [3, 6]])

570 >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]

571 array([[1, 2, 3, ..., 4, 5, 6]])

572

573 """

574

575 def __init__(self):

576 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)

579c_ = CClass()

582@set_module('numpy')

583class ndenumerate:

584 """

585 Multidimensional index iterator.

586

587 Return an iterator yielding pairs of array coordinates and values.

588

589 Parameters

590 ----------

591 arr : ndarray

592 Input array.

593

594 See Also

595 --------

596 ndindex, flatiter

597

598 Examples

599 --------

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

601 >>> for index, x in np.ndenumerate(a):

602 ... print(index, x)

603 (0, 0) 1

604 (0, 1) 2

605 (1, 0) 3

606 (1, 1) 4

607

608 """

609

610 def __init__(self, arr):

611 self.iter = np.asarray(arr).flat

612

613 def __next__(self):

614 """

615 Standard iterator method, returns the index tuple and array value.

616

617 Returns

618 -------

619 coords : tuple of ints

620 The indices of the current iteration.

621 val : scalar

622 The array element of the current iteration.

623

624 """

625 return self.iter.coords, next(self.iter)

626

627 def __iter__(self):

628 return self

629

630

631@set_module('numpy')

632class ndindex:

633 """

634 An N-dimensional iterator object to index arrays.

635

636 Given the shape of an array, an `ndindex` instance iterates over

637 the N-dimensional index of the array. At each iteration a tuple

638 of indices is returned, the last dimension is iterated over first.

639

640 Parameters

641 ----------

642 shape : ints, or a single tuple of ints

643 The size of each dimension of the array can be passed as

644 individual parameters or as the elements of a tuple.

645

646 See Also

647 --------

648 ndenumerate, flatiter

649

650 Examples

651 --------

652 Dimensions as individual arguments

653

654 >>> for index in np.ndindex(3, 2, 1):

655 ... print(index)

656 (0, 0, 0)

657 (0, 1, 0)

658 (1, 0, 0)

659 (1, 1, 0)

660 (2, 0, 0)

661 (2, 1, 0)

662

663 Same dimensions - but in a tuple ``(3, 2, 1)``

664

665 >>> for index in np.ndindex((3, 2, 1)):

666 ... print(index)

667 (0, 0, 0)

668 (0, 1, 0)

669 (1, 0, 0)

670 (1, 1, 0)

671 (2, 0, 0)

672 (2, 1, 0)

673

674 """

675

676 def __init__(self, *shape):

677 if len(shape) == 1 and isinstance(shape[0], tuple):

678 shape = shape[0]

679 x = as_strided(_nx.zeros(1), shape=shape,

680 strides=_nx.zeros_like(shape))

681 self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],

682 order='C')

683

684 def __iter__(self):

685 return self

686

687 def ndincr(self):

688 """

689 Increment the multi-dimensional index by one.

690

691 This method is for backward compatibility only: do not use.

692

693 .. deprecated:: 1.20.0

694 This method has been advised against since numpy 1.8.0, but only

695 started emitting DeprecationWarning as of this version.

696 """

697 # NumPy 1.20.0, 2020-09-08

698 warnings.warn(

699 "`ndindex.ndincr()` is deprecated, use `next(ndindex)` instead",

700 DeprecationWarning, stacklevel=2)

701 next(self)

702

703 def __next__(self):

704 """

705 Standard iterator method, updates the index and returns the index

706 tuple.

707

708 Returns

709 -------

710 val : tuple of ints

711 Returns a tuple containing the indices of the current

712 iteration.

713

714 """

715 next(self._it)

716 return self._it.multi_index

717

718

719# You can do all this with slice() plus a few special objects,

720# but there's a lot to remember. This version is simpler because

721# it uses the standard array indexing syntax.

722#

723# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>

724# last revision: 1999-7-23

725#

726# Cosmetic changes by T. Oliphant 2001

727#

728#

729

730class IndexExpression:

731 """

732 A nicer way to build up index tuples for arrays.

733

734 .. note::

735 Use one of the two predefined instances ``index_exp`` or `s_`

736 rather than directly using `IndexExpression`.

737

738 For any index combination, including slicing and axis insertion,

739 ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any

740 array `a`. However, ``np.index_exp[indices]`` can be used anywhere

741 in Python code and returns a tuple of slice objects that can be

742 used in the construction of complex index expressions.

743

744 Parameters

745 ----------

746 maketuple : bool

747 If True, always returns a tuple.

748

749 See Also

750 --------

751 s_ : Predefined instance without tuple conversion:

752 `s_ = IndexExpression(maketuple=False)`.

753 The ``index_exp`` is another predefined instance that

754 always returns a tuple:

755 `index_exp = IndexExpression(maketuple=True)`.

756

757 Notes

758 -----

759 You can do all this with :class:`slice` plus a few special objects,

760 but there's a lot to remember and this version is simpler because

761 it uses the standard array indexing syntax.

762

763 Examples

764 --------

765 >>> np.s_[2::2]

766 slice(2, None, 2)

767 >>> np.index_exp[2::2]

768 (slice(2, None, 2),)

769

770 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]

771 array([2, 4])

772

773 """

774

775 def __init__(self, maketuple):

776 self.maketuple = maketuple

777

778 def __getitem__(self, item):

779 if self.maketuple and not isinstance(item, tuple):

780 return (item,)

781 else:

782 return item

783

784

785index_exp = IndexExpression(maketuple=True)

786s_ = IndexExpression(maketuple=False)

787

788# End contribution from Konrad.

789

790

791# The following functions complement those in twodim_base, but are

792# applicable to N-dimensions.

793

794

795def _fill_diagonal_dispatcher(a, val, wrap=None):

796 return (a,)

797

798

799@array_function_dispatch(_fill_diagonal_dispatcher)

800def fill_diagonal(a, val, wrap=False):

801 """Fill the main diagonal of the given array of any dimensionality.

802

803 For an array `a` with ``a.ndim >= 2``, the diagonal is the list of

804 values ``a[i, ..., i]`` with indices ``i`` all identical. This function

805 modifies the input array in-place without returning a value.

806

807 Parameters

808 ----------

809 a : array, at least 2-D.

810 Array whose diagonal is to be filled in-place.

811 val : scalar or array_like

812 Value(s) to write on the diagonal. If `val` is scalar, the value is

813 written along the diagonal. If array-like, the flattened `val` is

814 written along the diagonal, repeating if necessary to fill all

815 diagonal entries.

816

817 wrap : bool

818 For tall matrices in NumPy version up to 1.6.2, the

819 diagonal "wrapped" after N columns. You can have this behavior

820 with this option. This affects only tall matrices.

821

822 See also

823 --------

824 diag_indices, diag_indices_from

825

826 Notes

827 -----

828 .. versionadded:: 1.4.0

829

830 This functionality can be obtained via `diag_indices`, but internally

831 this version uses a much faster implementation that never constructs the

832 indices and uses simple slicing.

833

834 Examples

835 --------

836 >>> a = np.zeros((3, 3), int)

837 >>> np.fill_diagonal(a, 5)

838 >>> a

839 array([[5, 0, 0],

840 [0, 5, 0],

841 [0, 0, 5]])

842

843 The same function can operate on a 4-D array:

844

845 >>> a = np.zeros((3, 3, 3, 3), int)

846 >>> np.fill_diagonal(a, 4)

847

848 We only show a few blocks for clarity:

849

850 >>> a[0, 0]

851 array([[4, 0, 0],

852 [0, 0, 0],

853 [0, 0, 0]])

854 >>> a[1, 1]

855 array([[0, 0, 0],

856 [0, 4, 0],

857 [0, 0, 0]])

858 >>> a[2, 2]

859 array([[0, 0, 0],

860 [0, 0, 0],

861 [0, 0, 4]])

862

863 The wrap option affects only tall matrices:

864

865 >>> # tall matrices no wrap

866 >>> a = np.zeros((5, 3), int)

867 >>> np.fill_diagonal(a, 4)

868 >>> a

869 array([[4, 0, 0],

870 [0, 4, 0],

871 [0, 0, 4],

872 [0, 0, 0],

873 [0, 0, 0]])

874

875 >>> # tall matrices wrap

876 >>> a = np.zeros((5, 3), int)

877 >>> np.fill_diagonal(a, 4, wrap=True)

878 >>> a

879 array([[4, 0, 0],

880 [0, 4, 0],

881 [0, 0, 4],

882 [0, 0, 0],

883 [4, 0, 0]])

884

885 >>> # wide matrices

886 >>> a = np.zeros((3, 5), int)

887 >>> np.fill_diagonal(a, 4, wrap=True)

888 >>> a

889 array([[4, 0, 0, 0, 0],

890 [0, 4, 0, 0, 0],

891 [0, 0, 4, 0, 0]])

892

893 The anti-diagonal can be filled by reversing the order of elements

894 using either `numpy.flipud` or `numpy.fliplr`.

895

896 >>> a = np.zeros((3, 3), int);

897 >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip

898 >>> a

899 array([[0, 0, 1],

900 [0, 2, 0],

901 [3, 0, 0]])

902 >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip

903 >>> a

904 array([[0, 0, 3],

905 [0, 2, 0],

906 [1, 0, 0]])

907

908 Note that the order in which the diagonal is filled varies depending

909 on the flip function.

910 """

911 if a.ndim < 2:

912 raise ValueError("array must be at least 2-d")

913 end = None

914 if a.ndim == 2:

915 # Explicit, fast formula for the common case. For 2-d arrays, we

916 # accept rectangular ones.

917 step = a.shape[1] + 1

918 # This is needed to don't have tall matrix have the diagonal wrap.

919 if not wrap:

920 end = a.shape[1] * a.shape[1]

921 else:

922 # For more than d=2, the strided formula is only valid for arrays with

923 # all dimensions equal, so we check first.

924 if not np.all(diff(a.shape) == 0):

925 raise ValueError("All dimensions of input must be of equal length")

926 step = 1 + (np.cumprod(a.shape[:-1])).sum()

927

928 # Write the value out into the diagonal.

929 a.flat[:end:step] = val

930

931

932@set_module('numpy')

933def diag_indices(n, ndim=2):

934 """

935 Return the indices to access the main diagonal of an array.

936

937 This returns a tuple of indices that can be used to access the main

938 diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape

939 (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for

940 ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``

941 for ``i = [0..n-1]``.

942

943 Parameters

944 ----------

945 n : int

946 The size, along each dimension, of the arrays for which the returned

947 indices can be used.

948

949 ndim : int, optional

950 The number of dimensions.

951

952 See Also

953 --------

954 diag_indices_from

955

956 Notes

957 -----

958 .. versionadded:: 1.4.0

959

960 Examples

961 --------

962 Create a set of indices to access the diagonal of a (4, 4) array:

963

964 >>> di = np.diag_indices(4)

965 >>> di

966 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))

967 >>> a = np.arange(16).reshape(4, 4)

968 >>> a

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

970 [ 4, 5, 6, 7],

971 [ 8, 9, 10, 11],

972 [12, 13, 14, 15]])

973 >>> a[di] = 100

974 >>> a

975 array([[100, 1, 2, 3],

976 [ 4, 100, 6, 7],

977 [ 8, 9, 100, 11],

978 [ 12, 13, 14, 100]])

979

980 Now, we create indices to manipulate a 3-D array:

981

982 >>> d3 = np.diag_indices(2, 3)

983 >>> d3

984 (array([0, 1]), array([0, 1]), array([0, 1]))

985

986 And use it to set the diagonal of an array of zeros to 1:

987

988 >>> a = np.zeros((2, 2, 2), dtype=int)

989 >>> a[d3] = 1

990 >>> a

991 array([[[1, 0],

992 [0, 0]],

993 [[0, 0],

994 [0, 1]]])

995

996 """

997 idx = np.arange(n)

998 return (idx,) * ndim

999

1000

1001def _diag_indices_from(arr):

1002 return (arr,)

1003

1004

1005@array_function_dispatch(_diag_indices_from)

1006def diag_indices_from(arr):

1007 """

1008 Return the indices to access the main diagonal of an n-dimensional array.

1009

1010 See `diag_indices` for full details.

1011

1012 Parameters

1013 ----------

1014 arr : array, at least 2-D

1015

1016 See Also

1017 --------

1018 diag_indices

1019

1020 Notes

1021 -----

1022 .. versionadded:: 1.4.0

1023

1024 Examples

1025 --------

1026

1027 Create a 4 by 4 array.

1028

1029 >>> a = np.arange(16).reshape(4, 4)

1030 >>> a

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

1032 [ 4, 5, 6, 7],

1033 [ 8, 9, 10, 11],

1034 [12, 13, 14, 15]])

1035

1036 Get the indices of the diagonal elements.

1037

1038 >>> di = np.diag_indices_from(a)

1039 >>> di

1040 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))

1041

1042 >>> a[di]

1043 array([ 0, 5, 10, 15])

1044

1045 This is simply syntactic sugar for diag_indices.

1046

1047 >>> np.diag_indices(a.shape[0])

1048 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))

1049

1050 """

1051

1052 if not arr.ndim >= 2:

1053 raise ValueError("input array must be at least 2-d")

1054 # For more than d=2, the strided formula is only valid for arrays with

1055 # all dimensions equal, so we check first.

1056 if not np.all(diff(arr.shape) == 0):

1057 raise ValueError("All dimensions of input must be of equal length")

1058

1059 return diag_indices(arr.shape[0], arr.ndim)

Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/numpy/lib/_index_tricks_impl.py: 29%

257 statements