358 lines
7.8 KiB
Python
358 lines
7.8 KiB
Python
|
import sys
|
||
|
|
import datetime as dt
|
||
|
|
from typing import Optional, Union, Sequence, Tuple, Any, overload, TypeVar
|
||
|
|
|
||
|
|
from numpy import (
|
||
|
|
ndarray,
|
||
|
|
signedinteger,
|
||
|
|
bool_,
|
||
|
|
generic,
|
||
|
|
_OrderKACF,
|
||
|
|
_OrderACF,
|
||
|
|
_ArrayLikeBool,
|
||
|
|
_ArrayLikeIntOrBool,
|
||
|
|
_ModeKind,
|
||
|
|
_PartitionKind,
|
||
|
|
_SortKind,
|
||
|
|
_SortSide,
|
||
|
|
)
|
||
|
|
from numpy.typing import (
|
||
|
|
DTypeLike,
|
||
|
|
ArrayLike,
|
||
|
|
_ShapeLike,
|
||
|
|
_Shape,
|
||
|
|
_NumberLike,
|
||
|
|
)
|
||
|
|
|
||
|
|
if sys.version_info >= (3, 8):
|
||
|
|
from typing import Literal
|
||
|
|
else:
|
||
|
|
from typing_extensions import Literal
|
||
|
|
|
||
|
|
# Various annotations for scalars
|
||
|
|
|
||
|
|
# While dt.datetime and dt.timedelta are not technically part of NumPy,
|
||
|
|
# they are one of the rare few builtin scalars which serve as valid return types.
|
||
|
|
# See https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113.
|
||
|
|
_ScalarNumpy = Union[generic, dt.datetime, dt.timedelta]
|
||
|
|
_ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex]
|
||
|
|
_Scalar = Union[_ScalarBuiltin, _ScalarNumpy]
|
||
|
|
|
||
|
|
# Integers and booleans can generally be used interchangeably
|
||
|
|
_ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic)
|
||
|
|
|
||
|
|
# The signature of take() follows a common theme with its overloads:
|
||
|
|
# 1. A generic comes in; the same generic comes out
|
||
|
|
# 2. A scalar comes in; a generic comes out
|
||
|
|
# 3. An array-like object comes in; some keyword ensures that a generic comes out
|
||
|
|
# 4. An array-like object comes in; an ndarray or generic comes out
|
||
|
|
def take(
|
||
|
|
a: ArrayLike,
|
||
|
|
indices: _ArrayLikeIntOrBool,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
mode: _ModeKind = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def reshape(
|
||
|
|
a: ArrayLike,
|
||
|
|
newshape: _ShapeLike,
|
||
|
|
order: _OrderACF = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def choose(
|
||
|
|
a: _ArrayLikeIntOrBool,
|
||
|
|
choices: ArrayLike,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
mode: _ModeKind = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def repeat(
|
||
|
|
a: ArrayLike,
|
||
|
|
repeats: _ArrayLikeIntOrBool,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def put(
|
||
|
|
a: ndarray,
|
||
|
|
ind: _ArrayLikeIntOrBool,
|
||
|
|
v: ArrayLike,
|
||
|
|
mode: _ModeKind = ...,
|
||
|
|
) -> None: ...
|
||
|
|
|
||
|
|
def swapaxes(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis1: int,
|
||
|
|
axis2: int,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def transpose(
|
||
|
|
a: ArrayLike,
|
||
|
|
axes: Union[None, Sequence[int], ndarray] = ...
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def partition(
|
||
|
|
a: ArrayLike,
|
||
|
|
kth: _ArrayLikeIntOrBool,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
kind: _PartitionKind = ...,
|
||
|
|
order: Union[None, str, Sequence[str]] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def argpartition(
|
||
|
|
a: ArrayLike,
|
||
|
|
kth: _ArrayLikeIntOrBool,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
kind: _PartitionKind = ...,
|
||
|
|
order: Union[None, str, Sequence[str]] = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def sort(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
kind: Optional[_SortKind] = ...,
|
||
|
|
order: Union[None, str, Sequence[str]] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def argsort(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
kind: Optional[_SortKind] = ...,
|
||
|
|
order: Union[None, str, Sequence[str]] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def argmax(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: None = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> signedinteger[Any]: ...
|
||
|
|
@overload
|
||
|
|
def argmax(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def argmin(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: None = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> signedinteger[Any]: ...
|
||
|
|
@overload
|
||
|
|
def argmin(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def searchsorted(
|
||
|
|
a: ArrayLike,
|
||
|
|
v: _Scalar,
|
||
|
|
side: _SortSide = ...,
|
||
|
|
sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array
|
||
|
|
) -> signedinteger[Any]: ...
|
||
|
|
@overload
|
||
|
|
def searchsorted(
|
||
|
|
a: ArrayLike,
|
||
|
|
v: ArrayLike,
|
||
|
|
side: _SortSide = ...,
|
||
|
|
sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def resize(
|
||
|
|
a: ArrayLike,
|
||
|
|
new_shape: _ShapeLike,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def squeeze(
|
||
|
|
a: _ScalarGeneric,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
) -> _ScalarGeneric: ...
|
||
|
|
@overload
|
||
|
|
def squeeze(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def diagonal(
|
||
|
|
a: ArrayLike,
|
||
|
|
offset: int = ...,
|
||
|
|
axis1: int = ...,
|
||
|
|
axis2: int = ..., # >= 2D array
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def trace(
|
||
|
|
a: ArrayLike, # >= 2D array
|
||
|
|
offset: int = ...,
|
||
|
|
axis1: int = ...,
|
||
|
|
axis2: int = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def ravel(a: ArrayLike, order: _OrderKACF = ...) -> ndarray: ...
|
||
|
|
|
||
|
|
def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ...
|
||
|
|
|
||
|
|
def shape(a: ArrayLike) -> _Shape: ...
|
||
|
|
|
||
|
|
def compress(
|
||
|
|
condition: ArrayLike, # 1D bool array
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def clip(
|
||
|
|
a: ArrayLike,
|
||
|
|
a_min: ArrayLike,
|
||
|
|
a_max: Optional[ArrayLike],
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> Any: ...
|
||
|
|
@overload
|
||
|
|
def clip(
|
||
|
|
a: ArrayLike,
|
||
|
|
a_min: None,
|
||
|
|
a_max: ArrayLike,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def sum(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: _ShapeLike = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
initial: _NumberLike = ...,
|
||
|
|
where: _ArrayLikeBool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def all(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: None = ...,
|
||
|
|
out: None = ...,
|
||
|
|
keepdims: Literal[False] = ...,
|
||
|
|
) -> bool_: ...
|
||
|
|
@overload
|
||
|
|
def all(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def any(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: None = ...,
|
||
|
|
out: None = ...,
|
||
|
|
keepdims: Literal[False] = ...,
|
||
|
|
) -> bool_: ...
|
||
|
|
@overload
|
||
|
|
def any(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def cumsum(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def ptp(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def amax(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
initial: _NumberLike = ...,
|
||
|
|
where: _ArrayLikeBool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def amin(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
initial: _NumberLike = ...,
|
||
|
|
where: _ArrayLikeBool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
# TODO: `np.prod()``: For object arrays `initial` does not necessarily
|
||
|
|
# have to be a numerical scalar.
|
||
|
|
# The only requirement is that it is compatible
|
||
|
|
# with the `.__mul__()` method(s) of the passed array's elements.
|
||
|
|
|
||
|
|
# Note that the same situation holds for all wrappers around
|
||
|
|
# `np.ufunc.reduce`, e.g. `np.sum()` (`.__add__()`).
|
||
|
|
def prod(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
initial: _NumberLike = ...,
|
||
|
|
where: _ArrayLikeBool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def cumprod(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[int] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> ndarray: ...
|
||
|
|
|
||
|
|
def ndim(a: ArrayLike) -> int: ...
|
||
|
|
|
||
|
|
def size(a: ArrayLike, axis: Optional[int] = ...) -> int: ...
|
||
|
|
|
||
|
|
def around(
|
||
|
|
a: ArrayLike,
|
||
|
|
decimals: int = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def mean(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def std(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
ddof: int = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|
||
|
|
|
||
|
|
def var(
|
||
|
|
a: ArrayLike,
|
||
|
|
axis: Optional[_ShapeLike] = ...,
|
||
|
|
dtype: DTypeLike = ...,
|
||
|
|
out: Optional[ndarray] = ...,
|
||
|
|
ddof: int = ...,
|
||
|
|
keepdims: bool = ...,
|
||
|
|
) -> Any: ...
|