1""" Basic functions for manipulating 2d arrays

2

3"""

4 import functools

5 import operator

6

7 from numpy.core.numeric import (

8 asanyarray, arange, zeros, greater_equal, multiply, ones,

9 asarray, where, int8, int16, int32, int64, intp, empty, promote_types,

10 diagonal, nonzero, indices

11 )

12 from numpy.core.overrides import set_array_function_like_doc, set_module

13 from numpy.core import overrides

14 from numpy.core import iinfo

15 from numpy.lib.stride_tricks import broadcast_to

16

17

18__all__ = [

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

20 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices',

21 'tril_indices_from', 'triu_indices', 'triu_indices_from', ]

22

23

24 array_function_dispatch = functools.partial(

25 overrides.array_function_dispatch, module='numpy')

26

27

28 i1 = iinfo(int8)

29 i2 = iinfo(int16)

30 i4 = iinfo(int32)

31

32

33 def _min_int(low, high):

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

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

36 return int8

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

38 return int16

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

40 return int32

41 return int64

42

43

44 def _flip_dispatcher(m):

45 return (m,)

46

47

48@array_function_dispatch(_flip_dispatcher)

49 def fliplr(m):

50 """

51 Reverse the order of elements along axis 1 (left/right).

52

53 For a 2-D array, this flips the entries in each row in the left/right

54 direction. Columns are preserved, but appear in a different order than

55 before.

56

57 Parameters

58 ----------

59 m : array_like

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

61

62 Returns

63 -------

64 f : ndarray

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

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

67

68 See Also

69 --------

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

71 flip : Flip array in one or more dimensions.

72 rot90 : Rotate array counterclockwise.

73

74 Notes

75 -----

76 Equivalent to ``m[:,::-1]`` or ``np.flip(m, axis=1)``.

77 Requires the array to be at least 2-D.

78

79 Examples

80 --------

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

82 >>> A

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

84 [0., 2., 0.],

85 [0., 0., 3.]])

86 >>> np.fliplr(A)

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

88 [0., 2., 0.],

89 [3., 0., 0.]])

90

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

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

93 True

94

95 """

96 m = asanyarray(m)

97 if m.ndim < 2:

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

99 return m[:, ::-1]

100

101

102@array_function_dispatch(_flip_dispatcher)

103 def flipud(m):

104 """

105 Reverse the order of elements along axis 0 (up/down).

106

107 For a 2-D array, this flips the entries in each column in the up/down

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

109

110 Parameters

111 ----------

112 m : array_like

113 Input array.

114

115 Returns

116 -------

117 out : array_like

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

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

120

121 See Also

122 --------

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

124 flip : Flip array in one or more dimensions.

125 rot90 : Rotate array counterclockwise.

126

127 Notes

128 -----

129 Equivalent to ``m[::-1, ...]`` or ``np.flip(m, axis=0)``.

130 Requires the array to be at least 1-D.

131

132 Examples

133 --------

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

135 >>> A

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

137 [0., 2., 0.],

138 [0., 0., 3.]])

139 >>> np.flipud(A)

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

141 [0., 2., 0.],

142 [1., 0., 0.]])

143

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

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

146 True

147

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

149 array([2, 1])

150

151 """

152 m = asanyarray(m)

153 if m.ndim < 1:

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

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

156

157

158 def _eye_dispatcher(N, M=None, k=None, dtype=None, order=None, *, like=None):

159 return (like,)

160

161

162@set_array_function_like_doc

163@set_module('numpy')

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

165 """

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

167

168 Parameters

169 ----------

170 N : int

171 Number of rows in the output.

172 M : int, optional

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

174 k : int, optional

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

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

177 to a lower diagonal.

178 dtype : data-type, optional

179 Data-type of the returned array.

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

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

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

183

184 .. versionadded:: 1.14.0

185 ${ARRAY_FUNCTION_LIKE}

186

187 .. versionadded:: 1.20.0

188

189 Returns

190 -------

191 I : ndarray of shape (N,M)

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

193 diagonal, whose values are equal to one.

194

195 See Also

196 --------

197 identity : (almost) equivalent function

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

199

200 Examples

201 --------

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

203 array([[1, 0],

204 [0, 1]])

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

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

207 [0., 0., 1.],

208 [0., 0., 0.]])

209

210 """

211 if like is not None:

212 return _eye_with_like(N, M=M, k=k, dtype=dtype, order=order, like=like)

213 if M is None:

214 M = N

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

216 if k >= M:

217 return m

218 # Ensure M and k are integers, so we don't get any surprise casting

219 # results in the expressions `M-k` and `M+1` used below. This avoids

