Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import functools 

2import sys 

3import math 

4 

5import numpy.core.numeric as _nx 

6from numpy.core.numeric import ( 

7 asarray, ScalarType, array, alltrue, cumprod, arange, ndim 

8 ) 

9from numpy.core.numerictypes import find_common_type, issubdtype 

10 

11import numpy.matrixlib as matrixlib 

12from .function_base import diff 

13from numpy.core.multiarray import ravel_multi_index, unravel_index 

14from numpy.core.overrides import set_module 

15from numpy.core import overrides, linspace 

16from numpy.lib.stride_tricks import as_strided 

17 

18 

19array_function_dispatch = functools.partial( 

20 overrides.array_function_dispatch, module='numpy') 

21 

22 

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 ] 

28 

29 

30def _ix__dispatcher(*args): 

31 return args 

32 

33 

34@array_function_dispatch(_ix__dispatcher) 

35def ix_(*args): 

36 """ 

37 Construct an open mesh from multiple sequences. 

38 

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. 

43 

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]]]``. 

47 

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)``). 

55 

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. 

61 

62 See Also 

63 -------- 

64 ogrid, mgrid, meshgrid 

65 

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]]) 

81 

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]]) 

90 

91 """ 

92 out = [] 

93 nd = len(args) 

94 for k, new in enumerate(args): 

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

96 new = 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 

108class nd_grid: 

109 """ 

110 Construct a multi-dimensional "meshgrid". 

111 

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

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

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

115 complex number, then the stop is not inclusive. 

116 

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

118 integer part of its magnitude is interpreted as specifying the 

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

120 the stop value **is inclusive**. 

121 

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

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

124 argument is greater than 1. 

125 

126 Parameters 

127 ---------- 

128 sparse : bool, optional 

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

130 

131 Notes 

132 ----- 

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

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

135 

136 mgrid = nd_grid(sparse=False) 

137 ogrid = nd_grid(sparse=True) 

138 

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

140 directly. 

141 """ 

142 

143 def __init__(self, sparse=False): 

144 self.sparse = sparse 

145 

146 def __getitem__(self, key): 

147 try: 

148 size = [] 

149 typ = int 

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

151 step = key[k].step 

152 start = key[k].start 

153 if start is None: 

154 start = 0 

155 if step is None: 

156 step = 1 

157 if isinstance(step, complex): 

158 size.append(int(abs(step))) 

159 typ = float 

160 else: 

161 size.append( 

162 int(math.ceil((key[k].stop - start)/(step*1.0)))) 

163 if (isinstance(step, float) or 

164 isinstance(start, float) or 

165 isinstance(key[k].stop, float)): 

166 typ = float 

167 if self.sparse: 

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

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

170 else: 

171 nn = _nx.indices(size, typ) 

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

173 step = key[k].step 

174 start = key[k].start 

175 if start is None: 

176 start = 0 

177 if step is None: 

178 step = 1 

179 if isinstance(step, complex): 

180 step = int(abs(step)) 

181 if step != 1: 

182 step = (key[k].stop - start)/float(step-1) 

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

184 if self.sparse: 

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

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

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

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

189 slobj[k] = _nx.newaxis 

190 return nn 

191 except (IndexError, TypeError): 

192 step = key.step 

193 stop = key.stop 

194 start = key.start 

195 if start is None: 

196 start = 0 

197 if isinstance(step, complex): 

198 step = abs(step) 

199 length = int(step) 

200 if step != 1: 

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

202 stop = key.stop + step 

203 return _nx.arange(0, length, 1, float)*step + start 

204 else: 

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

206 

207 

208class MGridClass(nd_grid): 

