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

1""" Basic functions for manipulating 2d arrays 

2 

3""" 

4import functools 

5 

6from numpy.core.numeric import ( 

7 asanyarray, arange, zeros, greater_equal, multiply, ones, 

8 asarray, where, int8, int16, int32, int64, empty, promote_types, diagonal, 

9 nonzero 

10 ) 

11from numpy.core.overrides import set_module 

12from numpy.core import overrides 

13from numpy.core import iinfo 

14 

15 

16__all__ = [ 

17 'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu', 

18 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices', 

19 'tril_indices_from', 'triu_indices', 'triu_indices_from', ] 

20 

21 

22array_function_dispatch = functools.partial( 

23 overrides.array_function_dispatch, module='numpy') 

24 

25 

26i1 = iinfo(int8) 

27i2 = iinfo(int16) 

28i4 = iinfo(int32) 

29 

30 

31def _min_int(low, high): 

32 """ get small int that fits the range """ 

33 if high <= i1.max and low >= i1.min: 

34 return int8 

35 if high <= i2.max and low >= i2.min: 

36 return int16 

37 if high <= i4.max and low >= i4.min: 

38 return int32 

39 return int64 

40 

41 

42def _flip_dispatcher(m): 

43 return (m,) 

44 

45 

46@array_function_dispatch(_flip_dispatcher) 

47def fliplr(m): 

48 """ 

49 Flip array in the left/right direction. 

50 

51 Flip the entries in each row in the left/right direction. 

52 Columns are preserved, but appear in a different order than before. 

53 

54 Parameters 

55 ---------- 

56 m : array_like 

57 Input array, must be at least 2-D. 

58 

59 Returns 

60 ------- 

61 f : ndarray 

62 A view of `m` with the columns reversed. Since a view 

63 is returned, this operation is :math:`\\mathcal O(1)`. 

64 

65 See Also 

66 -------- 

67 flipud : Flip array in the up/down direction. 

68 rot90 : Rotate array counterclockwise. 

69 

70 Notes 

71 ----- 

72 Equivalent to m[:,::-1]. Requires the array to be at least 2-D. 

73 

74 Examples 

75 -------- 

76 >>> A = np.diag([1.,2.,3.]) 

77 >>> A 

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

79 [0., 2., 0.], 

80 [0., 0., 3.]]) 

81 >>> np.fliplr(A) 

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

83 [0., 2., 0.], 

84 [3., 0., 0.]]) 

85 

86 >>> A = np.random.randn(2,3,5) 

87 >>> np.all(np.fliplr(A) == A[:,::-1,...]) 

88 True 

89 

90 """ 

91 m = asanyarray(m) 

92 if m.ndim < 2: 

93 raise ValueError("Input must be >= 2-d.") 

94 return m[:, ::-1] 

95 

96 

97@array_function_dispatch(_flip_dispatcher) 

98def flipud(m): 

99 """ 

100 Flip array in the up/down direction. 

101 

102 Flip the entries in each column in the up/down direction. 

103 Rows are preserved, but appear in a different order than before. 

104 

105 Parameters 

106 ---------- 

107 m : array_like 

108 Input array. 

109 

110 Returns 

111 ------- 

112 out : array_like 

113 A view of `m` with the rows reversed. Since a view is 

114 returned, this operation is :math:`\\mathcal O(1)`. 

115 

116 See Also 

117 -------- 

118 fliplr : Flip array in the left/right direction. 

119 rot90 : Rotate array counterclockwise. 

120 

121 Notes 

122 ----- 

123 Equivalent to ``m[::-1,...]``. 

124 Does not require the array to be two-dimensional. 

125 

126 Examples 

127 -------- 

128 >>> A = np.diag([1.0, 2, 3]) 

129 >>> A 

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

131 [0., 2., 0.], 

132 [0., 0., 3.]]) 

133 >>> np.flipud(A) 

134 array([[0., 0., 3.], 

135 [0., 2., 0.], 

136 [1., 0., 0.]]) 

137 

138 >>> A = np.random.randn(2,3,5) 

139 >>> np.all(np.flipud(A) == A[::-1,...]) 

140 True 

141 

142 >>> np.flipud([1,2]) 

143 array([2, 1]) 

144 

145 """ 

