diff --git a/scipy-stubs/sparse/linalg/_expm_multiply.pyi b/scipy-stubs/sparse/linalg/_expm_multiply.pyi index 25a3d0ae..2e003a5c 100644 --- a/scipy-stubs/sparse/linalg/_expm_multiply.pyi +++ b/scipy-stubs/sparse/linalg/_expm_multiply.pyi @@ -1,45 +1,115 @@ -from typing import Any, TypeVar, overload +from _typeshed import Incomplete +from typing import Any, Never, TypeAlias, TypeVar, overload import numpy as np import optype as op import optype.numpy as onp +import optype.numpy.compat as npc -from scipy.sparse._typing import Numeric -from scipy.sparse.linalg._interface import LinearOperator +from ._interface import LinearOperator +from scipy.sparse._base import _spbase, sparray __all__ = ["expm_multiply"] -_SCT = TypeVar("_SCT", bound=Numeric) +_ScalarT = TypeVar("_ScalarT", bound=npc.number | np.bool_) +_InexactT = TypeVar("_InexactT", bound=npc.inexact) +_ShapeT = TypeVar("_ShapeT", bound=tuple[Any, ...]) + +_ToLinearOperator: TypeAlias = LinearOperator[_ScalarT] | _spbase[_ScalarT, tuple[int, int]] | onp.ArrayND[_ScalarT] +_SparseOrDense: TypeAlias = sparray[_ScalarT, _ShapeT] | onp.ArrayND[_ScalarT, _ShapeT] + +_AsFloat64: TypeAlias = np.float64 | npc.integer | np.bool_ +_ToFloat64: TypeAlias = _AsFloat64 | np.float32 | np.float16 ### +@overload # workaround for mypy's and pyright's typing spec non-compliance regarding overloads +def expm_multiply( + A: _ToLinearOperator[_AsFloat64], + B: _SparseOrDense[_ToFloat64, tuple[Never] | tuple[Never, Never]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.ArrayND[np.float64]: ... +@overload +def expm_multiply( + A: _ToLinearOperator[_InexactT], + B: _SparseOrDense[_InexactT | npc.integer | np.bool_, tuple[Never] | tuple[Never, Never]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.ArrayND[_InexactT]: ... @overload # 1-d def expm_multiply( - A: LinearOperator[_SCT], - B: onp.Array1D[_SCT | np.integer[Any] | np.float16 | np.bool_], + A: _ToLinearOperator[_AsFloat64], + B: _SparseOrDense[_ToFloat64, tuple[int]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.Array1D[np.float64]: ... +@overload +def expm_multiply( + A: _ToLinearOperator[_InexactT], + B: _SparseOrDense[_InexactT | npc.integer | np.bool_, tuple[int]], start: onp.ToFloat | None = None, stop: onp.ToFloat | None = None, num: op.CanIndex | None = None, endpoint: bool | None = None, traceA: onp.ToComplex | None = None, -) -> onp.Array1D[_SCT]: ... +) -> onp.Array1D[_InexactT]: ... @overload # 2-d def expm_multiply( - A: LinearOperator[_SCT], - B: onp.Array2D[_SCT | np.integer[Any] | np.float16 | np.bool_], + A: _ToLinearOperator[_AsFloat64], + B: _SparseOrDense[_ToFloat64, tuple[int, int]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.Array2D[np.float64]: ... +@overload +def expm_multiply( + A: _ToLinearOperator[_InexactT], + B: _SparseOrDense[_InexactT | npc.integer | np.bool_, tuple[int, int]], start: onp.ToFloat | None = None, stop: onp.ToFloat | None = None, num: op.CanIndex | None = None, endpoint: bool | None = None, traceA: onp.ToComplex | None = None, -) -> onp.Array2D[_SCT]: ... +) -> onp.Array2D[_InexactT]: ... @overload # 1-d or 2-d def expm_multiply( - A: LinearOperator[_SCT], - B: onp.ArrayND[_SCT | np.float16 | np.integer[Any] | np.bool_], + A: _ToLinearOperator[_AsFloat64], + B: _SparseOrDense[_ToFloat64, tuple[Any, ...]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.ArrayND[np.float64]: ... +@overload +def expm_multiply( + A: _ToLinearOperator[_InexactT], + B: _SparseOrDense[_InexactT | npc.integer | np.bool_, tuple[Any, ...]], + start: onp.ToFloat | None = None, + stop: onp.ToFloat | None = None, + num: op.CanIndex | None = None, + endpoint: bool | None = None, + traceA: onp.ToComplex | None = None, +) -> onp.ArrayND[_InexactT]: ... +@overload # fallback +def expm_multiply( + A: _ToLinearOperator[npc.number], + B: _SparseOrDense[npc.number, tuple[Any, ...]], start: onp.ToFloat | None = None, stop: onp.ToFloat | None = None, num: op.CanIndex | None = None, endpoint: bool | None = None, traceA: onp.ToComplex | None = None, -) -> onp.Array1D[_SCT] | onp.Array2D[_SCT]: ... +) -> onp.ArrayND[Incomplete]: ... diff --git a/tests/sparse/sparse/__init__.pyi b/tests/sparse/linalg/__init__.pyi similarity index 100% rename from tests/sparse/sparse/__init__.pyi rename to tests/sparse/linalg/__init__.pyi diff --git a/tests/sparse/sparse/test_dsolve.pyi b/tests/sparse/linalg/test_dsolve.pyi similarity index 100% rename from tests/sparse/sparse/test_dsolve.pyi rename to tests/sparse/linalg/test_dsolve.pyi diff --git a/tests/sparse/linalg/test_expm_multiply.pyi b/tests/sparse/linalg/test_expm_multiply.pyi new file mode 100644 index 00000000..61183f9f --- /dev/null +++ b/tests/sparse/linalg/test_expm_multiply.pyi @@ -0,0 +1,80 @@ +from typing import Any, assert_type + +import numpy as np +import numpy.typing as npt + +from scipy.sparse import coo_array +from scipy.sparse.linalg import LinearOperator, expm_multiply + +_dense_i8_1d: np.ndarray[tuple[int], np.dtype[np.int8]] +_dense_i8_2d: np.ndarray[tuple[int, int], np.dtype[np.int8]] +_dense_i8_nd: npt.NDArray[np.int8] + +_dense_f32_1d: np.ndarray[tuple[int], np.dtype[np.float32]] +_dense_f32_2d: np.ndarray[tuple[int, int], np.dtype[np.float32]] +_dense_f32_nd: npt.NDArray[np.float32] + +_sparse_i8_1d: coo_array[np.int8, tuple[int]] +_sparse_i8_2d: coo_array[np.int8, tuple[int, int]] +_sparse_i8_nd: coo_array[np.int8] + +_sparse_f32_1d: coo_array[np.float32, tuple[int]] +_sparse_f32_2d: coo_array[np.float32, tuple[int, int]] +_sparse_f32_nd: coo_array[np.float32] + +_linop_i8: LinearOperator[np.int8] +_linop_f32: LinearOperator[np.float32] + +# + +assert_type(expm_multiply(_dense_i8_2d, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_2d, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_2d, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_dense_i8_nd, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_nd, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_nd, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_sparse_i8_2d, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_sparse_i8_2d, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_sparse_i8_2d, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_linop_i8, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_linop_i8, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_linop_i8, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +# + +assert_type(expm_multiply(_dense_i8_2d, _dense_i8_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_2d, _dense_i8_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_2d, _dense_i8_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_dense_i8_nd, _dense_i8_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_nd, _dense_i8_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_dense_i8_nd, _dense_i8_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_sparse_i8_2d, _dense_i8_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_sparse_i8_2d, _dense_i8_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_sparse_i8_2d, _dense_i8_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +assert_type(expm_multiply(_linop_i8, _dense_i8_1d), np.ndarray[tuple[int], np.dtype[np.float64]]) +assert_type(expm_multiply(_linop_i8, _dense_i8_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]]) +assert_type(expm_multiply(_linop_i8, _dense_i8_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float64]]) + +# + +assert_type(expm_multiply(_dense_f32_2d, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float32]]) +assert_type(expm_multiply(_dense_f32_2d, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float32]]) +assert_type(expm_multiply(_dense_f32_2d, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float32]]) + +assert_type(expm_multiply(_dense_f32_nd, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float32]]) +assert_type(expm_multiply(_dense_f32_nd, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float32]]) +assert_type(expm_multiply(_dense_f32_nd, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float32]]) + +assert_type(expm_multiply(_sparse_f32_2d, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float32]]) +assert_type(expm_multiply(_sparse_f32_2d, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float32]]) +assert_type(expm_multiply(_sparse_f32_2d, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float32]]) + +assert_type(expm_multiply(_linop_f32, _dense_f32_1d), np.ndarray[tuple[int], np.dtype[np.float32]]) +assert_type(expm_multiply(_linop_f32, _dense_f32_2d), np.ndarray[tuple[int, int], np.dtype[np.float32]]) +assert_type(expm_multiply(_linop_f32, _dense_f32_nd), np.ndarray[tuple[Any, ...], np.dtype[np.float32]])