🐛 integrate: several OdeSolver fixes (#880)

Author: jorenham

scipy-stubs/integrate/_ivp/base.pyi

@@ -1,75 +1,85 @@
 from collections.abc import Callable
-from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, overload
+from typing import Any, ClassVar, Final, Generic, Literal, TypeAlias, overload
+from typing_extensions import TypeVar
 
 import numpy as np
 import optype.numpy as onp
 import optype.numpy.compat as npc
 
-_VT = TypeVar("_VT", bound=onp.ArrayND[npc.inexact], default=onp.ArrayND[Any])
+_ScalarT = TypeVar("_ScalarT", bound=np.float64 | np.complex128, default=np.float64)
+_ScalarT_co = TypeVar("_ScalarT_co", bound=npc.inexact, default=np.float64 | Any, covariant=True)
 
-class OdeSolver:
+_ToFunReal: TypeAlias = Callable[[float, onp.ArrayND[np.float64]], onp.ToFloatND]
+_ToFunComplex: TypeAlias = Callable[[float, onp.ArrayND[np.complex128]], onp.ToComplexND]
+
+@overload
+def check_arguments(
+    fun: _ToFunReal, y0: onp.ToFloatND, support_complex: bool
+) -> Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.float64]]: ...
+@overload
+def check_arguments(
+    fun: _ToFunComplex, y0: onp.ToJustComplexND, support_complex: Literal[True]
+) -> Callable[[float, onp.ArrayND[np.complex128]], onp.ArrayND[np.complex128]]: ...
+
+class OdeSolver(Generic[_ScalarT]):
     TOO_SMALL_STEP: ClassVar[str] = ...
 
     t: float
     t_old: float
     t_bound: float
+    y: onp.ArrayND[_ScalarT]
     vectorized: bool
-    fun: Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.float64]]
-    fun_single: Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.float64]]
-    fun_vectorized: Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.float64]]
-    direction: float
-    n: int
+    fun: Callable[[float, onp.ArrayND[_ScalarT]], onp.ArrayND[_ScalarT]]
+    fun_single: Callable[[float, onp.Array1D[_ScalarT]], onp.Array1D[_ScalarT]]
+    fun_vectorized: Callable[[float, onp.Array2D[_ScalarT]], onp.Array2D[_ScalarT]]
+    direction: np.float64
     status: Literal["running", "finished", "failed"]
+    n: int
     nfev: int
     njev: int
     nlu: int
 
     @overload
     def __init__(
-        self,
+        self: OdeSolver[np.float64],
         /,
-        fun: Callable[[float, onp.ArrayND[np.float64]], onp.ToFloatND],
-        t0: onp.ToFloatND,
+        fun: _ToFunReal,
+        t0: float,
         y0: onp.ToFloatND,
-        t_bound: onp.ToFloat,
+        t_bound: float,
         vectorized: bool,
         support_complex: onp.ToBool = False,
     ) -> None: ...
     @overload
     def __init__(
-        self,
+        self: OdeSolver[np.complex128],
         /,
-        fun: Callable[[float, onp.ArrayND[np.float64 | np.complex128]], onp.ToComplexND],
-        t0: onp.ToFloat,
-        y0: onp.ToComplexND,
+        fun: _ToFunComplex,
+        t0: float,
+        y0: onp.ToJustComplexND,
         t_bound: onp.ToFloat,
         vectorized: bool,
         support_complex: onp.ToTrue,
     ) -> None: ...
     @property
-    def step_size(self, /) -> float | None: ...
+    def step_size(self, /) -> np.float64 | None: ...
     def step(self, /) -> str | None: ...
-    def dense_output(self, /) -> ConstantDenseOutput: ...
+    def dense_output(self, /) -> ConstantDenseOutput[_ScalarT]: ...
 
-class DenseOutput:
+class DenseOutput(Generic[_ScalarT_co]):
     t_old: Final[float]
     t: Final[float]
     t_min: Final[float]
     t_max: Final[float]
 
-    def __init__(self, /, t_old: onp.ToFloat, t: onp.ToFloat) -> None: ...
+    def __init__(self, /, t_old: float, t: float) -> None: ...
+
+    #
     @overload
-    def __call__(self, /, t: onp.ToFloat) -> onp.Array1D[npc.inexact]: ...
+    def __call__(self, /, t: onp.ToFloat) -> onp.Array1D[_ScalarT_co]: ...
     @overload
-    def __call__(self, /, t: onp.ToFloatND) -> onp.ArrayND[npc.inexact]: ...
+    def __call__(self, /, t: onp.ToFloatND) -> onp.ArrayND[_ScalarT_co]: ...
 