146 m = asanyarray(m) 

147 if m.ndim < 1: 

148 raise ValueError("Input must be >= 1-d.") 

149 return m[::-1, ...] 

150 

151 

152@set_module('numpy') 

153def eye(N, M=None, k=0, dtype=float, order='C'): 

154 """ 

155 Return a 2-D array with ones on the diagonal and zeros elsewhere. 

156 

157 Parameters 

158 ---------- 

159 N : int 

160 Number of rows in the output. 

161 M : int, optional 

162 Number of columns in the output. If None, defaults to `N`. 

163 k : int, optional 

164 Index of the diagonal: 0 (the default) refers to the main diagonal, 

165 a positive value refers to an upper diagonal, and a negative value 

166 to a lower diagonal. 

167 dtype : data-type, optional 

168 Data-type of the returned array. 

169 order : {'C', 'F'}, optional 

170 Whether the output should be stored in row-major (C-style) or 

171 column-major (Fortran-style) order in memory. 

172 

173 .. versionadded:: 1.14.0 

174 

175 Returns 

176 ------- 

177 I : ndarray of shape (N,M) 

178 An array where all elements are equal to zero, except for the `k`-th 

179 diagonal, whose values are equal to one. 

180 

181 See Also 

182 -------- 

183 identity : (almost) equivalent function 

184 diag : diagonal 2-D array from a 1-D array specified by the user. 

185 

186 Examples 

187 -------- 

188 >>> np.eye(2, dtype=int) 

189 array([[1, 0], 

190 [0, 1]]) 

191 >>> np.eye(3, k=1) 

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

193 [0., 0., 1.], 

194 [0., 0., 0.]]) 

195 

196 """ 

197 if M is None: 

198 M = N 

199 m = zeros((N, M), dtype=dtype, order=order) 

200 if k >= M: 

201 return m 

202 if k >= 0: 

203 i = k 

204 else: 

205 i = (-k) * M 

206 m[:M-k].flat[i::M+1] = 1 

207 return m 

208 

209 

210def _diag_dispatcher(v, k=None): 

211 return (v,) 

212 

213 

214@array_function_dispatch(_diag_dispatcher) 

215def diag(v, k=0): 

216 """ 

217 Extract a diagonal or construct a diagonal array. 

218 

219 See the more detailed documentation for ``numpy.diagonal`` if you use this 

220 function to extract a diagonal and wish to write to the resulting array; 

221 whether it returns a copy or a view depends on what version of numpy you 

222 are using. 

223 

224 Parameters 

225 ---------- 

226 v : array_like 

227 If `v` is a 2-D array, return a copy of its `k`-th diagonal. 

228 If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th 

229 diagonal. 

230 k : int, optional 

231 Diagonal in question. The default is 0. Use `k>0` for diagonals 

232 above the main diagonal, and `k<0` for diagonals below the main 

233 diagonal. 

234 

235 Returns 

236 ------- 

237 out : ndarray 

238 The extracted diagonal or constructed diagonal array. 

239 

240 See Also 

241 -------- 

242 diagonal : Return specified diagonals. 

243 diagflat : Create a 2-D array with the flattened input as a diagonal. 

244 trace : Sum along diagonals. 

245 triu : Upper triangle of an array. 

246 tril : Lower triangle of an array. 

247 

248 Examples 

249 -------- 

250 >>> x = np.arange(9).reshape((3,3)) 

251 >>> x 

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

253 [3, 4, 5], 

254 [6, 7, 8]]) 

255 

256 >>> np.diag(x) 

257 array([0, 4, 8]) 

258 >>> np.diag(x, k=1) 

259 array([1, 5]) 

260 >>> np.diag(x, k=-1) 

261 array([3, 7]) 

262 

263 >>> np.diag(np.diag(x)) 

264 array([[0, 0, 0], 

265 [0, 4, 0], 

266 [0, 0, 8]]) 

267 

268 """ 

269 v = asanyarray(v) 

270 s = v.shape 

271 if len(s) == 1: 

