🏷️ linalg: stub more of the internal scipy machinery (#710)

Author: jorenham

scipy-stubs/linalg/_linalg_pythran.pyi

@@ -0,0 +1,14 @@
+from typing import Final, TypeVar
+
+import numpy as np
+import optype.numpy as onp
+
+_InexactT = TypeVar("_InexactT", bound=np.float32 | np.float64 | np.complex64 | np.complex128)
+
+###
+
+__pythran__: Final[tuple[str, str]] = ...
+
+def _funm_loops(
+    F: onp.Array2D[_InexactT], T: onp.Array2D[_InexactT], n: int, minden: _InexactT
+) -> tuple[onp.Array2D[_InexactT], _InexactT]: ...

scipy-stubs/linalg/_matfuncs.pyi

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any, Literal, TypeAlias, TypeVar, overload
+from typing import Any, Final, Literal, TypeAlias, TypeVar, overload
 from typing_extensions import deprecated
 
 import numpy as np
@@ -28,6 +28,9 @@ __all__ = [
 ]
 
 _InexactT = TypeVar("_InexactT", bound=npc.inexact)
+_ComplexT = TypeVar("_ComplexT", bound=npc.complexfloating)
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
+
 _FuncND: TypeAlias = Callable[[onp.Array[Any, _InexactT]], onp.ToComplexND]  # return type is unsafely cast to the input type
 
 _ToPosInt: TypeAlias = npc.unsignedinteger | Literal[0, 1, 2, 4, 5, 6, 7, 8]
@@ -41,6 +44,23 @@ _InexactND: TypeAlias = onp.ArrayND[npc.inexact]
 
 ###
 
+eps: Final[np.float64] = ...  # undocumented
+feps: Final[np.float32] = ...  # undocumented
+_array_precision: Final[dict[Literal["i", "l", "f", "d", "F", "D"], Literal[0, 1]]] = ...  # undocumented
+
+def _asarray_square(A: onp.ToArray2D[_InexactT]) -> onp.Array2D[_InexactT]: ...  # undocumented
+
+#
+@overload
+def _maybe_real(
+    A: onp.ArrayND[npc.inexact], B: onp.ArrayND[npc.inexact64, _ShapeT], tol: float | None = None
+) -> onp.ArrayND[np.float64, _ShapeT]: ...  # undocumented
+@overload
+def _maybe_real(
+    A: onp.ArrayND[npc.inexact], B: onp.ArrayND[npc.inexact32, _ShapeT], tol: float | None = None
+) -> onp.ArrayND[np.float32, _ShapeT]: ...  # undocumented
+
+#
 @overload  # +integer, +unsignedinteger
 def fractional_matrix_power(A: onp.ToIntND, t: _ToPosInt) -> _IntND: ...
 @overload  # ~float64, +integer
@@ -118,6 +138,9 @@ def expm(A: onp.ToJustComplexND) -> _ComplexND: ...
 @overload  # +complexfloating
 def expm(A: onp.ToComplexND) -> _InexactND: ...
 
+#
+def _exp_sinch(x: onp.ArrayND[_ComplexT, _ShapeT]) -> onp.ArrayND[_ComplexT, _ShapeT]: ...  # undocumented
+
 #
 @overload  # +integer | ~float64
 def cosm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...

scipy-stubs/linalg/_matfuncs_schur_sqrtm.pyi

@@ -0,0 +1,15 @@
+from typing import Literal, TypeVar
+
+import numpy as np
+import optype.numpy as onp
+
+_InexactT = TypeVar("_InexactT", bound=np.float32 | np.float64 | np.complex64 | np.complex128)
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
+
+###
+
+class error(Exception): ...
+
+def recursive_schur_sqrtm(
+    A: onp.ArrayND[_InexactT, _ShapeT],
+) -> tuple[onp.ArrayND[_InexactT, _ShapeT], Literal[0, 1], Literal[0, 1], int]: ...

scipy-stubs/linalg/_misc.pyi

@@ -173,3 +173,6 @@ def norm(
     keepdims: AnyBool = False,
     check_finite: AnyBool = True,
 ) -> np.floating[Any] | onp.ArrayND[np.floating[Any]]: ...
+
+#
+def _datacopied(arr: onp.ArrayND, original: onp.CanArrayND) -> bool: ...  # undocumented

scipy-stubs/linalg/_solvers.pyi

@@ -1,4 +1,5 @@
-from typing import Final, Literal, TypeAlias, overload
+from _typeshed import Incomplete
+from typing import Final, Literal, TypeAlias, TypeVar, overload
 
 import numpy as np
 import optype as op
@@ -13,6 +14,8 @@ __all__ = [
     "solve_sylvester",
 ]
 
+_InexactT = TypeVar("_InexactT", bound=np.float32 | np.float64 | np.complex64 | np.complex128)
+
 _FloatND: TypeAlias = onp.ArrayND[np.float32 | np.float64]
 _ComplexND: TypeAlias = onp.ArrayND[np.complex64 | np.complex128]
 
@@ -40,6 +43,10 @@ def solve_continuous_lyapunov(a: onp.ToComplexND, q: onp.ToJustComplexND) -> _Co
 #
 solve_lyapunov: Final = solve_continuous_lyapunov
 
+#
+def _solve_discrete_lyapunov_direct(a: onp.Array2D[_InexactT], q: onp.Array2D[_InexactT]) -> onp.Array2D[_InexactT]: ...
+def _solve_discrete_lyapunov_bilinear(a: onp.Array2D[_InexactT], q: onp.Array2D[_InexactT]) -> onp.Array2D[_InexactT]: ...
+
 #
 @overload  # real
 def solve_discrete_lyapunov(a: onp.ToFloatND, q: onp.ToFloatND, method: _DiscreteMethod | None = None) -> _FloatND: ...
@@ -213,3 +220,25 @@ def solve_discrete_are(
     s: onp.ToJustComplexND,
     balanced: op.CanBool = True,
 ) -> _ComplexND: ...
+
+#
+def _are_validate_args(
+    a: onp.ToComplexND,
+    b: onp.ToComplexND,
+    q: onp.ToComplexND,
+    r: onp.ToComplexND,
+    e: onp.ToComplexND | None,
+    s: onp.ToComplexND | None,
+    eq_type: Literal["care", "dare"] = "care",
+) -> tuple[
+    onp.ArrayND[Incomplete],  # a
+    onp.ArrayND[Incomplete],  # b
+    onp.ArrayND[Incomplete],  # q
+    onp.ArrayND[Incomplete],  # r
+    onp.ArrayND[Incomplete],  # e
+    onp.ArrayND[Incomplete],  # s
+    int,  # m
+    int,  # n
+    type[float | complex],  # r_or_c
+    bool,  # gen_or_not
+]: ...

scipy-stubs/linalg/_testutils.pyi

@@ -0,0 +1,30 @@
+from collections.abc import Callable, Iterable
+from typing import Any, Final, Generic, TypeAlias, TypeVar
+
+import numpy as np
+import optype.numpy as onp
+
+_T = TypeVar("_T")
+_ScalarT = TypeVar("_ScalarT", bound=np.generic)
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
+
+_Ignored: TypeAlias = object
+
+class _FakeMatrix(Generic[_ScalarT, _ShapeT]):  # undocumented
+    _data: onp.ArrayND[_ScalarT, _ShapeT]
+
+    def __init__(self, /, data: onp.ArrayND[_ScalarT, _ShapeT]) -> None: ...
+    __array_interface__: Final[Callable[[], dict[str, Any]]] = ...
+
+class _FakeMatrix2(Generic[_ScalarT, _ShapeT]):  # undocumented
+    _data: onp.ArrayND[_ScalarT, _ShapeT]
+    def __init__(self, /, data: onp.ArrayND[_ScalarT, _ShapeT]) -> None: ...
+    def __array__(self, /, dtype: np.dtype | None = None, copy: bool | None = None) -> onp.ArrayND[_ScalarT, _ShapeT]: ...
+
+def _get_array(shape: _ShapeT, dtype: type[_ScalarT]) -> onp.ArrayND[_ScalarT, _ShapeT]: ...  # undocumented
+def _id(x: _T) -> _T: ...  # undocumented
+
+#
+def assert_no_overwrite(
+    call: Callable[..., _Ignored], shapes: Iterable[tuple[int, ...]], dtypes: Iterable[type[np.generic]] | None = None
+) -> None: ...  # undocumented