220 # a problem with inputs with type (for example) np.uint64.

221 M = operator.index(M)

222 k = operator.index(k)

223 if k >= 0:

224 i = k

225 else:

226 i = (-k) * M

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

228 return m

229

230

231_eye_with_like = array_function_dispatch(

232 _eye_dispatcher

233)(eye)

234

235

236 def _diag_dispatcher(v, k=None):

237 return (v,)

238

239

240@array_function_dispatch(_diag_dispatcher)

241 def diag(v, k=0):

242 """

243 Extract a diagonal or construct a diagonal array.

244

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

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

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

248 are using.

249

250 Parameters

251 ----------

252 v : array_like

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

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

255 diagonal.

256 k : int, optional

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

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

259 diagonal.

260

261 Returns

262 -------

263 out : ndarray

264 The extracted diagonal or constructed diagonal array.

265

266 See Also

267 --------

268 diagonal : Return specified diagonals.

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

270 trace : Sum along diagonals.

271 triu : Upper triangle of an array.

272 tril : Lower triangle of an array.

273

274 Examples

275 --------

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

277 >>> x

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

279 [3, 4, 5],

280 [6, 7, 8]])

281

282 >>> np.diag(x)

283 array([0, 4, 8])

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

285 array([1, 5])

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

287 array([3, 7])

288

289 >>> np.diag(np.diag(x))

290 array([[0, 0, 0],

291 [0, 4, 0],

292 [0, 0, 8]])

293

294 """

295 v = asanyarray(v)

296 s = v.shape

297 if len(s) == 1:

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

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

300 if k >= 0:

301 i = k

302 else:

303 i = (-k) * n

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

305 return res

306 elif len(s) == 2:

307 return diagonal(v, k)

308 else:

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

310

311

312@array_function_dispatch(_diag_dispatcher)

313 def diagflat(v, k=0):

314 """

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

316

317 Parameters

318 ----------

319 v : array_like

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

321 diagonal of the output.

322 k : int, optional

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

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

325 (below) the main.

326

327 Returns

328 -------

329 out : ndarray

330 The 2-D output array.

331

332 See Also

333 --------

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

335 diagonal : Return specified diagonals.

336 trace : Sum along diagonals.

337

338 Examples

339 --------

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

341 array([[1, 0, 0, 0],

342 [0, 2, 0, 0],

343 [0, 0, 3, 0],

344 [0, 0, 0, 4]])

345

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

347 array([[0, 1, 0],

348 [0, 0, 2],

349 [0, 0, 0]])

350

351 """

352 try:

353 wrap = v.__array_wrap__

354 except AttributeError:

355 wrap = None

356 v = asarray(v).ravel()

357 s = len(v)

358 n = s + abs(k)

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

360 if (k >= 0):

361 i = arange(0, n-k, dtype=intp)

362 fi = i+k+i*n

363 else:

364 i = arange(0, n+k, dtype=intp)

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

366 res.flat[fi] = v

367 if not wrap:

368 return res

369 return wrap(res)

370

371

372 def _tri_dispatcher(N, M=None, k=None, dtype=None, *, like=None):

373 return (like,)

374

375

376@set_array_function_like_doc

377@set_module('numpy')

378 def tri(N, M=None, k=0, dtype=float, *, like=None):

379 """

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

381

382 Parameters

383 ----------

384 N : int

385 Number of rows in the array.

386 M : int, optional

387 Number of columns in the array.

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

389 k : int, optional

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

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

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

393 dtype : dtype, optional

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

395 ${ARRAY_FUNCTION_LIKE}

396

397 .. versionadded:: 1.20.0

398

399 Returns

400 -------

401 tri : ndarray of shape (N, M)

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

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

404

405 Examples

406 --------

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

408 array([[1, 1, 1, 0, 0],

409 [1, 1, 1, 1, 0],

410 [1, 1, 1, 1, 1]])

411

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

413 array([[0., 0., 0., 0., 0.],

414 [1., 0., 0., 0., 0.],

415 [1., 1., 0., 0., 0.]])

416

417 """

418 if like is not None:

419 return _tri_with_like(N, M=M, k=k, dtype=dtype, like=like)

420

421 if M is None:

422 M = N

423

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

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

426

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

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

429

430 return m

431

432

433_tri_with_like = array_function_dispatch(

434 _tri_dispatcher

435)(tri)

436

437

438 def _trilu_dispatcher(m, k=None):

439 return (m,)

440

441

442@array_function_dispatch(_trilu_dispatcher)

443 def tril(m, k=0):