209 """ 

210 `nd_grid` instance which returns a dense multi-dimensional "meshgrid". 

211 

212 An instance of `numpy.lib.index_tricks.nd_grid` which returns an dense 

213 (or fleshed out) mesh-grid when indexed, so that each returned argument 

214 has the same shape. The dimensions and number of the output arrays are 

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

216 complex number, then the stop is not inclusive. 

217 

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

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

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

221 the stop value **is inclusive**. 

222 

223 Returns 

224 ---------- 

225 mesh-grid `ndarrays` all of the same dimensions 

226 

227 See Also 

228 -------- 

229 numpy.lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects 

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

231 r_ : array concatenator 

232 

233 Examples 

234 -------- 

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

236 array([[[0, 0, 0, 0, 0], 

237 [1, 1, 1, 1, 1], 

238 [2, 2, 2, 2, 2], 

239 [3, 3, 3, 3, 3], 

240 [4, 4, 4, 4, 4]], 

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

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

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

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

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

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

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

248 

249 """ 

250 def __init__(self): 

251 super(MGridClass, self).__init__(sparse=False) 

252 

253mgrid = MGridClass() 

254 

255class OGridClass(nd_grid): 

256 """ 

257 `nd_grid` instance which returns an open multi-dimensional "meshgrid". 

258 

259 An instance of `numpy.lib.index_tricks.nd_grid` which returns an open 

260 (i.e. not fleshed out) mesh-grid when indexed, so that only one dimension 

261 of each returned array is greater than 1. The dimension and number of the 

262 output arrays are equal to the number of indexing dimensions. If the step 

263 length is not a complex number, then the stop is not inclusive. 

264 

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

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

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

268 the stop value **is inclusive**. 

269 

270 Returns 

271 ------- 

272 mesh-grid 

273 `ndarrays` with only one dimension not equal to 1 

274 

275 See Also 

276 -------- 

277 np.lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects 

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

279 r_ : array concatenator 

280 

281 Examples 

282 -------- 

283 >>> from numpy import ogrid 

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

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

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

287 [array([[0], 

288 [1], 

289 [2], 

290 [3], 

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

292 

293 """ 

294 def __init__(self): 

295 super(OGridClass, self).__init__(sparse=True) 

296 

297ogrid = OGridClass() 

298 

299 

300class AxisConcatenator: 

301 """ 

302 Translates slice objects to concatenation along an axis. 

303 

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

305 """ 

306 # allow ma.mr_ to override this 

307 concatenate = staticmethod(_nx.concatenate) 

308 makemat = staticmethod(matrixlib.matrix) 

309 

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

311 self.axis = axis 

312 self.matrix = matrix 

313 self.trans1d = trans1d 

314 self.ndmin = ndmin 

315 

316 def __getitem__(self, key): 

317 # handle matrix builder syntax 

318 if isinstance(key, str): 

319 frame = sys._getframe().f_back 

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

321 return mymat 

322 

323 if not isinstance(key, tuple): 

324 key = (key,) 

325 

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

327 trans1d = self.trans1d 

328 ndmin = self.ndmin 

329 matrix = self.matrix 

330 axis = self.axis 

331 

332 objs = [] 

333 scalars = [] 

334 arraytypes = [] 

335 scalartypes = [] 

336 

337 for k, item in enumerate(key): 

338 scalar = False 

339 if isinstance(item, slice): 

340 step = item.step 

341 start = item.start 

342 stop = item.stop 

343 if start is None: 

344 start = 0 

345 if step is None: 

346 step = 1 

347 if isinstance(step, complex): 

348 size = int(abs(step)) 

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

350 else: 

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

352 if ndmin > 1: 

353 newobj = array(newobj, copy=False, ndmin=ndmin) 

354 if trans1d != -1: 

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

356 elif isinstance(item, str): 

357 if k != 0: 

358 raise ValueError("special directives must be the " 

359 "first entry.") 

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

361 matrix = True 

362 col = (item == 'c') 

363 continue 

364 if ',' in item: 

365 vec = item.split(',') 

366 try: 

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

368 if len(vec) == 3: 

369 trans1d = int(vec[2]) 

370 continue 

371 except Exception as e: 

372 raise ValueError( 

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

374 ) from e 

375 try: 

376 axis = int(item) 