272 n = s[0]+abs(k) 

273 res = zeros((n, n), v.dtype) 

274 if k >= 0: 

275 i = k 

276 else: 

277 i = (-k) * n 

278 res[:n-k].flat[i::n+1] = v 

279 return res 

280 elif len(s) == 2: 

281 return diagonal(v, k) 

282 else: 

283 raise ValueError("Input must be 1- or 2-d.") 

284 

285 

286@array_function_dispatch(_diag_dispatcher) 

287def diagflat(v, k=0): 

288 """ 

289 Create a two-dimensional array with the flattened input as a diagonal. 

290 

291 Parameters 

292 ---------- 

293 v : array_like 

294 Input data, which is flattened and set as the `k`-th 

295 diagonal of the output. 

296 k : int, optional 

297 Diagonal to set; 0, the default, corresponds to the "main" diagonal, 

298 a positive (negative) `k` giving the number of the diagonal above 

299 (below) the main. 

300 

301 Returns 

302 ------- 

303 out : ndarray 

304 The 2-D output array. 

305 

306 See Also 

307 -------- 

308 diag : MATLAB work-alike for 1-D and 2-D arrays. 

309 diagonal : Return specified diagonals. 

310 trace : Sum along diagonals. 

311 

312 Examples 

313 -------- 

314 >>> np.diagflat([[1,2], [3,4]]) 

315 array([[1, 0, 0, 0], 

316 [0, 2, 0, 0], 

317 [0, 0, 3, 0], 

318 [0, 0, 0, 4]]) 

319 

320 >>> np.diagflat([1,2], 1) 

321 array([[0, 1, 0], 

322 [0, 0, 2], 

323 [0, 0, 0]]) 

324 

325 """ 

326 try: 

327 wrap = v.__array_wrap__ 

328 except AttributeError: 

329 wrap = None 

330 v = asarray(v).ravel() 

331 s = len(v) 

332 n = s + abs(k) 

333 res = zeros((n, n), v.dtype) 

334 if (k >= 0): 

335 i = arange(0, n-k) 

336 fi = i+k+i*n 

337 else: 

338 i = arange(0, n+k) 

339 fi = i+(i-k)*n 

340 res.flat[fi] = v 

341 if not wrap: 

342 return res 

343 return wrap(res) 

344 

345 

346@set_module('numpy') 

347def tri(N, M=None, k=0, dtype=float): 

348 """ 

349 An array with ones at and below the given diagonal and zeros elsewhere. 

350 

351 Parameters 

352 ---------- 

353 N : int 

354 Number of rows in the array. 

355 M : int, optional 

356 Number of columns in the array. 

357 By default, `M` is taken equal to `N`. 

358 k : int, optional 

359 The sub-diagonal at and below which the array is filled. 

360 `k` = 0 is the main diagonal, while `k` < 0 is below it, 

361 and `k` > 0 is above. The default is 0. 

362 dtype : dtype, optional 

363 Data type of the returned array. The default is float. 

364 

365 Returns 

366 ------- 

367 tri : ndarray of shape (N, M) 

368 Array with its lower triangle filled with ones and zero elsewhere; 

369 in other words ``T[i,j] == 1`` for ``j <= i + k``, 0 otherwise. 

370 

371 Examples 

372 -------- 

373 >>> np.tri(3, 5, 2, dtype=int) 

374 array([[1, 1, 1, 0, 0], 

375 [1, 1, 1, 1, 0], 

376 [1, 1, 1, 1, 1]]) 

377 

378 >>> np.tri(3, 5, -1) 

379 array([[0., 0., 0., 0., 0.], 

380 [1., 0., 0., 0., 0.], 

381 [1., 1., 0., 0., 0.]]) 

382 

383 """ 

384 if M is None: 

385 M = N 

386 

387 m = greater_equal.outer(arange(N, dtype=_min_int(0, N)), 

388 arange(-k, M-k, dtype=_min_int(-k, M - k))) 

389 

390 # Avoid making a copy if the requested type is already bool 

391 m = m.astype(dtype, copy=False) 

392 

393 return m 

394 

395 

396def _trilu_dispatcher(m, k=None): 