444 """

445 Lower triangle of an array.

446

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

448 For arrays with ``ndim`` exceeding 2, `tril` will apply to the final two

449 axes.

450

451 Parameters

452 ----------

453 m : array_like, shape (..., M, N)

454 Input array.

455 k : int, optional

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

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

458

459 Returns

460 -------

461 tril : ndarray, shape (..., M, N)

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

463

464 See Also

465 --------

466 triu : same thing, only for the upper triangle

467

468 Examples

469 --------

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

471 array([[ 0, 0, 0],

472 [ 4, 0, 0],

473 [ 7, 8, 0],

474 [10, 11, 12]])

475

476 >>> np.tril(np.arange(3*4*5).reshape(3, 4, 5))

477 array([[[ 0, 0, 0, 0, 0],

478 [ 5, 6, 0, 0, 0],

479 [10, 11, 12, 0, 0],

480 [15, 16, 17, 18, 0]],

481 [[20, 0, 0, 0, 0],

482 [25, 26, 0, 0, 0],

483 [30, 31, 32, 0, 0],

484 [35, 36, 37, 38, 0]],

485 [[40, 0, 0, 0, 0],

486 [45, 46, 0, 0, 0],

487 [50, 51, 52, 0, 0],

488 [55, 56, 57, 58, 0]]])

489

490 """

491 m = asanyarray(m)

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

493

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

495

496

497@array_function_dispatch(_trilu_dispatcher)

498 def triu(m, k=0):

499 """

500 Upper triangle of an array.

501

502 Return a copy of an array with the elements below the `k`-th diagonal

503 zeroed. For arrays with ``ndim`` exceeding 2, `triu` will apply to the

504 final two axes.

505

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

507

508 See Also

509 --------

510 tril : lower triangle of an array

511

512 Examples

513 --------

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

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

516 [ 4, 5, 6],

517 [ 0, 8, 9],

518 [ 0, 0, 12]])

519

520 >>> np.triu(np.arange(3*4*5).reshape(3, 4, 5))

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

522 [ 0, 6, 7, 8, 9],

523 [ 0, 0, 12, 13, 14],

524 [ 0, 0, 0, 18, 19]],

525 [[20, 21, 22, 23, 24],

526 [ 0, 26, 27, 28, 29],

527 [ 0, 0, 32, 33, 34],

528 [ 0, 0, 0, 38, 39]],

529 [[40, 41, 42, 43, 44],

530 [ 0, 46, 47, 48, 49],

531 [ 0, 0, 52, 53, 54],

532 [ 0, 0, 0, 58, 59]]])

533

534 """

535 m = asanyarray(m)

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

537

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

539

540

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

542 return (x,)

543

544

545# Originally borrowed from John Hunter and matplotlib

546@array_function_dispatch(_vander_dispatcher)

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

548 """

549 Generate a Vandermonde matrix.

550

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

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

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

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

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

556 Theophile Vandermonde.

557

558 Parameters

559 ----------

560 x : array_like

561 1-D input array.

562 N : int, optional

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

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

565 increasing : bool, optional

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

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

568

569 .. versionadded:: 1.9.0

570

571 Returns

572 -------

573 out : ndarray

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

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

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

577

578 See Also

579 --------

580 polynomial.polynomial.polyvander

581

582 Examples

583 --------

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

585 >>> N = 3

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

587 array([[ 1, 1, 1],

588 [ 4, 2, 1],

589 [ 9, 3, 1],

590 [25, 5, 1]])

591

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

593 array([[ 1, 1, 1],

594 [ 4, 2, 1],

595 [ 9, 3, 1],

596 [25, 5, 1]])

597

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

599 >>> np.vander(x)

600 array([[ 1, 1, 1, 1],

601 [ 8, 4, 2, 1],

602 [ 27, 9, 3, 1],

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

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

605 array([[ 1, 1, 1, 1],

606 [ 1, 2, 4, 8],

607 [ 1, 3, 9, 27],

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

609

610 The determinant of a square Vandermonde matrix is the product

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

612

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

614 48.000000000000043 # may vary

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

616 48

617

618 """

619 x = asarray(x)

620 if x.ndim != 1:

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

622 if N is None:

623 N = len(x)

624

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

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

627

628 if N > 0:

629 tmp[:, 0] = 1

630 if N > 1:

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

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

633

634 return v

635

636

637 def _histogram2d_dispatcher(x, y, bins=None, range=None, density=None,

638 weights=None):

639 yield x

640 yield y

641

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

643 try:

644 N = len(bins)

645 except TypeError:

646 N = 1

647 if N == 2:

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

649 else:

650 yield bins

651

652 yield weights

653

654

655@array_function_dispatch(_histogram2d_dispatcher)