-class ConstantDenseOutput(DenseOutput, Generic[_VT]):
-    value: _VT
-    def __init__(self, /, t_old: onp.ToFloat, t: onp.ToFloat, value: _VT) -> None: ...
-
-def check_arguments(
-    fun: Callable[[float, onp.ArrayND[np.float64]], onp.ToComplexND], y0: onp.ToComplexND, support_complex: bool
-) -> (
-    Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.float64]]
-    | Callable[[float, onp.ArrayND[np.float64]], onp.ArrayND[np.complex128]]
-): ...
+class ConstantDenseOutput(DenseOutput[_ScalarT_co], Generic[_ScalarT_co]):
+    value: onp.ArrayND[_ScalarT_co]
+    def __init__(self, /, t_old: float, t: float, value: onp.ArrayND[_ScalarT_co]) -> None: ...

scipy-stubs/integrate/_ivp/bdf.pyi

@@ -63,7 +63,7 @@ class BDF(OdeSolver, Generic[_SCT_co]):
         **extraneous: Never,
     ) -> None: ...
 
-class BdfDenseOutput(DenseOutput):
+class BdfDenseOutput(DenseOutput[np.float64]):
     order: int
     t_shift: onp.ArrayND[np.float64]
     denom: onp.ArrayND[np.float64]

scipy-stubs/integrate/_ivp/common.pyi

@@ -12,7 +12,7 @@ from scipy.sparse import csc_matrix
 
 _FloatingT = TypeVar("_FloatingT", bound=npc.floating)
 _ToFloatT = TypeVar("_ToFloatT", bound=onp.ToFloat)
-_InterpT_co = TypeVar("_InterpT_co", bound=DenseOutput, default=DenseOutput, covariant=True)
+_InterpT_co = TypeVar("_InterpT_co", bound=DenseOutput[npc.inexact], default=DenseOutput[Any], covariant=True)
 
 _ToFloat64: TypeAlias = np.float16 | np.float32 | np.float64 | npc.integer | np.bool_
 

scipy-stubs/integrate/_ivp/ivp.pyi

@@ -72,7 +72,7 @@ class OdeResult(_RichResult[Any], Generic[_Inexact64T_co]):
 def prepare_events(events: _Events[_Inexact64T]) -> tuple[_Events[_Inexact64T], _Float1D, _Float1D]: ...
 def solve_event_equation(event: _FuncEvent[_Inexact64T], sol: _FuncSol[_Inexact64T], t_old: float, t: float) -> float: ...
 def handle_events(
-    sol: DenseOutput,
+    sol: DenseOutput[Any],
     events: Sequence[_FuncEvent[_Inexact64T]],
     active_events: onp.ArrayND[np.intp],
     event_count: onp.ArrayND[np.intp | np.float64],

scipy-stubs/integrate/_ivp/lsoda.pyi

@@ -26,7 +26,7 @@ class LSODA(OdeSolver):
         **extraneous: Never,
     ) -> None: ...
 
-class LsodaDenseOutput(DenseOutput):
+class LsodaDenseOutput(DenseOutput[np.float64]):
     h: float
     yh: onp.Array1D[np.float64]
     p: onp.Array1D[np.intp]

scipy-stubs/integrate/_ivp/radau.pyi

@@ -69,7 +69,7 @@ class Radau(OdeSolver):
         **extraneous: Never,
     ) -> None: ...
 
-class RadauDenseOutput(DenseOutput):
+class RadauDenseOutput(DenseOutput[np.float64]):
     order: int
     h: float
     Q: onp.ArrayND[np.float64]

scipy-stubs/integrate/_ivp/rk.pyi

@@ -60,15 +60,15 @@ class DOP853(RungeKutta[_SCT_fc], Generic[_SCT_fc]):
 
     K_extended: onp.ArrayND[_SCT_fc]
 
-class RkDenseOutput(DenseOutput, Generic[_SCT_fc]):
+class RkDenseOutput(DenseOutput[_SCT_fc], Generic[_SCT_fc]):
     h: float
     order: int
     Q: onp.ArrayND[_SCT_fc]
     y_old: onp.ArrayND[_SCT_fc]
 
     def __init__(self, /, t_old: float, t: float, y_old: onp.ArrayND[_SCT_fc], Q: onp.ArrayND[_SCT_fc]) -> None: ...
 
-class Dop853DenseOutput(DenseOutput, Generic[_SCT_fc]):
+class Dop853DenseOutput(DenseOutput[_SCT_fc], Generic[_SCT_fc]):
     h: float
     F: onp.ArrayND[_SCT_fc]
     y_old: onp.ArrayND[_SCT_fc]