397 return (m,) 

398 

399 

400@array_function_dispatch(_trilu_dispatcher) 

401def tril(m, k=0): 

402 """ 

403 Lower triangle of an array. 

404 

405 Return a copy of an array with elements above the `k`-th diagonal zeroed. 

406 

407 Parameters 

408 ---------- 

409 m : array_like, shape (M, N) 

410 Input array. 

411 k : int, optional 

412 Diagonal above which to zero elements. `k = 0` (the default) is the 

413 main diagonal, `k < 0` is below it and `k > 0` is above. 

414 

415 Returns 

416 ------- 

417 tril : ndarray, shape (M, N) 

418 Lower triangle of `m`, of same shape and data-type as `m`. 

419 

420 See Also 

421 -------- 

422 triu : same thing, only for the upper triangle 

423 

424 Examples 

425 -------- 

426 >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) 

427 array([[ 0, 0, 0], 

428 [ 4, 0, 0], 

429 [ 7, 8, 0], 

430 [10, 11, 12]]) 

431 

432 """ 

433 m = asanyarray(m) 

434 mask = tri(*m.shape[-2:], k=k, dtype=bool) 

435 

436 return where(mask, m, zeros(1, m.dtype)) 

437 

438 

439@array_function_dispatch(_trilu_dispatcher) 

440def triu(m, k=0): 

441 """ 

442 Upper triangle of an array. 

443 

444 Return a copy of a matrix with the elements below the `k`-th diagonal 

445 zeroed. 

446 

447 Please refer to the documentation for `tril` for further details. 

448 

449 See Also 

450 -------- 

451 tril : lower triangle of an array 

452 

453 Examples 

454 -------- 

455 >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) 

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

457 [ 4, 5, 6], 

458 [ 0, 8, 9], 

459 [ 0, 0, 12]]) 

460 

461 """ 

462 m = asanyarray(m) 

463 mask = tri(*m.shape[-2:], k=k-1, dtype=bool) 

464 

465 return where(mask, zeros(1, m.dtype), m) 

466 

467 

468def _vander_dispatcher(x, N=None, increasing=None): 

469 return (x,) 

470 

471 

472# Originally borrowed from John Hunter and matplotlib 

473@array_function_dispatch(_vander_dispatcher) 

474def vander(x, N=None, increasing=False): 

475 """ 

476 Generate a Vandermonde matrix. 

477 

478 The columns of the output matrix are powers of the input vector. The 

479 order of the powers is determined by the `increasing` boolean argument. 

480 Specifically, when `increasing` is False, the `i`-th output column is 

481 the input vector raised element-wise to the power of ``N - i - 1``. Such 

482 a matrix with a geometric progression in each row is named for Alexandre- 

483 Theophile Vandermonde. 

484 

485 Parameters 

486 ---------- 

487 x : array_like 

488 1-D input array. 

489 N : int, optional 

490 Number of columns in the output. If `N` is not specified, a square 

491 array is returned (``N = len(x)``). 

492 increasing : bool, optional 

493 Order of the powers of the columns. If True, the powers increase 

494 from left to right, if False (the default) they are reversed. 

495 

496 .. versionadded:: 1.9.0 

497 

498 Returns 

499 ------- 

500 out : ndarray 

501 Vandermonde matrix. If `increasing` is False, the first column is 

502 ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is 

503 True, the columns are ``x^0, x^1, ..., x^(N-1)``. 

504 

505 See Also 

506 -------- 

507 polynomial.polynomial.polyvander 

508 

509 Examples 

510 -------- 

511 >>> x = np.array([1, 2, 3, 5]) 

512 >>> N = 3 

513 >>> np.vander(x, N) 

514 array([[ 1, 1, 1], 

515 [ 4, 2, 1], 

516 [ 9, 3, 1], 

517 [25, 5, 1]]) 

518 

519 >>> np.column_stack([x**(N-1-i) for i in range(N)]) 

520 array([[ 1, 1, 1], 

521 [ 4, 2, 1], 

522 [ 9, 3, 1], 

523 [25, 5, 1]]) 

524 

525 >>> x = np.array([1, 2, 3, 5]) 

526 >>> np.vander(x) 

527 array([[ 1, 1, 1, 1], 

528 [ 8, 4, 2, 1], 

529 [ 27, 9, 3, 1], 

530 [125, 25, 5, 1]]) 

531 >>> np.vander(x, increasing=True) 

532 array([[ 1, 1, 1, 1], 

533 [ 1, 2, 4, 8], 

534 [ 1, 3, 9, 27], 

535 [ 1, 5, 25, 125]]) 

536 

537 The determinant of a square Vandermonde matrix is the product 

538 of the differences between the values of the input vector: 

539 

540 >>> np.linalg.det(np.vander(x)) 

541 48.000000000000043 # may vary 

542 >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1) 

543 48 

544 

545 """ 

