diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 7 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.py | 37 | ||||
| -rw-r--r-- | numpy/lib/twodim_base.pyi | 265 |
3 files changed, 283 insertions, 26 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 783d45c2f..2e9ae6644 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -88,8 +88,11 @@ def rot90(m, k=1, axes=(0, 1)): Notes ----- - rot90(m, k=1, axes=(1,0)) is the reverse of rot90(m, k=1, axes=(0,1)) - rot90(m, k=1, axes=(1,0)) is equivalent to rot90(m, k=-1, axes=(0,1)) + ``rot90(m, k=1, axes=(1,0))`` is the reverse of + ``rot90(m, k=1, axes=(0,1))`` + + ``rot90(m, k=1, axes=(1,0))`` is equivalent to + ``rot90(m, k=-1, axes=(0,1))`` Examples -------- diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index ac5fe1c6b..83c028061 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -439,10 +439,12 @@ def tril(m, k=0): Lower triangle of an array. Return a copy of an array with elements above the `k`-th diagonal zeroed. + For arrays with ``ndim`` exceeding 2, `tril` will apply to the final two + axes. Parameters ---------- - m : array_like, shape (M, N) + m : array_like, shape (..., M, N) Input array. k : int, optional Diagonal above which to zero elements. `k = 0` (the default) is the @@ -450,7 +452,7 @@ def tril(m, k=0): Returns ------- - tril : ndarray, shape (M, N) + tril : ndarray, shape (..., M, N) Lower triangle of `m`, of same shape and data-type as `m`. See Also @@ -465,6 +467,20 @@ def tril(m, k=0): [ 7, 8, 0], [10, 11, 12]]) + >>> np.tril(np.arange(3*4*5).reshape(3, 4, 5)) + array([[[ 0, 0, 0, 0, 0], + [ 5, 6, 0, 0, 0], + [10, 11, 12, 0, 0], + [15, 16, 17, 18, 0]], + [[20, 0, 0, 0, 0], + [25, 26, 0, 0, 0], + [30, 31, 32, 0, 0], + [35, 36, 37, 38, 0]], + [[40, 0, 0, 0, 0], + [45, 46, 0, 0, 0], + [50, 51, 52, 0, 0], + [55, 56, 57, 58, 0]]]) + """ m = asanyarray(m) mask = tri(*m.shape[-2:], k=k, dtype=bool) @@ -478,7 +494,8 @@ def triu(m, k=0): Upper triangle of an array. Return a copy of an array with the elements below the `k`-th diagonal - zeroed. + zeroed. For arrays with ``ndim`` exceeding 2, `triu` will apply to the final + two axes. Please refer to the documentation for `tril` for further details. @@ -494,6 +511,20 @@ def triu(m, k=0): [ 0, 8, 9], [ 0, 0, 12]]) + >>> np.triu(np.arange(3*4*5).reshape(3, 4, 5)) + array([[[ 0, 1, 2, 3, 4], + [ 0, 6, 7, 8, 9], + [ 0, 0, 12, 13, 14], + [ 0, 0, 0, 18, 19]], + [[20, 21, 22, 23, 24], + [ 0, 26, 27, 28, 29], + [ 0, 0, 32, 33, 34], + [ 0, 0, 0, 38, 39]], + [[40, 41, 42, 43, 44], + [ 0, 46, 47, 48, 49], + [ 0, 0, 52, 53, 54], + [ 0, 0, 0, 58, 59]]]) + """ m = asanyarray(m) mask = tri(*m.shape[-2:], k=k-1, dtype=bool) diff --git a/numpy/lib/twodim_base.pyi b/numpy/lib/twodim_base.pyi index 79b9511b8..007338d77 100644 --- a/numpy/lib/twodim_base.pyi +++ b/numpy/lib/twodim_base.pyi @@ -1,32 +1,255 @@ -from typing import List, Optional, Any +from typing import ( + Any, + Callable, + List, + Sequence, + overload, + Tuple, + Type, + TypeVar, + Union, +) -from numpy import ndarray, _OrderCF -from numpy.typing import ArrayLike, DTypeLike +from numpy import ( + ndarray, + dtype, + generic, + number, + bool_, + timedelta64, + datetime64, + int_, + intp, + float64, + signedinteger, + floating, + complexfloating, + object_, + _OrderCF, +) + +from numpy.typing import ( + DTypeLike, + _SupportsDType, + ArrayLike, + NDArray, + _NestedSequence, + _SupportsArray, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, +) + +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=generic) + +# The returned arrays dtype must be compatible with `np.equal` +_MaskFunc = Callable[ + [NDArray[int_], _T], + NDArray[Union[number[Any], bool_, timedelta64, datetime64, object_]], +] + +_DTypeLike = Union[ + Type[_SCT], + dtype[_SCT], + _SupportsDType[dtype[_SCT]], +] +_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]] __all__: List[str] -def fliplr(m): ... -def flipud(m): ... +@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: Optional[int] = ..., + M: None | int = ..., + k: int = ..., + dtype: None = ..., + order: _OrderCF = ..., + *, + like: None | ArrayLike = ..., +) -> NDArray[float64]: ... +@overload +def eye( + N: int, + M: None | int = ..., + k: int = ..., + dtype: _DTypeLike[_SCT] = ..., + order: _OrderCF = ..., + *, + like: None | ArrayLike = ..., +) -> NDArray[_SCT]: ... +@overload +def eye( + N: int, + M: None | int = ..., k: int = ..., dtype: DTypeLike = ..., order: _OrderCF = ..., *, - like: Optional[ArrayLike] = ... -) -> ndarray[Any, Any]: ... - -def diag(v, k=...): ... -def diagflat(v, k=...): ... -def tri(N, M=..., k=..., dtype = ..., *, like=...): ... -def tril(m, k=...): ... -def triu(m, k=...): ... -def vander(x, N=..., increasing=...): ... -def histogram2d(x, y, bins=..., range=..., normed=..., weights=..., density=...): ... -def mask_indices(n, mask_func, k=...): ... -def tril_indices(n, k=..., m=...): ... -def tril_indices_from(arr, k=...): ... -def triu_indices(n, k=..., m=...): ... -def triu_indices_from(arr, k=...): ... + like: None | ArrayLike = ..., +) -> 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 | ArrayLike = ... +) -> NDArray[float64]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: _DTypeLike[_SCT] = ..., + *, + like: None | ArrayLike = ... +) -> NDArray[_SCT]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: DTypeLike = ..., + *, + like: None | ArrayLike = ... +) -> NDArray[Any]: ... + +@overload +def tril(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def tril(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def triu(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def triu(v: ArrayLike, k: int = ...) -> 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( # type: ignore[misc] + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + bins: int | Sequence[int] = ..., + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[floating[Any]], + NDArray[floating[Any]], +]: ... +@overload +def histogram2d( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + bins: int | Sequence[int] = ..., + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[complexfloating[Any, Any]], + NDArray[complexfloating[Any, Any]], +]: ... +@overload # TODO: Sort out `bins` +def histogram2d( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + bins: Sequence[_ArrayLikeInt_co], + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[Any], + NDArray[Any], +]: ... + +# 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_]]: ... |