Skip to content

Commit

Permalink
✨ sparse: complete the conrete sparse array- and matrix-types
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorenham
jorenham committed Dec 13, 2024
1 parent 26857bc commit 7cac054
Show file tree
Hide file tree
Showing 14 changed files with 902 additions and 212 deletions.
2 changes: 2 additions & 0 deletions scipy-stubs/_typing.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ __all__ = [
"FileModeRWA",
"FileName",
"NanPolicy",
"OrderCF",
"OrderKACF",
"Seed",
"Untyped",
Expand Down Expand Up @@ -72,6 +73,7 @@ AnyShape: TypeAlias = op.CanIndex | Sequence[op.CanIndex]
RNG: TypeAlias = np.random.BitGenerator | np.random.Generator | np.random.RandomState
Seed: TypeAlias = RNG | np.random.SeedSequence | onp.ToJustInt | onp.ToJustIntND
ByteOrder: TypeAlias = Literal["S", "<", "little", ">", "big", "=", "native", "|", "I"]
OrderCF: TypeAlias = Literal["C", "F"]
OrderKACF: TypeAlias = Literal["K", "A", "C", "F"]
Casting: TypeAlias = Literal["no", "equiv", "safe", "same_kind", "unsafe"]
CorrelateMode: TypeAlias = Literal["valid", "same", "full"]
Expand Down
162 changes: 118 additions & 44 deletions scipy-stubs/sparse/_base.pyi
View file
Original file line number Diff line number Diff line change
@@ -1,38 +1,39 @@
# mypy: disable-error-code="misc"
# pyright: reportUnannotatedClassAttribute=false

import abc
from collections.abc import Iterator, Sequence
from typing import Any, Final, Generic, Literal, TypeAlias, overload
from typing_extensions import Self, TypeAliasType, TypeIs, TypeVar
from typing import Final, Generic, Literal, TypeAlias, overload
from typing_extensions import Self, TypeIs, TypeVar

import numpy as np
import optype as op
import optype.numpy as onp
import optype.typing as opt
from scipy._typing import Untyped
from scipy._typing import Casting, OrderCF, Untyped
from ._bsr import bsr_array, bsr_matrix
from ._coo import coo_array, coo_matrix
from ._csc import csc_array, csc_matrix
from ._csr import csr_array, csr_matrix
from ._dia import dia_array, dia_matrix
from ._dok import dok_array, dok_matrix
from ._lil import lil_array, lil_matrix
from ._matrix import spmatrix as spmatrix
from ._typing import Float, Int, Scalar, SPFormat, ToShape

__all__ = ["SparseEfficiencyWarning", "SparseWarning", "issparse", "isspmatrix", "sparray"]

# NOTE: This is roughly speaking equivalent to `[u]int8 | [u]int16 | [u]int32 | [u]int64` (on most modern platforms)
_Int: TypeAlias = np.integer[Any]
# NOTE: NOT equivalent to `floating`! It's considered invalid to use `float16` in `scipy.sparse`.
_Float: TypeAlias = np.float32 | np.float64 | np.longdouble
# NOTE: This is (almost always) equivalent to `complex64 | complex128 | clongdouble`.
_Complex: TypeAlias = np.complexfloating[Any, Any]
# NOTE: Roughly speaking, this is equivalent to `number[Any] - float16` (the `-` denotes the set difference analogue for types)
# NOTE: This (almost always) matches `._sputils.supported_dtypes`
# NOTE: The `TypeAliasType` is used to avoid long error messages.
_Scalar = TypeAliasType("_Scalar", np.bool_ | _Int | _Float | _Complex)

_T = TypeVar("_T")
_SCT_co = TypeVar("_SCT_co", bound=_Scalar, default=_Scalar, covariant=True)
_OutT = TypeVar("_OutT", bound=onp.ArrayND)
_SCT = TypeVar("_SCT", bound=np.generic)
_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True)

_Matrix: TypeAlias = np.matrix[tuple[int, int], np.dtype[_SCT]]

_ToShape: TypeAlias = tuple[op.CanIndex] | tuple[op.CanIndex, op.CanIndex]
_ToSparseFromPy: TypeAlias = Sequence[Sequence[_T]] | Sequence[_T]
_ToSparseFromArrayLike: TypeAlias = onp.CanArrayND[_SCT_co] | _ToSparseFromPy[_SCT_co]

_Format: TypeAlias = Literal["bsr", "coo", "csc", "csr", "dia", "dok", "lil"]
_Real: TypeAlias = np.bool_ | Int | Float

###

