function_base.pyi
1 import sys 2 from collections.abc import Sequence, Iterator, Callable, Iterable 3 from typing import ( 4 Literal as L, 5 Any, 6 TypeVar, 7 overload, 8 Protocol, 9 SupportsIndex, 10 SupportsInt, 11 ) 12 13 if sys.version_info >= (3, 10): 14 from typing import TypeGuard 15 else: 16 from typing_extensions import TypeGuard 17 18 from numpy import ( 19 vectorize as vectorize, 20 ufunc, 21 generic, 22 floating, 23 complexfloating, 24 intp, 25 float64, 26 complex128, 27 timedelta64, 28 datetime64, 29 object_, 30 _OrderKACF, 31 ) 32 33 from numpy._typing import ( 34 NDArray, 35 ArrayLike, 36 DTypeLike, 37 _ShapeLike, 38 _ScalarLike_co, 39 _DTypeLike, 40 _ArrayLike, 41 _ArrayLikeInt_co, 42 _ArrayLikeFloat_co, 43 _ArrayLikeComplex_co, 44 _ArrayLikeTD64_co, 45 _ArrayLikeDT64_co, 46 _ArrayLikeObject_co, 47 _FloatLike_co, 48 _ComplexLike_co, 49 ) 50 51 from numpy.core.function_base import ( 52 add_newdoc as add_newdoc, 53 ) 54 55 from numpy.core.multiarray import ( 56 add_docstring as add_docstring, 57 bincount as bincount, 58 ) 59 60 from numpy.core.umath import _add_newdoc_ufunc 61 62 _T = TypeVar("_T") 63 _T_co = TypeVar("_T_co", covariant=True) 64 _SCT = TypeVar("_SCT", bound=generic) 65 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) 66 67 _2Tuple = tuple[_T, _T] 68 69 class _TrimZerosSequence(Protocol[_T_co]): 70 def __len__(self) -> int: ... 71 def __getitem__(self, key: slice, /) -> _T_co: ... 72 def __iter__(self) -> Iterator[Any]: ... 73 74 class _SupportsWriteFlush(Protocol): 75 def write(self, s: str, /) -> object: ... 76 def flush(self) -> object: ... 77 78 __all__: list[str] 79 80 # NOTE: This is in reality a re-export of `np.core.umath._add_newdoc_ufunc` 81 def add_newdoc_ufunc(ufunc: ufunc, new_docstring: str, /) -> None: ... 82 83 @overload 84 def rot90( 85 m: _ArrayLike[_SCT], 86 k: int = ..., 87 axes: tuple[int, int] = ..., 88 ) -> NDArray[_SCT]: ... 89 @overload 90 def rot90( 91 m: ArrayLike, 92 k: int = ..., 93 axes: tuple[int, int] = ..., 94 ) -> NDArray[Any]: ... 95 96 @overload 97 def flip(m: _SCT, axis: None = ...) -> _SCT: ... 98 @overload 99 def flip(m: _ScalarLike_co, axis: None = ...) -> Any: ... 100 @overload 101 def flip(m: _ArrayLike[_SCT], axis: None | _ShapeLike = ...) -> NDArray[_SCT]: ... 102 @overload 103 def flip(m: ArrayLike, axis: None | _ShapeLike = ...) -> NDArray[Any]: ... 104 105 def iterable(y: object) -> TypeGuard[Iterable[Any]]: ... 106 107 @overload 108 def average( 109 a: _ArrayLikeFloat_co, 110 axis: None = ..., 111 weights: None | _ArrayLikeFloat_co= ..., 112 returned: L[False] = ..., 113 keepdims: L[False] = ..., 114 ) -> floating[Any]: ... 115 @overload 116 def average( 117 a: _ArrayLikeComplex_co, 118 axis: None = ..., 119 weights: None | _ArrayLikeComplex_co = ..., 120 returned: L[False] = ..., 121 keepdims: L[False] = ..., 122 ) -> complexfloating[Any, Any]: ... 123 @overload 124 def average( 125 a: _ArrayLikeObject_co, 126 axis: None = ..., 127 weights: None | Any = ..., 128 returned: L[False] = ..., 129 keepdims: L[False] = ..., 130 ) -> Any: ... 131 @overload 132 def average( 133 a: _ArrayLikeFloat_co, 134 axis: None = ..., 135 weights: None | _ArrayLikeFloat_co= ..., 136 returned: L[True] = ..., 137 keepdims: L[False] = ..., 138 ) -> _2Tuple[floating[Any]]: ... 139 @overload 140 def average( 141 a: _ArrayLikeComplex_co, 142 axis: None = ..., 143 weights: None | _ArrayLikeComplex_co = ..., 144 returned: L[True] = ..., 145 keepdims: L[False] = ..., 146 ) -> _2Tuple[complexfloating[Any, Any]]: ... 147 @overload 148 def average( 149 a: _ArrayLikeObject_co, 150 axis: None = ..., 151 weights: None | Any = ..., 152 returned: L[True] = ..., 153 keepdims: L[False] = ..., 154 ) -> _2Tuple[Any]: ... 155 @overload 156 def average( 157 a: _ArrayLikeComplex_co | _ArrayLikeObject_co, 158 axis: None | _ShapeLike = ..., 159 weights: None | Any = ..., 160 returned: L[False] = ..., 161 keepdims: bool = ..., 162 ) -> Any: ... 163 @overload 164 def average( 165 a: _ArrayLikeComplex_co | _ArrayLikeObject_co, 166 axis: None | _ShapeLike = ..., 167 weights: None | Any = ..., 168 returned: L[True] = ..., 169 keepdims: bool = ..., 170 ) -> _2Tuple[Any]: ... 171 172 @overload 173 def asarray_chkfinite( 174 a: _ArrayLike[_SCT], 175 dtype: None = ..., 176 order: _OrderKACF = ..., 177 ) -> NDArray[_SCT]: ... 178 @overload 179 def asarray_chkfinite( 180 a: object, 181 dtype: None = ..., 182 order: _OrderKACF = ..., 183 ) -> NDArray[Any]: ... 184 @overload 185 def asarray_chkfinite( 186 a: Any, 187 dtype: _DTypeLike[_SCT], 188 order: _OrderKACF = ..., 189 ) -> NDArray[_SCT]: ... 190 @overload 191 def asarray_chkfinite( 192 a: Any, 193 dtype: DTypeLike, 194 order: _OrderKACF = ..., 195 ) -> NDArray[Any]: ... 196 197 # TODO: Use PEP 612 `ParamSpec` once mypy supports `Concatenate` 198 # xref python/mypy#8645 199 @overload 200 def piecewise( 201 x: _ArrayLike[_SCT], 202 condlist: ArrayLike, 203 funclist: Sequence[Any | Callable[..., Any]], 204 *args: Any, 205 **kw: Any, 206 ) -> NDArray[_SCT]: ... 207 @overload 208 def piecewise( 209 x: ArrayLike, 210 condlist: ArrayLike, 211 funclist: Sequence[Any | Callable[..., Any]], 212 *args: Any, 213 **kw: Any, 214 ) -> NDArray[Any]: ... 215 216 def select( 217 condlist: Sequence[ArrayLike], 218 choicelist: Sequence[ArrayLike], 219 default: ArrayLike = ..., 220 ) -> NDArray[Any]: ... 221 222 @overload 223 def copy( 224 a: _ArrayType, 225 order: _OrderKACF, 226 subok: L[True], 227 ) -> _ArrayType: ... 228 @overload 229 def copy( 230 a: _ArrayType, 231 order: _OrderKACF = ..., 232 *, 233 subok: L[True], 234 ) -> _ArrayType: ... 235 @overload 236 def copy( 237 a: _ArrayLike[_SCT], 238 order: _OrderKACF = ..., 239 subok: L[False] = ..., 240 ) -> NDArray[_SCT]: ... 241 @overload 242 def copy( 243 a: ArrayLike, 244 order: _OrderKACF = ..., 245 subok: L[False] = ..., 246 ) -> NDArray[Any]: ... 247 248 def gradient( 249 f: ArrayLike, 250 *varargs: ArrayLike, 251 axis: None | _ShapeLike = ..., 252 edge_order: L[1, 2] = ..., 253 ) -> Any: ... 254 255 @overload 256 def diff( 257 a: _T, 258 n: L[0], 259 axis: SupportsIndex = ..., 260 prepend: ArrayLike = ..., 261 append: ArrayLike = ..., 262 ) -> _T: ... 263 @overload 264 def diff( 265 a: ArrayLike, 266 n: int = ..., 267 axis: SupportsIndex = ..., 268 prepend: ArrayLike = ..., 269 append: ArrayLike = ..., 270 ) -> NDArray[Any]: ... 271 272 @overload 273 def interp( 274 x: _ArrayLikeFloat_co, 275 xp: _ArrayLikeFloat_co, 276 fp: _ArrayLikeFloat_co, 277 left: None | _FloatLike_co = ..., 278 right: None | _FloatLike_co = ..., 279 period: None | _FloatLike_co = ..., 280 ) -> NDArray[float64]: ... 281 @overload 282 def interp( 283 x: _ArrayLikeFloat_co, 284 xp: _ArrayLikeFloat_co, 285 fp: _ArrayLikeComplex_co, 286 left: None | _ComplexLike_co = ..., 287 right: None | _ComplexLike_co = ..., 288 period: None | _FloatLike_co = ..., 289 ) -> NDArray[complex128]: ... 290 291 @overload 292 def angle(z: _ComplexLike_co, deg: bool = ...) -> floating[Any]: ... 293 @overload 294 def angle(z: object_, deg: bool = ...) -> Any: ... 295 @overload 296 def angle(z: _ArrayLikeComplex_co, deg: bool = ...) -> NDArray[floating[Any]]: ... 297 @overload 298 def angle(z: _ArrayLikeObject_co, deg: bool = ...) -> NDArray[object_]: ... 299 300 @overload 301 def unwrap( 302 p: _ArrayLikeFloat_co, 303 discont: None | float = ..., 304 axis: int = ..., 305 *, 306 period: float = ..., 307 ) -> NDArray[floating[Any]]: ... 308 @overload 309 def unwrap( 310 p: _ArrayLikeObject_co, 311 discont: None | float = ..., 312 axis: int = ..., 313 *, 314 period: float = ..., 315 ) -> NDArray[object_]: ... 316 317 def sort_complex(a: ArrayLike) -> NDArray[complexfloating[Any, Any]]: ... 318 319 def trim_zeros( 320 filt: _TrimZerosSequence[_T], 321 trim: L["f", "b", "fb", "bf"] = ..., 322 ) -> _T: ... 323 324 @overload 325 def extract(condition: ArrayLike, arr: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... 326 @overload 327 def extract(condition: ArrayLike, arr: ArrayLike) -> NDArray[Any]: ... 328 329 def place(arr: NDArray[Any], mask: ArrayLike, vals: Any) -> None: ... 330 331 def disp( 332 mesg: object, 333 device: None | _SupportsWriteFlush = ..., 334 linefeed: bool = ..., 335 ) -> None: ... 336 337 @overload 338 def cov( 339 m: _ArrayLikeFloat_co, 340 y: None | _ArrayLikeFloat_co = ..., 341 rowvar: bool = ..., 342 bias: bool = ..., 343 ddof: None | SupportsIndex | SupportsInt = ..., 344 fweights: None | ArrayLike = ..., 345 aweights: None | ArrayLike = ..., 346 *, 347 dtype: None = ..., 348 ) -> NDArray[floating[Any]]: ... 349 @overload 350 def cov( 351 m: _ArrayLikeComplex_co, 352 y: None | _ArrayLikeComplex_co = ..., 353 rowvar: bool = ..., 354 bias: bool = ..., 355 ddof: None | SupportsIndex | SupportsInt = ..., 356 fweights: None | ArrayLike = ..., 357 aweights: None | ArrayLike = ..., 358 *, 359 dtype: None = ..., 360 ) -> NDArray[complexfloating[Any, Any]]: ... 361 @overload 362 def cov( 363 m: _ArrayLikeComplex_co, 364 y: None | _ArrayLikeComplex_co = ..., 365 rowvar: bool = ..., 366 bias: bool = ..., 367 ddof: None | SupportsIndex | SupportsInt = ..., 368 fweights: None | ArrayLike = ..., 369 aweights: None | ArrayLike = ..., 370 *, 371 dtype: _DTypeLike[_SCT], 372 ) -> NDArray[_SCT]: ... 373 @overload 374 def cov( 375 m: _ArrayLikeComplex_co, 376 y: None | _ArrayLikeComplex_co = ..., 377 rowvar: bool = ..., 378 bias: bool = ..., 379 ddof: None | SupportsIndex | SupportsInt = ..., 380 fweights: None | ArrayLike = ..., 381 aweights: None | ArrayLike = ..., 382 *, 383 dtype: DTypeLike, 384 ) -> NDArray[Any]: ... 385 386 # NOTE `bias` and `ddof` have been deprecated 387 @overload 388 def corrcoef( 389 m: _ArrayLikeFloat_co, 390 y: None | _ArrayLikeFloat_co = ..., 391 rowvar: bool = ..., 392 *, 393 dtype: None = ..., 394 ) -> NDArray[floating[Any]]: ... 395 @overload 396 def corrcoef( 397 m: _ArrayLikeComplex_co, 398 y: None | _ArrayLikeComplex_co = ..., 399 rowvar: bool = ..., 400 *, 401 dtype: None = ..., 402 ) -> NDArray[complexfloating[Any, Any]]: ... 403 @overload 404 def corrcoef( 405 m: _ArrayLikeComplex_co, 406 y: None | _ArrayLikeComplex_co = ..., 407 rowvar: bool = ..., 408 *, 409 dtype: _DTypeLike[_SCT], 410 ) -> NDArray[_SCT]: ... 411 @overload 412 def corrcoef( 413 m: _ArrayLikeComplex_co, 414 y: None | _ArrayLikeComplex_co = ..., 415 rowvar: bool = ..., 416 *, 417 dtype: DTypeLike, 418 ) -> NDArray[Any]: ... 419 420 def blackman(M: _FloatLike_co) -> NDArray[floating[Any]]: ... 421 422 def bartlett(M: _FloatLike_co) -> NDArray[floating[Any]]: ... 423 424 def hanning(M: _FloatLike_co) -> NDArray[floating[Any]]: ... 425 426 def hamming(M: _FloatLike_co) -> NDArray[floating[Any]]: ... 427 428 def i0(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... 429 430 def kaiser( 431 M: _FloatLike_co, 432 beta: _FloatLike_co, 433 ) -> NDArray[floating[Any]]: ... 434 435 @overload 436 def sinc(x: _FloatLike_co) -> floating[Any]: ... 437 @overload 438 def sinc(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... 439 @overload 440 def sinc(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... 441 @overload 442 def sinc(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... 443 444 # NOTE: Deprecated 445 # def msort(a: ArrayLike) -> NDArray[Any]: ... 446 447 @overload 448 def median( 449 a: _ArrayLikeFloat_co, 450 axis: None = ..., 451 out: None = ..., 452 overwrite_input: bool = ..., 453 keepdims: L[False] = ..., 454 ) -> floating[Any]: ... 455 @overload 456 def median( 457 a: _ArrayLikeComplex_co, 458 axis: None = ..., 459 out: None = ..., 460 overwrite_input: bool = ..., 461 keepdims: L[False] = ..., 462 ) -> complexfloating[Any, Any]: ... 463 @overload 464 def median( 465 a: _ArrayLikeTD64_co, 466 axis: None = ..., 467 out: None = ..., 468 overwrite_input: bool = ..., 469 keepdims: L[False] = ..., 470 ) -> timedelta64: ... 471 @overload 472 def median( 473 a: _ArrayLikeObject_co, 474 axis: None = ..., 475 out: None = ..., 476 overwrite_input: bool = ..., 477 keepdims: L[False] = ..., 478 ) -> Any: ... 479 @overload 480 def median( 481 a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, 482 axis: None | _ShapeLike = ..., 483 out: None = ..., 484 overwrite_input: bool = ..., 485 keepdims: bool = ..., 486 ) -> Any: ... 487 @overload 488 def median( 489 a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, 490 axis: None | _ShapeLike = ..., 491 out: _ArrayType = ..., 492 overwrite_input: bool = ..., 493 keepdims: bool = ..., 494 ) -> _ArrayType: ... 495 496 _MethodKind = L[ 497 "inverted_cdf", 498 "averaged_inverted_cdf", 499 "closest_observation", 500 "interpolated_inverted_cdf", 501 "hazen", 502 "weibull", 503 "linear", 504 "median_unbiased", 505 "normal_unbiased", 506 "lower", 507 "higher", 508 "midpoint", 509 "nearest", 510 ] 511 512 @overload 513 def percentile( 514 a: _ArrayLikeFloat_co, 515 q: _FloatLike_co, 516 axis: None = ..., 517 out: None = ..., 518 overwrite_input: bool = ..., 519 method: _MethodKind = ..., 520 keepdims: L[False] = ..., 521 ) -> floating[Any]: ... 522 @overload 523 def percentile( 524 a: _ArrayLikeComplex_co, 525 q: _FloatLike_co, 526 axis: None = ..., 527 out: None = ..., 528 overwrite_input: bool = ..., 529 method: _MethodKind = ..., 530 keepdims: L[False] = ..., 531 ) -> complexfloating[Any, Any]: ... 532 @overload 533 def percentile( 534 a: _ArrayLikeTD64_co, 535 q: _FloatLike_co, 536 axis: None = ..., 537 out: None = ..., 538 overwrite_input: bool = ..., 539 method: _MethodKind = ..., 540 keepdims: L[False] = ..., 541 ) -> timedelta64: ... 542 @overload 543 def percentile( 544 a: _ArrayLikeDT64_co, 545 q: _FloatLike_co, 546 axis: None = ..., 547 out: None = ..., 548 overwrite_input: bool = ..., 549 method: _MethodKind = ..., 550 keepdims: L[False] = ..., 551 ) -> datetime64: ... 552 @overload 553 def percentile( 554 a: _ArrayLikeObject_co, 555 q: _FloatLike_co, 556 axis: None = ..., 557 out: None = ..., 558 overwrite_input: bool = ..., 559 method: _MethodKind = ..., 560 keepdims: L[False] = ..., 561 ) -> Any: ... 562 @overload 563 def percentile( 564 a: _ArrayLikeFloat_co, 565 q: _ArrayLikeFloat_co, 566 axis: None = ..., 567 out: None = ..., 568 overwrite_input: bool = ..., 569 method: _MethodKind = ..., 570 keepdims: L[False] = ..., 571 ) -> NDArray[floating[Any]]: ... 572 @overload 573 def percentile( 574 a: _ArrayLikeComplex_co, 575 q: _ArrayLikeFloat_co, 576 axis: None = ..., 577 out: None = ..., 578 overwrite_input: bool = ..., 579 method: _MethodKind = ..., 580 keepdims: L[False] = ..., 581 ) -> NDArray[complexfloating[Any, Any]]: ... 582 @overload 583 def percentile( 584 a: _ArrayLikeTD64_co, 585 q: _ArrayLikeFloat_co, 586 axis: None = ..., 587 out: None = ..., 588 overwrite_input: bool = ..., 589 method: _MethodKind = ..., 590 keepdims: L[False] = ..., 591 ) -> NDArray[timedelta64]: ... 592 @overload 593 def percentile( 594 a: _ArrayLikeDT64_co, 595 q: _ArrayLikeFloat_co, 596 axis: None = ..., 597 out: None = ..., 598 overwrite_input: bool = ..., 599 method: _MethodKind = ..., 600 keepdims: L[False] = ..., 601 ) -> NDArray[datetime64]: ... 602 @overload 603 def percentile( 604 a: _ArrayLikeObject_co, 605 q: _ArrayLikeFloat_co, 606 axis: None = ..., 607 out: None = ..., 608 overwrite_input: bool = ..., 609 method: _MethodKind = ..., 610 keepdims: L[False] = ..., 611 ) -> NDArray[object_]: ... 612 @overload 613 def percentile( 614 a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, 615 q: _ArrayLikeFloat_co, 616 axis: None | _ShapeLike = ..., 617 out: None = ..., 618 overwrite_input: bool = ..., 619 method: _MethodKind = ..., 620 keepdims: bool = ..., 621 ) -> Any: ... 622 @overload 623 def percentile( 624 a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, 625 q: _ArrayLikeFloat_co, 626 axis: None | _ShapeLike = ..., 627 out: _ArrayType = ..., 628 overwrite_input: bool = ..., 629 method: _MethodKind = ..., 630 keepdims: bool = ..., 631 ) -> _ArrayType: ... 632 633 # NOTE: Not an alias, but they do have identical signatures 634 # (that we can reuse) 635 quantile = percentile 636 637 # TODO: Returns a scalar for <= 1D array-likes; returns an ndarray otherwise 638 def trapz( 639 y: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, 640 x: None | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co = ..., 641 dx: float = ..., 642 axis: SupportsIndex = ..., 643 ) -> Any: ... 644 645 def meshgrid( 646 *xi: ArrayLike, 647 copy: bool = ..., 648 sparse: bool = ..., 649 indexing: L["xy", "ij"] = ..., 650 ) -> list[NDArray[Any]]: ... 651 652 @overload 653 def delete( 654 arr: _ArrayLike[_SCT], 655 obj: slice | _ArrayLikeInt_co, 656 axis: None | SupportsIndex = ..., 657 ) -> NDArray[_SCT]: ... 658 @overload 659 def delete( 660 arr: ArrayLike, 661 obj: slice | _ArrayLikeInt_co, 662 axis: None | SupportsIndex = ..., 663 ) -> NDArray[Any]: ... 664 665 @overload 666 def insert( 667 arr: _ArrayLike[_SCT], 668 obj: slice | _ArrayLikeInt_co, 669 values: ArrayLike, 670 axis: None | SupportsIndex = ..., 671 ) -> NDArray[_SCT]: ... 672 @overload 673 def insert( 674 arr: ArrayLike, 675 obj: slice | _ArrayLikeInt_co, 676 values: ArrayLike, 677 axis: None | SupportsIndex = ..., 678 ) -> NDArray[Any]: ... 679 680 def append( 681 arr: ArrayLike, 682 values: ArrayLike, 683 axis: None | SupportsIndex = ..., 684 ) -> NDArray[Any]: ... 685 686 @overload 687 def digitize( 688 x: _FloatLike_co, 689 bins: _ArrayLikeFloat_co, 690 right: bool = ..., 691 ) -> intp: ... 692 @overload 693 def digitize( 694 x: _ArrayLikeFloat_co, 695 bins: _ArrayLikeFloat_co, 696 right: bool = ..., 697 ) -> NDArray[intp]: ...