546 x = asarray(x) 

547 if x.ndim != 1: 

548 raise ValueError("x must be a one-dimensional array or sequence.") 

549 if N is None: 

550 N = len(x) 

551 

552 v = empty((len(x), N), dtype=promote_types(x.dtype, int)) 

553 tmp = v[:, ::-1] if not increasing else v 

554 

555 if N > 0: 

556 tmp[:, 0] = 1 

557 if N > 1: 

558 tmp[:, 1:] = x[:, None] 

559 multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1) 

560 

561 return v 

562 

563 

564def _histogram2d_dispatcher(x, y, bins=None, range=None, normed=None, 

565 weights=None, density=None): 

566 yield x 

567 yield y 

568 

569 # This terrible logic is adapted from the checks in histogram2d 

570 try: 

571 N = len(bins) 

572 except TypeError: 

573 N = 1 

574 if N == 2: 

575 yield from bins # bins=[x, y] 

576 else: 

577 yield bins 

578 

579 yield weights 

580 

581 

582@array_function_dispatch(_histogram2d_dispatcher) 

583def histogram2d(x, y, bins=10, range=None, normed=None, weights=None, 

584 density=None): 

585 """ 

586 Compute the bi-dimensional histogram of two data samples. 

587 

588 Parameters 

589 ---------- 

590 x : array_like, shape (N,) 

591 An array containing the x coordinates of the points to be 

592 histogrammed. 

593 y : array_like, shape (N,) 

594 An array containing the y coordinates of the points to be 

595 histogrammed. 

596 bins : int or array_like or [int, int] or [array, array], optional 

597 The bin specification: 

598 

599 * If int, the number of bins for the two dimensions (nx=ny=bins). 

600 * If array_like, the bin edges for the two dimensions 

601 (x_edges=y_edges=bins). 

602 * If [int, int], the number of bins in each dimension 

603 (nx, ny = bins). 

604 * If [array, array], the bin edges in each dimension 

605 (x_edges, y_edges = bins). 

606 * A combination [int, array] or [array, int], where int 

607 is the number of bins and array is the bin edges. 

608 

609 range : array_like, shape(2,2), optional 

610 The leftmost and rightmost edges of the bins along each dimension 

611 (if not specified explicitly in the `bins` parameters): 

612 ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range 

613 will be considered outliers and not tallied in the histogram. 

614 density : bool, optional 

615 If False, the default, returns the number of samples in each bin. 

616 If True, returns the probability *density* function at the bin, 

617 ``bin_count / sample_count / bin_area``. 

618 normed : bool, optional 

619 An alias for the density argument that behaves identically. To avoid 

620 confusion with the broken normed argument to `histogram`, `density` 

621 should be preferred. 

622 weights : array_like, shape(N,), optional 

623 An array of values ``w_i`` weighing each sample ``(x_i, y_i)``. 

624 Weights are normalized to 1 if `normed` is True. If `normed` is 

625 False, the values of the returned histogram are equal to the sum of 

626 the weights belonging to the samples falling into each bin. 

627 

628 Returns 

629 ------- 

630 H : ndarray, shape(nx, ny) 

631 The bi-dimensional histogram of samples `x` and `y`. Values in `x` 

632 are histogrammed along the first dimension and values in `y` are 

633 histogrammed along the second dimension. 

634 xedges : ndarray, shape(nx+1,) 

635 The bin edges along the first dimension. 

636 yedges : ndarray, shape(ny+1,) 

637 The bin edges along the second dimension. 

638 

639 See Also 

640 -------- 

641 histogram : 1D histogram 

642 histogramdd : Multidimensional histogram 

643 

644 Notes 

645 ----- 

646 When `normed` is True, then the returned histogram is the sample 

647 density, defined such that the sum over bins of the product 

648 ``bin_value * bin_area`` is 1. 

649 

650 Please note that the histogram does not follow the Cartesian convention 

651 where `x` values are on the abscissa and `y` values on the ordinate 

652 axis. Rather, `x` is histogrammed along the first dimension of the 

653 array (vertical), and `y` along the second dimension of the array 

654 (horizontal). This ensures compatibility with `histogramdd`. 

655 

656 Examples 

657 -------- 

658 >>> from matplotlib.image import NonUniformImage 

659 >>> import matplotlib.pyplot as plt 

660 

661 Construct a 2-D histogram with variable bin width. First define the bin 

662 edges: 

663 

664 >>> xedges = [0, 1, 3, 5] 

665 >>> yedges = [0, 2, 3, 4, 6] 

666 

667 Next we create a histogram H with random bin content: 

668 

669 >>> x = np.random.normal(2, 1, 100) 

670 >>> y = np.random.normal(1, 1, 100) 

671 >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges)) 

672 >>> H = H.T # Let each row list bins with common y range. 

673 

674 :func:`imshow <matplotlib.pyplot.imshow>` can only display square bins: 

675 

676 >>> fig = plt.figure(figsize=(7, 3)) 

677 >>> ax = fig.add_subplot(131, title='imshow: square bins') 

678 >>> plt.imshow(H, interpolation='nearest', origin='low', 

679 ... extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]]) 

680 <matplotlib.image.AxesImage object at 0x...> 

681 

682 :func:`pcolormesh <matplotlib.pyplot.pcolormesh>` can display actual edges: 

683 

684 >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges', 

685 ... aspect='equal') 

686 >>> X, Y = np.meshgrid(xedges, yedges) 

687 >>> ax.pcolormesh(X, Y, H) 

688 <matplotlib.collections.QuadMesh object at 0x...> 

689 

690 :class:`NonUniformImage <matplotlib.image.NonUniformImage>` can be used to 

691 display actual bin edges with interpolation: 

692 

693 >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated', 

694 ... aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]]) 

695 >>> im = NonUniformImage(ax, interpolation='bilinear') 

696 >>> xcenters = (xedges[:-1] + xedges[1:]) / 2 

697 >>> ycenters = (yedges[:-1] + yedges[1:]) / 2 

698 >>> im.set_data(xcenters, ycenters, H) 

699 >>> ax.images.append(im) 

700 >>> plt.show() 

701 

702 """ 