656 def histogram2d(x, y, bins=10, range=None, density=None, weights=None):

657 """

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

659

660 Parameters

661 ----------

662 x : array_like, shape (N,)

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

664 histogrammed.

665 y : array_like, shape (N,)

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

667 histogrammed.

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

669 The bin specification:

670

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

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

673 (x_edges=y_edges=bins).

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

675 (nx, ny = bins).

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

677 (x_edges, y_edges = bins).

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

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

680

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

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

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

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

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

686 density : bool, optional

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

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

689 ``bin_count / sample_count / bin_area``.

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

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

692 Weights are normalized to 1 if `density` is True. If `density` is

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

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

695

696 Returns

697 -------

698 H : ndarray, shape(nx, ny)

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

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

701 histogrammed along the second dimension.

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

703 The bin edges along the first dimension.

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

705 The bin edges along the second dimension.

706

707 See Also

708 --------

709 histogram : 1D histogram

710 histogramdd : Multidimensional histogram

711

712 Notes

713 -----

714 When `density` is True, then the returned histogram is the sample

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

716 ``bin_value * bin_area`` is 1.

717

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

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

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

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

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

723

724 Examples

725 --------

726 >>> from matplotlib.image import NonUniformImage

727 >>> import matplotlib.pyplot as plt

728

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

730 edges:

731

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

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

734

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

736

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

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

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

740 >>> # Histogram does not follow Cartesian convention (see Notes),

741 >>> # therefore transpose H for visualization purposes.

742 >>> H = H.T

743

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

745

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

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

748 >>> plt.imshow(H, interpolation='nearest', origin='lower',

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

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

751

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

753

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

755 ... aspect='equal')

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

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

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

759

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

761 display actual bin edges with interpolation:

762

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

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

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

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

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

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

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

770 >>> plt.show()

771

772 It is also possible to construct a 2-D histogram without specifying bin

773 edges:

774

775 >>> # Generate non-symmetric test data

776 >>> n = 10000

777 >>> x = np.linspace(1, 100, n)

778 >>> y = 2*np.log(x) + np.random.rand(n) - 0.5

779 >>> # Compute 2d histogram. Note the order of x/y and xedges/yedges

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

781

782 Now we can plot the histogram using

783 :func:`pcolormesh <matplotlib.pyplot.pcolormesh>`, and a

784 :func:`hexbin <matplotlib.pyplot.hexbin>` for comparison.

785

786 >>> # Plot histogram using pcolormesh

787 >>> fig, (ax1, ax2) = plt.subplots(ncols=2, sharey=True)

788 >>> ax1.pcolormesh(xedges, yedges, H, cmap='rainbow')

789 >>> ax1.plot(x, 2*np.log(x), 'k-')

790 >>> ax1.set_xlim(x.min(), x.max())

791 >>> ax1.set_ylim(y.min(), y.max())

792 >>> ax1.set_xlabel('x')

793 >>> ax1.set_ylabel('y')

794 >>> ax1.set_title('histogram2d')

795 >>> ax1.grid()

796

797 >>> # Create hexbin plot for comparison

798 >>> ax2.hexbin(x, y, gridsize=20, cmap='rainbow')

799 >>> ax2.plot(x, 2*np.log(x), 'k-')

800 >>> ax2.set_title('hexbin')

801 >>> ax2.set_xlim(x.min(), x.max())

802 >>> ax2.set_xlabel('x')

803 >>> ax2.grid()

804

805 >>> plt.show()

806 """

807 from numpy import histogramdd

808

809 if len(x) != len(y):

810 raise ValueError('x and y must have the same length.')

811

812 try:

813 N = len(bins)

814 except TypeError:

815 N = 1

816

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

818 xedges = yedges = asarray(bins)

819 bins = [xedges, yedges]

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

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

822

823

824@set_module('numpy')

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

826 """

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

828

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

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

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

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

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

834

835 Parameters

836 ----------

837 n : int

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

839 mask_func : callable

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

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

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

843 k : scalar

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

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

846 offset.

847

848 Returns

849 -------

850 indices : tuple of arrays.

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

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

853

854 See Also

855 --------

856 triu, tril, triu_indices, tril_indices

857

858 Notes

859 -----

860 .. versionadded:: 1.4.0

861

862 Examples

863 --------

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

865 part of any 3x3 array:

866

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

868

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

870

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

872 >>> a

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

874 [3, 4, 5],

875 [6, 7, 8]])

876 >>> a[iu]

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

878

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

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

881

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

883

884 with which we now extract only three elements:

885

886 >>> a[iu1]

887 array([1, 2, 5])

888

889 """

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