Expand All @@ -42,8 +43,7 @@ class SparseWarning(Warning): ...
class SparseFormatWarning(SparseWarning): ...
class SparseEfficiencyWarning(SparseWarning): ...

# TODO(jorenham): Generic `shape` property
# TODO(jorenham): Generic `data` attribute
# TODO(jorenham): Generic `shape` and `data`
class _spbase(Generic[_SCT_co]):
__array_priority__: float = 10.1
maxprint: Final[int | None]
Expand All @@ -61,7 +61,7 @@ class _spbase(Generic[_SCT_co]):
# NOTE: At runtime this isn't abstract, but returns `und` instead.
@property
@abc.abstractmethod
def format(self, /) -> _Format: ...
def format(self, /) -> SPFormat: ...

#
@property
Expand All @@ -73,7 +73,7 @@ class _spbase(Generic[_SCT_co]):

# NOTE: In `scipy>=1.15.0` the `maxprint` param will become keyword-only.
@overload # shape
def __init__(self: _spbase[np.float64], /, arg1: _ToShape, maxprint: int | None = 50) -> None: ...
def __init__(self: _spbase[np.float64], /, arg1: ToShape, maxprint: int | None = 50) -> None: ...
@overload # sparse
def __init__(self, /, arg1: _spbase[_SCT_co], maxprint: int | None = 50) -> None: ...
@overload # dense array-like
Expand Down Expand Up @@ -108,13 +108,38 @@ class _spbase(Generic[_SCT_co]):

#
def __bool__(self, /) -> bool: ...
def __iter__(self, /) -> Iterator[Untyped]: ...

#
def __lt__(self, other: Untyped, /) -> Untyped: ...
def __gt__(self, other: Untyped, /) -> Untyped: ...
def __le__(self, other: Untyped, /) -> Untyped: ...
def __ge__(self, other: Untyped, /) -> Untyped: ...
@overload # {csc,csr}_array
def __iter__(self: csc_array | csr_array, /) -> Iterator[csr_array[_SCT_co]]: ...
@overload # {csc,csr}_matrix
def __iter__(self: csc_matrix | csr_matrix, /) -> Iterator[csr_matrix[_SCT_co]]: ...
@overload # dok_array
def __iter__(self: dok_array, /) -> Iterator[dok_array[_SCT_co]]: ...
@overload # dok_matrix
def __iter__(self: dok_matrix, /) -> Iterator[dok_matrix[_SCT_co]]: ...
@overload # dok_array
def __iter__(self: lil_array, /) -> Iterator[lil_array[_SCT_co]]: ...
@overload # dok_array
def __iter__(self: lil_matrix, /) -> Iterator[lil_matrix[_SCT_co]]: ...

#
@overload
def __lt__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ...
@overload
def __lt__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ...
@overload
def __gt__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ...
@overload
def __gt__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ...
@overload
def __le__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ...
@overload
def __le__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ...
@overload
def __ge__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ...
@overload
def __ge__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ...

#
def __neg__(self, /) -> Self: ...
Expand All @@ -131,14 +156,14 @@ class _spbase(Generic[_SCT_co]):
def __matmul__(self, other: Untyped, /) -> Untyped: ...
def __rmatmul__(self, other: Untyped, /) -> Untyped: ...
def __truediv__(self, other: Untyped, /) -> Untyped: ...
def __div__(self, other: Untyped, /) -> Untyped: ...
def __div__(self, other: Untyped, /) -> Untyped: ... # que?
def __pow__(self, other: Untyped, /) -> Untyped: ...

#
def nonzero(self, /) -> Untyped: ...
def nonzero(self, /) -> tuple[onp.Array1D[np.int32], onp.Array1D[np.int32]]: ...
def count_nonzero(self, /) -> int: ...
def conjugate(self, /, copy: bool = True) -> Self: ...
conj = conjugate
def conj(self, /, copy: bool = True) -> Self: ...
def transpose(self, /, axes: Untyped | None = None, copy: bool = False) -> Self: ...

#
Expand All @@ -149,31 +174,80 @@ class _spbase(Generic[_SCT_co]):
def power(self, /, n: Untyped, dtype: Untyped | None = None) -> Self: ...