703 from numpy import histogramdd 

704 

705 try: 

706 N = len(bins) 

707 except TypeError: 

708 N = 1 

709 

710 if N != 1 and N != 2: 

711 xedges = yedges = asarray(bins) 

712 bins = [xedges, yedges] 

713 hist, edges = histogramdd([x, y], bins, range, normed, weights, density) 

714 return hist, edges[0], edges[1] 

715 

716 

717@set_module('numpy') 

718def mask_indices(n, mask_func, k=0): 

719 """ 

720 Return the indices to access (n, n) arrays, given a masking function. 

721 

722 Assume `mask_func` is a function that, for a square array a of size 

723 ``(n, n)`` with a possible offset argument `k`, when called as 

724 ``mask_func(a, k)`` returns a new array with zeros in certain locations 

725 (functions like `triu` or `tril` do precisely this). Then this function 

726 returns the indices where the non-zero values would be located. 

727 

728 Parameters 

729 ---------- 

730 n : int 

731 The returned indices will be valid to access arrays of shape (n, n). 

732 mask_func : callable 

733 A function whose call signature is similar to that of `triu`, `tril`. 

734 That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. 

735 `k` is an optional argument to the function. 

736 k : scalar 

737 An optional argument which is passed through to `mask_func`. Functions 

738 like `triu`, `tril` take a second argument that is interpreted as an 

739 offset. 

740 

741 Returns 

742 ------- 

743 indices : tuple of arrays. 

744 The `n` arrays of indices corresponding to the locations where 

745 ``mask_func(np.ones((n, n)), k)`` is True. 

746 

747 See Also 

748 -------- 

749 triu, tril, triu_indices, tril_indices 

750 

751 Notes 

752 ----- 

753 .. versionadded:: 1.4.0 

754 

755 Examples 

756 -------- 

757 These are the indices that would allow you to access the upper triangular 

758 part of any 3x3 array: 

759 

760 >>> iu = np.mask_indices(3, np.triu) 

761 

762 For example, if `a` is a 3x3 array: 

763 

764 >>> a = np.arange(9).reshape(3, 3) 

765 >>> a 

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

767 [3, 4, 5], 

768 [6, 7, 8]]) 

769 >>> a[iu] 

770 array([0, 1, 2, 4, 5, 8]) 

771 

772 An offset can be passed also to the masking function. This gets us the 

773 indices starting on the first diagonal right of the main one: 

774 

775 >>> iu1 = np.mask_indices(3, np.triu, 1) 

776 

777 with which we now extract only three elements: 

778 

779 >>> a[iu1] 

780 array([1, 2, 5]) 

781 

782 """ 