891 a = mask_func(m, k)

892 return nonzero(a != 0)

893

894

895@set_module('numpy')

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

897 """

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

899

900 Parameters

901 ----------

902 n : int

903 The row dimension of the arrays for which the returned

904 indices will be valid.

905 k : int, optional

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

907 m : int, optional

908 .. versionadded:: 1.9.0

909

910 The column dimension of the arrays for which the returned

911 arrays will be valid.

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

913

914

915 Returns

916 -------

917 inds : tuple of arrays

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

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

920

921 See also

922 --------

923 triu_indices : similar function, for upper-triangular.

924 mask_indices : generic function accepting an arbitrary mask function.

925 tril, triu

926

927 Notes

928 -----

929 .. versionadded:: 1.4.0

930

931 Examples

932 --------

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

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

935 diagonals further right:

936

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

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

939

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

941

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

943 >>> a

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

945 [ 4, 5, 6, 7],

946 [ 8, 9, 10, 11],

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

948

949 Both for indexing:

950

951 >>> a[il1]

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

953

954 And for assigning values:

955

956 >>> a[il1] = -1

957 >>> a

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

959 [-1, -1, 6, 7],

960 [-1, -1, -1, 11],

961 [-1, -1, -1, -1]])

962

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

964

965 >>> a[il2] = -10

966 >>> a

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

968 [-10, -10, -10, -10],

969 [-10, -10, -10, -10],

970 [-10, -10, -10, -10]])

971

972 """

973 tri_ = tri(n, m, k=k, dtype=bool)

974

975 return tuple(broadcast_to(inds, tri_.shape)[tri_]

976 for inds in indices(tri_.shape, sparse=True))

977

978

979 def _trilu_indices_form_dispatcher(arr, k=None):

980 return (arr,)

981

982

983@array_function_dispatch(_trilu_indices_form_dispatcher)

984 def tril_indices_from(arr, k=0):

985 """

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

987

988 See `tril_indices` for full details.

989

990 Parameters

991 ----------

992 arr : array_like

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

994 the same as arr.

995 k : int, optional

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

997

998 See Also

999 --------

1000 tril_indices, tril

1001

1002 Notes

1003 -----

1004 .. versionadded:: 1.4.0

1005

1006 """

1007 if arr.ndim != 2:

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

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

1010

1011

1012@set_module('numpy')

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

1014 """

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

1016

1017 Parameters

1018 ----------

1019 n : int

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

1021 be valid.

1022 k : int, optional

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

1024 m : int, optional

1025 .. versionadded:: 1.9.0

1026

1027 The column dimension of the arrays for which the returned

1028 arrays will be valid.

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

1030

1031

1032 Returns

1033 -------

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

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

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

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

1038

1039 See also

1040 --------

1041 tril_indices : similar function, for lower-triangular.

1042 mask_indices : generic function accepting an arbitrary mask function.

1043 triu, tril

1044

1045 Notes

1046 -----

1047 .. versionadded:: 1.4.0

1048

1049 Examples

1050 --------

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

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

1053 diagonals further right:

1054

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

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

1057

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

1059

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

1061 >>> a

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

1063 [ 4, 5, 6, 7],

1064 [ 8, 9, 10, 11],

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

1066

1067 Both for indexing:

1068

1069 >>> a[iu1]

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

1071

1072 And for assigning values:

1073

1074 >>> a[iu1] = -1

1075 >>> a

1076 array([[-1, -1, -1, -1],

1077 [ 4, -1, -1, -1],

1078 [ 8, 9, -1, -1],

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

1080

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

1082 of the main one):

1083

1084 >>> a[iu2] = -10

1085 >>> a

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

1087 [ 4, -1, -1, -10],

1088 [ 8, 9, -1, -1],

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

1090

1091 """

1092 tri_ = ~tri(n, m, k=k - 1, dtype=bool)

1093

1094 return tuple(broadcast_to(inds, tri_.shape)[tri_]

1095 for inds in indices(tri_.shape, sparse=True))

1096

1097

1098@array_function_dispatch(_trilu_indices_form_dispatcher)

1099 def triu_indices_from(arr, k=0):

1100 """

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

1102

1103 See `triu_indices` for full details.

1104

1105 Parameters

1106 ----------

1107 arr : ndarray, shape(N, N)

1108 The indices will be valid for square arrays.

1109 k : int, optional

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

1111

1112 Returns

1113 -------

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

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

1116

1117 See Also

1118 --------

1119 triu_indices, triu

1120

1121 Notes

1122 -----

1123 .. versionadded:: 1.4.0

1124

1125 """

1126 if arr.ndim != 2:

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

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