#
def diagonal(self, /, k: int = 0) -> Untyped: ...
def trace(self, /, offset: int = 0) -> Untyped: ...
def sum(self, /, axis: Untyped | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...
def mean(self, /, axis: Untyped | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...
def diagonal(self, /, k: int = 0) -> onp.Array1D[_SCT_co]: ... # only if 2-d
def trace(self, /, offset: int = 0) -> _SCT_co: ...

#
def sum(self, /, axis: op.CanIndex | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...
def mean(self, /, axis: op.CanIndex | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...

#
def copy(self, /) -> Self: ...
def reshape(self, /, *args: Untyped, **kwargs: Untyped) -> Self: ...
def astype(self, /, dtype: Untyped, casting: str = "unsafe", copy: bool = True) -> Untyped: ...
def astype(self, /, dtype: Untyped, casting: Casting = "unsafe", copy: bool = True) -> Untyped: ...
def asformat(self, /, format: Untyped, copy: bool = False) -> Untyped: ...

#
def todense(self, /, order: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...
def toarray(self, /, order: Untyped | None = None, out: Untyped | None = None) -> Untyped: ...
def tocsr(self, /, copy: bool = False) -> Untyped: ...
def todok(self, /, copy: bool = False) -> Untyped: ...
def tocoo(self, /, copy: bool = False) -> Untyped: ...
def tolil(self, /, copy: bool = False) -> Untyped: ...
def todia(self, /, copy: bool = False) -> Untyped: ...
def tobsr(self, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> Untyped: ...
def tocsc(self, /, copy: bool = False) -> Untyped: ...
@overload # self: spmatrix, out: None
def todense(self: spmatrix, /, order: OrderCF | None = None, out: None = None) -> _Matrix[_SCT_co]: ...
@overload # self: spmatrix, out: array (positional)
def todense(self: spmatrix, /, order: OrderCF | None, out: onp.ArrayND[_SCT]) -> _Matrix[_SCT]: ...
@overload # self: spmatrix, out: array (keyword)
def todense(self: spmatrix, /, order: OrderCF | None = None, *, out: onp.ArrayND[_SCT]) -> _Matrix[_SCT]: ...
@overload # out: None
def todense(self, /, order: OrderCF | None = None, out: None = None) -> onp.Array2D[_SCT_co]: ...
@overload # out: array (positional)
def todense(self, /, order: OrderCF | None, out: _OutT) -> _OutT: ...
@overload # out: array (keyword)
def todense(self, /, order: OrderCF | None = None, *, out: _OutT) -> _OutT: ...

#
@overload # out: None
def toarray(self, /, order: OrderCF | None = None, out: None = None) -> onp.Array2D[_SCT_co]: ...
@overload # out: array (positional)
def toarray(self, /, order: OrderCF | None, out: _OutT) -> _OutT: ...
@overload # out: array (keyword)
def toarray(self, /, order: OrderCF | None = None, *, out: _OutT) -> _OutT: ...

#
@overload
def tobsr(self: sparray, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> bsr_array[_SCT_co]: ...
@overload
def tobsr(self: spmatrix, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> bsr_matrix[_SCT_co]: ...
#
@overload
def tocoo(self: sparray, /, copy: bool = False) -> coo_array[_SCT_co]: ...
@overload
def tocoo(self: spmatrix, /, copy: bool = False) -> coo_matrix[_SCT_co]: ...
#
@overload
def tocsc(self: sparray, /, copy: bool = False) -> csc_array[_SCT_co]: ...
@overload
def tocsc(self: spmatrix, /, copy: bool = False) -> csc_matrix[_SCT_co]: ...
#
@overload
def tocsr(self: sparray, /, copy: bool = False) -> csr_array[_SCT_co]: ...
@overload
def tocsr(self: spmatrix, /, copy: bool = False) -> csr_matrix[_SCT_co]: ...
#
@overload
def todia(self: sparray, /, copy: bool = False) -> dia_array[_SCT_co]: ...
@overload
def todia(self: spmatrix, /, copy: bool = False) -> dia_matrix[_SCT_co]: ...
#
@overload
def todok(self: sparray, /, copy: bool = False) -> dok_array[_SCT_co]: ...
@overload
def todok(self: spmatrix, /, copy: bool = False) -> dok_matrix[_SCT_co]: ...
#
@overload
def tolil(self: sparray, /, copy: bool = False) -> lil_array[_SCT_co]: ...
@overload
def tolil(self: spmatrix, /, copy: bool = False) -> lil_matrix[_SCT_co]: ...

#
def resize(self, /, shape: tuple[op.CanIndex, op.CanIndex]) -> None: ...
def setdiag(self, /, values: Untyped, k: int = 0) -> None: ...
def setdiag(self, /, values: onp.ToComplex1D, k: int = 0) -> None: ...

class sparray: ...

Expand Down
Loading

0 comments on commit 7cac054

Please sign in to comment.