783 m = ones((n, n), int) 

784 a = mask_func(m, k) 

785 return nonzero(a != 0) 

786 

787 

788@set_module('numpy') 

789def tril_indices(n, k=0, m=None): 

790 """ 

791 Return the indices for the lower-triangle of an (n, m) array. 

792 

793 Parameters 

794 ---------- 

795 n : int 

796 The row dimension of the arrays for which the returned 

797 indices will be valid. 

798 k : int, optional 

799 Diagonal offset (see `tril` for details). 

800 m : int, optional 

801 .. versionadded:: 1.9.0 

802 

803 The column dimension of the arrays for which the returned 

804 arrays will be valid. 

805 By default `m` is taken equal to `n`. 

806 

807 

808 Returns 

809 ------- 

810 inds : tuple of arrays 

811 The indices for the triangle. The returned tuple contains two arrays, 

812 each with the indices along one dimension of the array. 

813 

814 See also 

815 -------- 

816 triu_indices : similar function, for upper-triangular. 

817 mask_indices : generic function accepting an arbitrary mask function. 

818 tril, triu 

819 

820 Notes 

821 ----- 

822 .. versionadded:: 1.4.0 

823 

824 Examples 

825 -------- 

826 Compute two different sets of indices to access 4x4 arrays, one for the 

827 lower triangular part starting at the main diagonal, and one starting two 

828 diagonals further right: 

829 

830 >>> il1 = np.tril_indices(4) 

831 >>> il2 = np.tril_indices(4, 2) 

832 

833 Here is how they can be used with a sample array: 

834 

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

836 >>> a 

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

838 [ 4, 5, 6, 7], 

839 [ 8, 9, 10, 11], 

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

841 

842 Both for indexing: 

843 

844 >>> a[il1] 

845 array([ 0, 4, 5, ..., 13, 14, 15]) 

846 

847 And for assigning values: 

848 

849 >>> a[il1] = -1 

850 >>> a 

851 array([[-1, 1, 2, 3], 

852 [-1, -1, 6, 7], 

853 [-1, -1, -1, 11], 

854 [-1, -1, -1, -1]]) 

855 

856 These cover almost the whole array (two diagonals right of the main one): 

857 

858 >>> a[il2] = -10 

859 >>> a 

860 array([[-10, -10, -10, 3], 

861 [-10, -10, -10, -10], 

862 [-10, -10, -10, -10], 

863 [-10, -10, -10, -10]]) 

864 

865 """ 

866 return nonzero(tri(n, m, k=k, dtype=bool)) 

867 

868 

869def _trilu_indices_form_dispatcher(arr, k=None): 

870 return (arr,) 

871 

872 

873@array_function_dispatch(_trilu_indices_form_dispatcher) 

874def tril_indices_from(arr, k=0): 

