diff --git a/scipy-stubs/integrate/_ivp/base.pyi b/scipy-stubs/integrate/_ivp/base.pyi index 38567b82..a0be153d 100644 --- a/scipy-stubs/integrate/_ivp/base.pyi +++ b/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: ... diff --git a/scipy-stubs/integrate/_ivp/bdf.pyi b/scipy-stubs/integrate/_ivp/bdf.pyi index 540a5270..91b5a874 100644 --- a/scipy-stubs/integrate/_ivp/bdf.pyi +++ b/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] diff --git a/scipy-stubs/integrate/_ivp/common.pyi b/scipy-stubs/integrate/_ivp/common.pyi index dab97812..6dc55cfd 100644 --- a/scipy-stubs/integrate/_ivp/common.pyi +++ b/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_ diff --git a/scipy-stubs/integrate/_ivp/ivp.pyi b/scipy-stubs/integrate/_ivp/ivp.pyi index 7cc70cf3..a0aadf8c 100644 --- a/scipy-stubs/integrate/_ivp/ivp.pyi +++ b/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], diff --git a/scipy-stubs/integrate/_ivp/lsoda.pyi b/scipy-stubs/integrate/_ivp/lsoda.pyi index 06f1faeb..31e26e62 100644 --- a/scipy-stubs/integrate/_ivp/lsoda.pyi +++ b/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] diff --git a/scipy-stubs/integrate/_ivp/radau.pyi b/scipy-stubs/integrate/_ivp/radau.pyi index 74d871c1..d91d202b 100644 --- a/scipy-stubs/integrate/_ivp/radau.pyi +++ b/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] diff --git a/scipy-stubs/integrate/_ivp/rk.pyi b/scipy-stubs/integrate/_ivp/rk.pyi index c6f9e6a3..8c4f770e 100644 --- a/scipy-stubs/integrate/_ivp/rk.pyi +++ b/scipy-stubs/integrate/_ivp/rk.pyi @@ -60,7 +60,7 @@ 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] @@ -68,7 +68,7 @@ class RkDenseOutput(DenseOutput, Generic[_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]