377 continue 

378 except (ValueError, TypeError): 

379 raise ValueError("unknown special directive") 

380 elif type(item) in ScalarType: 

381 newobj = array(item, ndmin=ndmin) 

382 scalars.append(len(objs)) 

383 scalar = True 

384 scalartypes.append(newobj.dtype) 

385 else: 

386 item_ndim = ndim(item) 

387 newobj = array(item, copy=False, subok=True, ndmin=ndmin) 

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

389 k2 = ndmin - item_ndim 

390 k1 = trans1d 

391 if k1 < 0: 

392 k1 += k2 + 1 

393 defaxes = list(range(ndmin)) 

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

395 newobj = newobj.transpose(axes) 

396 objs.append(newobj) 

397 if not scalar and isinstance(newobj, _nx.ndarray): 

398 arraytypes.append(newobj.dtype) 

399 

400 # Ensure that scalars won't up-cast unless warranted 

401 final_dtype = find_common_type(arraytypes, scalartypes) 

402 if final_dtype is not None: 

403 for k in scalars: 

404 objs[k] = objs[k].astype(final_dtype) 

405 

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

407 

408 if matrix: 

409 oldndim = res.ndim 

410 res = self.makemat(res) 

411 if oldndim == 1 and col: 

412 res = res.T 

413 return res 

414 

415 def __len__(self): 

416 return 0 

417 

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

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

420# in help(r_) 

421 

422class RClass(AxisConcatenator): 

423 """ 

424 Translates slice objects to concatenation along the first axis. 

425 

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

427 

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

429 them along their first axis. 

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

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

432 

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

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

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

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

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

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

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

440 concatenated together. 

441 

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

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

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

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

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

447 same matrix result. 

448 

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

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

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

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

453 

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

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

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

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

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

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

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

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

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

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

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

465 

466 Parameters 

467 ---------- 

468 Not a function, so takes no parameters 

469 

470 

471 Returns 

472 ------- 

473 A concatenated ndarray or matrix. 

474 

475 See Also 

476 -------- 

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

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

479 

480 Examples 

481 -------- 

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

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

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

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

486 

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

488 number of dimensions to force entries into. 

489 

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

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

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

493 [3, 4, 5, 3, 4, 5]]) 

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

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

496 [4, 5, 6]]) 

497 

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

499 array([[1], 

500 [2], 

501 [3], 

502 [4], 

503 [5], 

504 [6]]) 

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

506 array([[1, 4], 

507 [2, 5], 

508 [3, 6]]) 

509 

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

511 

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

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

514 

515 """ 

516 

517 def __init__(self): 

518 AxisConcatenator.__init__(self, 0) 

519 

520r_ = RClass() 

521 

522class CClass(AxisConcatenator): 

523 """ 

524 Translates slice objects to concatenation along the second axis. 

525 

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

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

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

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

530  

531 See Also 

532 -------- 

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

534 r_ : For more detailed documentation. 

535 

536 Examples 

537 -------- 

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

539 array([[1, 4], 

540 [2, 5], 

541 [3, 6]]) 

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

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

544 

545 """ 

546 

547 def __init__(self): 

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

549 

550 

551c_ = CClass() 

552 

553 

554@set_module('numpy') 

555class ndenumerate: 

556 """ 

557 Multidimensional index iterator. 

558 

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

560 

561 Parameters 

562 ---------- 

563 arr : ndarray 

564 Input array. 

565 

566 See Also 

567 -------- 

568 ndindex, flatiter 

569 

570 Examples 

571 -------- 

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

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

574 ... print(index, x) 

575 (0, 0) 1 

576 (0, 1) 2 

577 (1, 0) 3 

578 (1, 1) 4 

579 

580 """ 

581 

582 def __init__(self, arr): 

583 self.iter = asarray(arr).flat 

584 

585 def __next__(self): 

