298 lines
8.7 KiB
Python
298 lines
8.7 KiB
Python
from _typeshed import Incomplete
|
|
from collections.abc import Sequence
|
|
from typing import SupportsIndex, TypeAlias, TypeVar, overload
|
|
|
|
import numpy as np
|
|
from numpy import _CastingKind
|
|
from numpy._typing import (
|
|
ArrayLike,
|
|
DTypeLike,
|
|
_AnyShape,
|
|
_ArrayLike,
|
|
_DTypeLike,
|
|
_ShapeLike,
|
|
)
|
|
from numpy.lib._function_base_impl import average
|
|
from numpy.lib._index_tricks_impl import AxisConcatenator
|
|
|
|
from .core import MaskedArray, dot
|
|
|
|
__all__ = [
|
|
"apply_along_axis",
|
|
"apply_over_axes",
|
|
"atleast_1d",
|
|
"atleast_2d",
|
|
"atleast_3d",
|
|
"average",
|
|
"clump_masked",
|
|
"clump_unmasked",
|
|
"column_stack",
|
|
"compress_cols",
|
|
"compress_nd",
|
|
"compress_rowcols",
|
|
"compress_rows",
|
|
"corrcoef",
|
|
"count_masked",
|
|
"cov",
|
|
"diagflat",
|
|
"dot",
|
|
"dstack",
|
|
"ediff1d",
|
|
"flatnotmasked_contiguous",
|
|
"flatnotmasked_edges",
|
|
"hsplit",
|
|
"hstack",
|
|
"in1d",
|
|
"intersect1d",
|
|
"isin",
|
|
"mask_cols",
|
|
"mask_rowcols",
|
|
"mask_rows",
|
|
"masked_all",
|
|
"masked_all_like",
|
|
"median",
|
|
"mr_",
|
|
"ndenumerate",
|
|
"notmasked_contiguous",
|
|
"notmasked_edges",
|
|
"polyfit",
|
|
"row_stack",
|
|
"setdiff1d",
|
|
"setxor1d",
|
|
"stack",
|
|
"union1d",
|
|
"unique",
|
|
"vander",
|
|
"vstack",
|
|
]
|
|
|
|
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
|
|
_ScalarT1 = TypeVar("_ScalarT1", bound=np.generic)
|
|
_ScalarT2 = TypeVar("_ScalarT2", bound=np.generic)
|
|
_MArrayT = TypeVar("_MArrayT", bound=MaskedArray)
|
|
|
|
_MArray: TypeAlias = MaskedArray[_AnyShape, np.dtype[_ScalarT]]
|
|
|
|
###
|
|
|
|
# keep in sync with `numpy._core.shape_base.atleast_1d`
|
|
@overload
|
|
def atleast_1d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def atleast_1d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
|
|
@overload
|
|
def atleast_1d(
|
|
a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
|
|
) -> tuple[_MArray[_ScalarT], ...]: ...
|
|
@overload
|
|
def atleast_1d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
|
|
@overload
|
|
def atleast_1d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
|
|
@overload
|
|
def atleast_1d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base.atleast_2d`
|
|
@overload
|
|
def atleast_2d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def atleast_2d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
|
|
@overload
|
|
def atleast_2d(
|
|
a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
|
|
) -> tuple[_MArray[_ScalarT], ...]: ...
|
|
@overload
|
|
def atleast_2d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
|
|
@overload
|
|
def atleast_2d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
|
|
@overload
|
|
def atleast_2d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base.atleast_2d`
|
|
@overload
|
|
def atleast_3d(a0: _ArrayLike[_ScalarT], /) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def atleast_3d(a0: _ArrayLike[_ScalarT1], a1: _ArrayLike[_ScalarT2], /) -> tuple[_MArray[_ScalarT1], _MArray[_ScalarT2]]: ...
|
|
@overload
|
|
def atleast_3d(
|
|
a0: _ArrayLike[_ScalarT], a1: _ArrayLike[_ScalarT], /, *arys: _ArrayLike[_ScalarT]
|
|
) -> tuple[_MArray[_ScalarT], ...]: ...
|
|
@overload
|
|
def atleast_3d(a0: ArrayLike, /) -> _MArray[Incomplete]: ...
|
|
@overload
|
|
def atleast_3d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[_MArray[Incomplete], _MArray[Incomplete]]: ...
|
|
@overload
|
|
def atleast_3d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[_MArray[Incomplete], ...]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base.vstack`
|
|
@overload
|
|
def vstack(
|
|
tup: Sequence[_ArrayLike[_ScalarT]],
|
|
*,
|
|
dtype: None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def vstack(
|
|
tup: Sequence[ArrayLike],
|
|
*,
|
|
dtype: _DTypeLike[_ScalarT],
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def vstack(
|
|
tup: Sequence[ArrayLike],
|
|
*,
|
|
dtype: DTypeLike | None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[Incomplete]: ...
|
|
|
|
row_stack = vstack
|
|
|
|
# keep in sync with `numpy._core.shape_base.hstack`
|
|
@overload
|
|
def hstack(
|
|
tup: Sequence[_ArrayLike[_ScalarT]],
|
|
*,
|
|
dtype: None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def hstack(
|
|
tup: Sequence[ArrayLike],
|
|
*,
|
|
dtype: _DTypeLike[_ScalarT],
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def hstack(
|
|
tup: Sequence[ArrayLike],
|
|
*,
|
|
dtype: DTypeLike | None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[Incomplete]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base_impl.column_stack`
|
|
@overload
|
|
def column_stack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def column_stack(tup: Sequence[ArrayLike]) -> _MArray[Incomplete]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base_impl.dstack`
|
|
@overload
|
|
def dstack(tup: Sequence[_ArrayLike[_ScalarT]]) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def dstack(tup: Sequence[ArrayLike]) -> _MArray[Incomplete]: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base.stack`
|
|
@overload
|
|
def stack(
|
|
arrays: Sequence[_ArrayLike[_ScalarT]],
|
|
axis: SupportsIndex = 0,
|
|
out: None = None,
|
|
*,
|
|
dtype: None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def stack(
|
|
arrays: Sequence[ArrayLike],
|
|
axis: SupportsIndex = 0,
|
|
out: None = None,
|
|
*,
|
|
dtype: _DTypeLike[_ScalarT],
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def stack(
|
|
arrays: Sequence[ArrayLike],
|
|
axis: SupportsIndex = 0,
|
|
out: None = None,
|
|
*,
|
|
dtype: DTypeLike | None = None,
|
|
casting: _CastingKind = "same_kind"
|
|
) -> _MArray[Incomplete]: ...
|
|
@overload
|
|
def stack(
|
|
arrays: Sequence[ArrayLike],
|
|
axis: SupportsIndex,
|
|
out: _MArrayT,
|
|
*,
|
|
dtype: DTypeLike | None = None,
|
|
casting: _CastingKind = "same_kind",
|
|
) -> _MArrayT: ...
|
|
@overload
|
|
def stack(
|
|
arrays: Sequence[ArrayLike],
|
|
axis: SupportsIndex = 0,
|
|
*,
|
|
out: _MArrayT,
|
|
dtype: DTypeLike | None = None,
|
|
casting: _CastingKind = "same_kind",
|
|
) -> _MArrayT: ...
|
|
|
|
# keep in sync with `numpy._core.shape_base_impl.hsplit`
|
|
@overload
|
|
def hsplit(ary: _ArrayLike[_ScalarT], indices_or_sections: _ShapeLike) -> list[_MArray[_ScalarT]]: ...
|
|
@overload
|
|
def hsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[_MArray[Incomplete]]: ...
|
|
|
|
# keep in sync with `numpy._core.twodim_base_impl.hsplit`
|
|
@overload
|
|
def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> _MArray[_ScalarT]: ...
|
|
@overload
|
|
def diagflat(v: ArrayLike, k: int = 0) -> _MArray[Incomplete]: ...
|
|
|
|
# TODO: everything below
|
|
|
|
def count_masked(arr, axis=None): ...
|
|
def masked_all(shape, dtype=float): ... # noqa: PYI014
|
|
def masked_all_like(arr): ...
|
|
|
|
def apply_along_axis(func1d, axis, arr, *args, **kwargs): ...
|
|
def apply_over_axes(func, a, axes): ...
|
|
def median(a, axis=None, out=None, overwrite_input=False, keepdims=False): ...
|
|
def compress_nd(x, axis=None): ...
|
|
def compress_rowcols(x, axis=None): ...
|
|
def compress_rows(a): ...
|
|
def compress_cols(a): ...
|
|
def mask_rows(a, axis=...): ...
|
|
def mask_cols(a, axis=...): ...
|
|
def ediff1d(arr, to_end=None, to_begin=None): ...
|
|
def unique(ar1, return_index=False, return_inverse=False): ...
|
|
def intersect1d(ar1, ar2, assume_unique=False): ...
|
|
def setxor1d(ar1, ar2, assume_unique=False): ...
|
|
def in1d(ar1, ar2, assume_unique=False, invert=False): ...
|
|
def isin(element, test_elements, assume_unique=False, invert=False): ...
|
|
def union1d(ar1, ar2): ...
|
|
def setdiff1d(ar1, ar2, assume_unique=False): ...
|
|
def cov(x, y=None, rowvar=True, bias=False, allow_masked=True, ddof=None): ...
|
|
def corrcoef(x, y=None, rowvar=True, allow_masked=True): ...
|
|
|
|
class MAxisConcatenator(AxisConcatenator):
|
|
__slots__ = ()
|
|
|
|
@staticmethod
|
|
def concatenate(arrays: Incomplete, axis: int = 0) -> Incomplete: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
|
|
@classmethod
|
|
def makemat(cls, arr: Incomplete) -> Incomplete: ... # type: ignore[override] # pyright: ignore[reportIncompatibleVariableOverride]
|
|
|
|
class mr_class(MAxisConcatenator):
|
|
__slots__ = ()
|
|
|
|
def __init__(self) -> None: ...
|
|
|
|
mr_: mr_class
|
|
|
|
def ndenumerate(a, compressed=True): ...
|
|
def flatnotmasked_edges(a): ...
|
|
def notmasked_edges(a, axis=None): ...
|
|
def flatnotmasked_contiguous(a): ...
|
|
def notmasked_contiguous(a, axis=None): ...
|
|
def clump_unmasked(a): ...
|
|
def clump_masked(a): ...
|
|
def vander(x, n=None): ...
|
|
def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False): ...
|
|
|
|
#
|
|
def mask_rowcols(a: Incomplete, axis: Incomplete | None = None) -> MaskedArray[Incomplete, np.dtype[Incomplete]]: ...
|