| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437 |
- from collections.abc import Callable, Sequence
- from typing import (
- Any,
- TypeAlias,
- overload,
- TypeVar,
- Literal as L,
- )
- import numpy as np
- from numpy import (
- generic,
- timedelta64,
- datetime64,
- int_,
- intp,
- float64,
- complex128,
- signedinteger,
- floating,
- complexfloating,
- object_,
- _OrderCF,
- )
- from numpy._typing import (
- DTypeLike,
- _DTypeLike,
- ArrayLike,
- _ArrayLike,
- NDArray,
- _SupportsArray,
- _SupportsArrayFunc,
- _ArrayLikeInt_co,
- _ArrayLikeFloat_co,
- _ArrayLikeComplex_co,
- _ArrayLikeObject_co,
- )
- __all__ = [
- "diag",
- "diagflat",
- "eye",
- "fliplr",
- "flipud",
- "tri",
- "triu",
- "tril",
- "vander",
- "histogram2d",
- "mask_indices",
- "tril_indices",
- "tril_indices_from",
- "triu_indices",
- "triu_indices_from",
- ]
- ###
- _T = TypeVar("_T")
- _SCT = TypeVar("_SCT", bound=generic)
- _SCT_complex = TypeVar("_SCT_complex", bound=np.complexfloating)
- _SCT_inexact = TypeVar("_SCT_inexact", bound=np.inexact)
- _SCT_number_co = TypeVar("_SCT_number_co", bound=_Number_co)
- # The returned arrays dtype must be compatible with `np.equal`
- _MaskFunc: TypeAlias = Callable[[NDArray[int_], _T], NDArray[_Number_co | timedelta64 | datetime64 | object_]]
- _Int_co: TypeAlias = np.integer | np.bool
- _Float_co: TypeAlias = np.floating | _Int_co
- _Number_co: TypeAlias = np.number | np.bool
- _ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_SCT]] | Sequence[_SCT]
- _ArrayLike1DInt_co: TypeAlias = _SupportsArray[np.dtype[_Int_co]] | Sequence[int | _Int_co]
- _ArrayLike1DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[float | _Float_co]
- _ArrayLike2DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[_ArrayLike1DFloat_co]
- _ArrayLike1DNumber_co: TypeAlias = _SupportsArray[np.dtype[_Number_co]] | Sequence[complex | _Number_co]
- ###
- @overload
- def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
- @overload
- def fliplr(m: ArrayLike) -> NDArray[Any]: ...
- @overload
- def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
- @overload
- def flipud(m: ArrayLike) -> NDArray[Any]: ...
- @overload
- def eye(
- N: int,
- M: None | int = ...,
- k: int = ...,
- dtype: None = ...,
- order: _OrderCF = ...,
- *,
- device: None | L["cpu"] = ...,
- like: None | _SupportsArrayFunc = ...,
- ) -> NDArray[float64]: ...
- @overload
- def eye(
- N: int,
- M: None | int,
- k: int,
- dtype: _DTypeLike[_SCT],
- order: _OrderCF = ...,
- *,
- device: None | L["cpu"] = ...,
- like: None | _SupportsArrayFunc = ...,
- ) -> NDArray[_SCT]: ...
- @overload
- def eye(
- N: int,
- M: None | int = ...,
- k: int = ...,
- *,
- dtype: _DTypeLike[_SCT],
- order: _OrderCF = ...,
- device: None | L["cpu"] = ...,
- like: None | _SupportsArrayFunc = ...,
- ) -> NDArray[_SCT]: ...
- @overload
- def eye(
- N: int,
- M: None | int = ...,
- k: int = ...,
- dtype: DTypeLike = ...,
- order: _OrderCF = ...,
- *,
- device: None | L["cpu"] = ...,
- like: None | _SupportsArrayFunc = ...,
- ) -> NDArray[Any]: ...
- @overload
- def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
- @overload
- def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
- @overload
- def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ...
- @overload
- def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ...
- @overload
- def tri(
- N: int,
- M: None | int = ...,
- k: int = ...,
- dtype: None = ...,
- *,
- like: None | _SupportsArrayFunc = ...
- ) -> NDArray[float64]: ...
- @overload
- def tri(
- N: int,
- M: None | int,
- k: int,
- dtype: _DTypeLike[_SCT],
- *,
- like: None | _SupportsArrayFunc = ...
- ) -> NDArray[_SCT]: ...
- @overload
- def tri(
- N: int,
- M: None | int = ...,
- k: int = ...,
- *,
- dtype: _DTypeLike[_SCT],
- like: None | _SupportsArrayFunc = ...
- ) -> NDArray[_SCT]: ...
- @overload
- def tri(
- N: int,
- M: None | int = ...,
- k: int = ...,
- dtype: DTypeLike = ...,
- *,
- like: None | _SupportsArrayFunc = ...
- ) -> NDArray[Any]: ...
- @overload
- def tril(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ...
- @overload
- def tril(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
- @overload
- def triu(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ...
- @overload
- def triu(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
- @overload
- def vander( # type: ignore[misc]
- x: _ArrayLikeInt_co,
- N: None | int = ...,
- increasing: bool = ...,
- ) -> NDArray[signedinteger[Any]]: ...
- @overload
- def vander( # type: ignore[misc]
- x: _ArrayLikeFloat_co,
- N: None | int = ...,
- increasing: bool = ...,
- ) -> NDArray[floating[Any]]: ...
- @overload
- def vander(
- x: _ArrayLikeComplex_co,
- N: None | int = ...,
- increasing: bool = ...,
- ) -> NDArray[complexfloating[Any, Any]]: ...
- @overload
- def vander(
- x: _ArrayLikeObject_co,
- N: None | int = ...,
- increasing: bool = ...,
- ) -> NDArray[object_]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1D[_SCT_complex],
- y: _ArrayLike1D[_SCT_complex | _Float_co],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_complex],
- NDArray[_SCT_complex],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1D[_SCT_complex | _Float_co],
- y: _ArrayLike1D[_SCT_complex],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_complex],
- NDArray[_SCT_complex],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1D[_SCT_inexact],
- y: _ArrayLike1D[_SCT_inexact | _Int_co],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_inexact],
- NDArray[_SCT_inexact],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1D[_SCT_inexact | _Int_co],
- y: _ArrayLike1D[_SCT_inexact],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_inexact],
- NDArray[_SCT_inexact],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DInt_co | Sequence[float | int],
- y: _ArrayLike1DInt_co | Sequence[float | int],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[float64],
- NDArray[float64],
- ]: ...
- @overload
- def histogram2d(
- x: Sequence[complex | float | int],
- y: Sequence[complex | float | int],
- bins: int | Sequence[int] = ...,
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[complex128 | float64],
- NDArray[complex128 | float64],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DNumber_co,
- y: _ArrayLike1DNumber_co,
- bins: _ArrayLike1D[_SCT_number_co] | Sequence[_ArrayLike1D[_SCT_number_co]],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_number_co],
- NDArray[_SCT_number_co],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1D[_SCT_inexact],
- y: _ArrayLike1D[_SCT_inexact],
- bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_number_co | _SCT_inexact],
- NDArray[_SCT_number_co | _SCT_inexact],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DInt_co | Sequence[float | int],
- y: _ArrayLike1DInt_co | Sequence[float | int],
- bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_number_co | float64],
- NDArray[_SCT_number_co | float64],
- ]: ...
- @overload
- def histogram2d(
- x: Sequence[complex | float | int],
- y: Sequence[complex | float | int],
- bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[_SCT_number_co | complex128 | float64],
- NDArray[_SCT_number_co | complex128 | float64] ,
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DNumber_co,
- y: _ArrayLike1DNumber_co,
- bins: Sequence[Sequence[bool]],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[np.bool],
- NDArray[np.bool],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DNumber_co,
- y: _ArrayLike1DNumber_co,
- bins: Sequence[Sequence[int | bool]],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[np.int_ | np.bool],
- NDArray[np.int_ | np.bool],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DNumber_co,
- y: _ArrayLike1DNumber_co,
- bins: Sequence[Sequence[float | int | bool]],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[np.float64 | np.int_ | np.bool],
- NDArray[np.float64 | np.int_ | np.bool],
- ]: ...
- @overload
- def histogram2d(
- x: _ArrayLike1DNumber_co,
- y: _ArrayLike1DNumber_co,
- bins: Sequence[Sequence[complex | float | int | bool]],
- range: None | _ArrayLike2DFloat_co = ...,
- density: None | bool = ...,
- weights: None | _ArrayLike1DFloat_co = ...,
- ) -> tuple[
- NDArray[float64],
- NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
- NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
- ]: ...
- # NOTE: we're assuming/demanding here the `mask_func` returns
- # an ndarray of shape `(n, n)`; otherwise there is the possibility
- # of the output tuple having more or less than 2 elements
- @overload
- def mask_indices(
- n: int,
- mask_func: _MaskFunc[int],
- k: int = ...,
- ) -> tuple[NDArray[intp], NDArray[intp]]: ...
- @overload
- def mask_indices(
- n: int,
- mask_func: _MaskFunc[_T],
- k: _T,
- ) -> tuple[NDArray[intp], NDArray[intp]]: ...
- def tril_indices(
- n: int,
- k: int = ...,
- m: None | int = ...,
- ) -> tuple[NDArray[int_], NDArray[int_]]: ...
- def tril_indices_from(
- arr: NDArray[Any],
- k: int = ...,
- ) -> tuple[NDArray[int_], NDArray[int_]]: ...
- def triu_indices(
- n: int,
- k: int = ...,
- m: None | int = ...,
- ) -> tuple[NDArray[int_], NDArray[int_]]: ...
- def triu_indices_from(
- arr: NDArray[Any],
- k: int = ...,
- ) -> tuple[NDArray[int_], NDArray[int_]]: ...
|