586 """ 

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

588 

589 Returns 

590 ------- 

591 coords : tuple of ints 

592 The indices of the current iteration. 

593 val : scalar 

594 The array element of the current iteration. 

595 

596 """ 

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

598 

599 def __iter__(self): 

600 return self 

601 

602 

603@set_module('numpy') 

604class ndindex: 

605 """ 

606 An N-dimensional iterator object to index arrays. 

607 

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

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

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

611 

612 Parameters 

613 ---------- 

614 `*args` : ints 

615 The size of each dimension of the array. 

616 

617 See Also 

618 -------- 

619 ndenumerate, flatiter 

620 

621 Examples 

622 -------- 

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

624 ... print(index) 

625 (0, 0, 0) 

626 (0, 1, 0) 

627 (1, 0, 0) 

628 (1, 1, 0) 

629 (2, 0, 0) 

630 (2, 1, 0) 

631 

632 """ 

633 

634 def __init__(self, *shape): 

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

636 shape = shape[0] 

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

638 strides=_nx.zeros_like(shape)) 

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

640 order='C') 

641 

642 def __iter__(self): 

643 return self 

644 

645 def ndincr(self): 

646 """ 

647 Increment the multi-dimensional index by one. 

648 

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

650 """ 

651 next(self) 

652 

653 def __next__(self): 

654 """ 

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

656 tuple. 

657 

658 Returns 

659 ------- 

660 val : tuple of ints 

661 Returns a tuple containing the indices of the current 

662 iteration. 

663 

664 """ 

665 next(self._it) 

666 return self._it.multi_index 

667 

668 

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

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

671# it uses the standard array indexing syntax. 

672# 

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

674# last revision: 1999-7-23 

675# 

676# Cosmetic changes by T. Oliphant 2001 

677# 

678# 

679 

680class IndexExpression: 

681 """ 

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

683 

684 .. note:: 

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

686 rather than directly using `IndexExpression`. 

687 

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

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

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

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

692 used in the construction of complex index expressions. 

693 

694 Parameters 

695 ---------- 

696 maketuple : bool 

697 If True, always returns a tuple. 

698 

699 See Also 

700 -------- 

701 index_exp : Predefined instance that always returns a tuple: 

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

703 s_ : Predefined instance without tuple conversion: 

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

705 

706 Notes 

707 ----- 

708 You can do all this with `slice()` plus a few special objects, 

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

710 it uses the standard array indexing syntax. 

711 

712 Examples 

713 -------- 

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

715 slice(2, None, 2) 

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

717 (slice(2, None, 2),) 

718 

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

720 array([2, 4]) 

721 

722 """ 

723 

724 def __init__(self, maketuple): 

725 self.maketuple = maketuple 

726 

727 def __getitem__(self, item): 

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

729 return (item,) 

730 else: 

731 return item 

732 

733index_exp = IndexExpression(maketuple=True) 

734s_ = IndexExpression(maketuple=False) 

735 

736# End contribution from Konrad. 

737 

738 

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

740# applicable to N-dimensions. 

741 

742 

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

744 return (a,) 

745 

746 

747@array_function_dispatch(_fill_diagonal_dispatcher) 

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

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

750 

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

752 locations with indices ``a[i, ..., i]`` all identical. This function 

753 modifies the input array in-place, it does not return a value. 

754 

755 Parameters 

756 ---------- 

757 a : array, at least 2-D. 

758 Array whose diagonal is to be filled, it gets modified in-place. 

759 

760 val : scalar 

761 Value to be written on the diagonal, its type must be compatible with 

762 that of the array a. 

763 

764 wrap : bool 

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

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

767 with this option. This affects only tall matrices. 

768 

769 See also 

770 -------- 

771 diag_indices, diag_indices_from 

772 

773 Notes 

774 ----- 

775 .. versionadded:: 1.4.0 

776 

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

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

779 indices and uses simple slicing. 

780 

781 Examples 

782 -------- 

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

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

785 >>> a 

786 array([[5, 0, 0], 

787 [0, 5, 0], 

788 [0, 0, 5]]) 