875 """ 

876 Return the indices for the lower-triangle of arr. 

877 

878 See `tril_indices` for full details. 

879 

880 Parameters 

881 ---------- 

882 arr : array_like 

883 The indices will be valid for square arrays whose dimensions are 

884 the same as arr. 

885 k : int, optional 

886 Diagonal offset (see `tril` for details). 

887 

888 See Also 

889 -------- 

890 tril_indices, tril 

891 

892 Notes 

893 ----- 

894 .. versionadded:: 1.4.0 

895 

896 """ 

897 if arr.ndim != 2: 

898 raise ValueError("input array must be 2-d") 

899 return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) 

900 

901 

902@set_module('numpy') 

903def triu_indices(n, k=0, m=None): 

904 """ 

905 Return the indices for the upper-triangle of an (n, m) array. 

906 

907 Parameters 

908 ---------- 

909 n : int 

910 The size of the arrays for which the returned indices will 

911 be valid. 

912 k : int, optional 

913 Diagonal offset (see `triu` for details). 

914 m : int, optional 

915 .. versionadded:: 1.9.0 

916 

917 The column dimension of the arrays for which the returned 

918 arrays will be valid. 

919 By default `m` is taken equal to `n`. 

920 

921 

922 Returns 

923 ------- 

924 inds : tuple, shape(2) of ndarrays, shape(`n`) 

925 The indices for the triangle. The returned tuple contains two arrays, 

926 each with the indices along one dimension of the array. Can be used 

927 to slice a ndarray of shape(`n`, `n`). 

928 

929 See also 

930 -------- 

931 tril_indices : similar function, for lower-triangular. 

932 mask_indices : generic function accepting an arbitrary mask function. 

933 triu, tril 

934 

935 Notes 

936 ----- 

937 .. versionadded:: 1.4.0 

938 

939 Examples 

940 -------- 

941 Compute two different sets of indices to access 4x4 arrays, one for the 

942 upper triangular part starting at the main diagonal, and one starting two 

943 diagonals further right: 

944 

945 >>> iu1 = np.triu_indices(4) 

946 >>> iu2 = np.triu_indices(4, 2) 

947 

948 Here is how they can be used with a sample array: 

949 

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

951 >>> a 

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

953 [ 4, 5, 6, 7], 

954 [ 8, 9, 10, 11], 

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

956 

957 Both for indexing: 

958 

959 >>> a[iu1] 

960 array([ 0, 1, 2, ..., 10, 11, 15]) 

961 

962 And for assigning values: 

963 

964 >>> a[iu1] = -1 

965 >>> a 

966 array([[-1, -1, -1, -1], 

967 [ 4, -1, -1, -1], 

968 [ 8, 9, -1, -1], 

969 [12, 13, 14, -1]]) 

970 

971 These cover only a small part of the whole array (two diagonals right 

972 of the main one): 

973 

974 >>> a[iu2] = -10 

975 >>> a 

976 array([[ -1, -1, -10, -10], 

977 [ 4, -1, -1, -10], 

978 [ 8, 9, -1, -1], 

979 [ 12, 13, 14, -1]]) 

980 

981 """ 

982 return nonzero(~tri(n, m, k=k-1, dtype=bool)) 

983 

984 

985@array_function_dispatch(_trilu_indices_form_dispatcher) 

986def triu_indices_from(arr, k=0): 

987 """ 

988 Return the indices for the upper-triangle of arr. 

989 

990 See `triu_indices` for full details. 

991 

992 Parameters 

993 ---------- 

994 arr : ndarray, shape(N, N) 

995 The indices will be valid for square arrays. 

996 k : int, optional 

997 Diagonal offset (see `triu` for details). 

998 

999 Returns 

1000 ------- 

1001 triu_indices_from : tuple, shape(2) of ndarray, shape(N) 

1002 Indices for the upper-triangle of `arr`. 

1003 

1004 See Also 

1005 -------- 

1006 triu_indices, triu 

1007 

1008 Notes 

1009 ----- 

1010 .. versionadded:: 1.4.0 

1011 

1012 """ 

1013 if arr.ndim != 2: 

1014 raise ValueError("input array must be 2-d") 

1015 return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])