789 

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

791 

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

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

794 

795 We only show a few blocks for clarity: 

796 

797 >>> a[0, 0] 

798 array([[4, 0, 0], 

799 [0, 0, 0], 

800 [0, 0, 0]]) 

801 >>> a[1, 1] 

802 array([[0, 0, 0], 

803 [0, 4, 0], 

804 [0, 0, 0]]) 

805 >>> a[2, 2] 

806 array([[0, 0, 0], 

807 [0, 0, 0], 

808 [0, 0, 4]]) 

809 

810 The wrap option affects only tall matrices: 

811 

812 >>> # tall matrices no wrap 

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

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

815 >>> a 

816 array([[4, 0, 0], 

817 [0, 4, 0], 

818 [0, 0, 4], 

819 [0, 0, 0], 

820 [0, 0, 0]]) 

821 

822 >>> # tall matrices wrap 

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

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

825 >>> a 

826 array([[4, 0, 0], 

827 [0, 4, 0], 

828 [0, 0, 4], 

829 [0, 0, 0], 

830 [4, 0, 0]]) 

831 

832 >>> # wide matrices 

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

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

835 >>> a 

836 array([[4, 0, 0, 0, 0], 

837 [0, 4, 0, 0, 0], 

838 [0, 0, 4, 0, 0]]) 

839 

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

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

842 

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

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

845 >>> a 

846 array([[0, 0, 1], 

847 [0, 2, 0], 

848 [3, 0, 0]]) 

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

850 >>> a 

851 array([[0, 0, 3], 

852 [0, 2, 0], 

853 [1, 0, 0]]) 

854 

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

856 on the flip function. 

857 """ 

858 if a.ndim < 2: 

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

860 end = None 

861 if a.ndim == 2: 

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

863 # accept rectangular ones. 

864 step = a.shape[1] + 1 

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

866 if not wrap: 

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

868 else: 

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

870 # all dimensions equal, so we check first. 

871 if not alltrue(diff(a.shape) == 0): 

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

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

874 

875 # Write the value out into the diagonal. 

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

877 

878 

879@set_module('numpy') 

880def diag_indices(n, ndim=2): 

881 """ 

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

883 

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

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

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

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

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

889 

890 Parameters 

891 ---------- 

892 n : int 

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

894 indices can be used. 

895 

896 ndim : int, optional 

897 The number of dimensions. 

898 

899 See also 

900 -------- 

901 diag_indices_from 

902 

903 Notes 

904 ----- 

905 .. versionadded:: 1.4.0 

906 

907 Examples 

908 -------- 

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

910 

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

912 >>> di 

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

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

915 >>> a 

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

917 [ 4, 5, 6, 7], 

918 [ 8, 9, 10, 11], 

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

920 >>> a[di] = 100 

921 >>> a 

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

923 [ 4, 100, 6, 7], 

924 [ 8, 9, 100, 11], 

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

926 

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

928 

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

930 >>> d3 

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

932 

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

934 

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

936 >>> a[d3] = 1 

937 >>> a 

938 array([[[1, 0], 

939 [0, 0]], 

940 [[0, 0], 

941 [0, 1]]]) 

942 

943 """ 

944 idx = arange(n) 

945 return (idx,) * ndim 

946 

947 

948def _diag_indices_from(arr): 

949 return (arr,) 

950 

951 

952@array_function_dispatch(_diag_indices_from) 

953def diag_indices_from(arr): 

954 """ 

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

956 

957 See `diag_indices` for full details. 

958 

959 Parameters 

960 ---------- 

961 arr : array, at least 2-D 

962 

963 See Also 

964 -------- 

965 diag_indices 

966 

967 Notes 

968 ----- 

969 .. versionadded:: 1.4.0 

970 

971 """ 

972 

973 if not arr.ndim >= 2: 

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

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

976 # all dimensions equal, so we check first. 

977 if not alltrue(diff(arr.shape) == 0): 

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

979 

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