diff --git a/scipy-stubs/integrate/_odepack.pyi b/scipy-stubs/integrate/_odepack.pyi index 947a4d78..344291b6 100644 --- a/scipy-stubs/integrate/_odepack.pyi +++ b/scipy-stubs/integrate/_odepack.pyi @@ -1,29 +1,77 @@ -from _typeshed import Incomplete +# defined in scipy/integrate/_odepackmodule.c + from collections.abc import Callable +from typing import Literal, TypedDict, overload, type_check_only +import numpy as np import optype.numpy as onp class error(Exception): ... +@type_check_only +class _InfoDict(TypedDict): + hu: onp.Array1D[np.float64] + tcur: onp.Array1D[np.float64] + tolsf: onp.Array1D[np.float64] + tsw: onp.Array1D[np.float64] + nst: onp.Array1D[np.int32] + nfe: onp.Array1D[np.int32] + nje: onp.Array1D[np.int32] + nqu: onp.Array1D[np.int32] + imxer: int + lenrw: int + leniw: int + mused: onp.Array1D[np.int32] + +### + +# NOTE: This (private) `odeint` differs from the one in `_odepack_py` because it +# doesn't have a `printmessg` kwarg, and has a different return type. +@overload def odeint( - fun: Callable[..., Incomplete], - y0: onp.ToFloat | onp.ToFloatND, - t: float, + fun: Callable[..., onp.ToFloat1D | float], + y0: onp.ToFloatND | float, + t: onp.ToFloat1D, args: tuple[object, ...] = (), - Dfun: Callable[..., Incomplete] | None = None, - col_deriv: int = 0, - ml: int = ..., - mu: int = ..., - full_output: onp.ToBool = 0, - rtol: float = ..., - atol: float = ..., - tcrit: float = ..., + Dfun: Callable[..., onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, + ml: int = -1, + mu: int = -1, + full_output: Literal[0, False] = 0, + rtol: float | None = None, + atol: float | None = None, + tcrit: onp.ToFloat1D | None = None, + h0: float = 0.0, + hmax: float = 0.0, + hmin: float = 0.0, + ixpr: onp.ToBool = 0, + mxstep: int = 0, + mxhnil: int = 0, + mxordn: int = 12, + mxords: int = 5, + tfirst: int = 0, +) -> tuple[onp.Array2D[np.float64], int]: ... +@overload +def odeint( + fun: Callable[..., onp.ToFloat1D | float], + y0: onp.ToFloatND | float, + t: onp.ToFloat1D, + args: tuple[object, ...], + Dfun: Callable[..., onp.ToFloat2D] | None, + col_deriv: onp.ToBool, + ml: int, + mu: int, + full_output: Literal[True, 1], + rtol: float | None = None, + atol: float | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: float = 0.0, - mxstep: float = 0.0, + ixpr: onp.ToBool = 0, + mxstep: int = 0, mxhnil: int = 0, - mxordn: int = 0, - mxords: int = 0, -) -> tuple[Incomplete, ...]: ... + mxordn: int = 12, + mxords: int = 5, + tfirst: int = 0, +) -> tuple[onp.Array2D[np.float64], _InfoDict, int]: ... diff --git a/scipy-stubs/integrate/_odepack_py.pyi b/scipy-stubs/integrate/_odepack_py.pyi index 8a964417..68a729b5 100644 --- a/scipy-stubs/integrate/_odepack_py.pyi +++ b/scipy-stubs/integrate/_odepack_py.pyi @@ -1,234 +1,243 @@ -from typing import Literal, Protocol, overload, type_check_only -from typing_extensions import TypeVar, TypeVarTuple, Unpack +import threading +from collections.abc import Callable +from typing import Final, Literal, TypeAlias, TypeVar, TypeVarTuple, TypedDict, overload, type_check_only +import numpy as np import optype.numpy as onp -import optype.numpy.compat as npc - -from ._typing import ODEInfoDict __all__ = ["ODEintWarning", "odeint"] -_Ts = TypeVarTuple("_Ts", default=Unpack[tuple[()]]) -_YT = TypeVar("_YT", bound=onp.ToFloat1D | onp.ToFloat) +_Ts = TypeVarTuple("_Ts") +_R = TypeVar("_R") -@type_check_only -class _ODEFunc(Protocol[_YT, *_Ts]): - def __call__(self, y: _YT, t: float, /, *args: *_Ts) -> _YT: ... +_FuncYT: TypeAlias = Callable[[onp.Array1D[np.float64], float, *_Ts], _R] +_FuncTY: TypeAlias = Callable[[float, onp.Array1D[np.float64], *_Ts], _R] @type_check_only -class _ODEFuncInv(Protocol[_YT, *_Ts]): - def __call__(self, t: float, y: _YT, /, *args: *_Ts) -> _YT: ... +class _InfoDict(TypedDict): + hu: onp.Array1D[np.float64] + tcur: onp.Array1D[np.float64] + tolsf: onp.Array1D[np.float64] + tsw: onp.Array1D[np.float64] + nst: onp.Array1D[np.int32] + nfe: onp.Array1D[np.int32] + nje: onp.Array1D[np.int32] + nqu: onp.Array1D[np.int32] + imxer: int + lenrw: int + leniw: int + mused: onp.Array1D[np.int32] + message: str ### +ODE_LOCK: Final[threading.Lock] = ... + class ODEintWarning(Warning): ... -# unspecified args -@overload +@overload # args=() (default), full_output=False (default), tfirst=False (default) def odeint( - func: _ODEFunc[_YT], - y0: _YT, - t: onp.ToInt1D, + func: _FuncYT[*tuple[()], onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, args: tuple[()] = (), - Dfun: _ODEFunc[_YT] | None = None, - col_deriv: Literal[0, 1] | bool = 0, - full_output: onp.ToFalse = 0, + Dfun: _FuncYT[*tuple[()], onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, + full_output: Literal[False, 0] = 0, ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToFalse = False, -) -> onp.Array2D[npc.floating]: ... -@overload + printmessg: onp.ToBool = 0, + tfirst: Literal[False, 0] = False, +) -> onp.Array2D[np.float64]: ... +@overload # args=() (default), full_output=False (default), *, tfirst=True def odeint( - func: _ODEFunc[_YT], - y0: _YT, - t: onp.ToInt1D, + func: _FuncTY[*tuple[()], onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, args: tuple[()] = (), - Dfun: _ODEFunc[_YT] | None = None, - col_deriv: Literal[0, 1] | bool = 0, - full_output: onp.ToFalse = 0, + Dfun: _FuncTY[*tuple[()], onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, + full_output: Literal[False, 0] = 0, ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, + printmessg: onp.ToBool = 0, *, - tfirst: onp.ToTrue, -) -> onp.Array2D[npc.floating]: ... -@overload + tfirst: Literal[True, 1], +) -> onp.Array2D[np.float64]: ... +@overload # args=() (default), *, full_output=True, tfirst=False (default) def odeint( - func: _ODEFunc[_YT], - y0: _YT, - t: onp.ToInt1D, + func: _FuncYT[*tuple[()], onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, args: tuple[()] = (), - Dfun: _ODEFunc[_YT] | None = None, - col_deriv: Literal[0, 1] | bool = 0, + Dfun: _FuncYT[*tuple[()], onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, *, - full_output: onp.ToTrue, + full_output: Literal[True, 1], ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToFalse = False, -) -> tuple[onp.Array2D[npc.floating], ODEInfoDict]: ... -@overload + printmessg: onp.ToBool = 0, + tfirst: Literal[False, 0] = False, +) -> tuple[onp.Array2D[np.float64], _InfoDict]: ... +@overload # args=() (default), full_output=True, *, tfirst=True def odeint( - func: _ODEFunc[_YT], - y0: _YT, - t: onp.ToInt1D, + func: _FuncTY[*tuple[()], onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, args: tuple[()] = (), - Dfun: _ODEFunc[_YT] | None = None, - col_deriv: Literal[0, 1] | bool = 0, + Dfun: _FuncTY[*tuple[()], onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, *, - full_output: onp.ToTrue, + full_output: Literal[True, 1], ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToTrue, -) -> tuple[onp.Array2D[npc.floating], ODEInfoDict]: ... - -# specified args -@overload + printmessg: onp.ToBool = 0, + tfirst: Literal[True, 1], +) -> tuple[onp.Array2D[np.float64], _InfoDict]: ... +@overload # args=, full_output=False (default), tfirst=False (default) def odeint( - func: _ODEFunc[_YT, *_Ts], - y0: _YT, - t: onp.ToInt1D, - args: tuple[*_Ts] = ..., - Dfun: _ODEFunc[_YT, *_Ts] | None = None, - col_deriv: Literal[0, 1] | bool = 0, - full_output: onp.ToFalse = 0, + func: _FuncYT[*_Ts, onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, + args: tuple[*_Ts], + Dfun: _FuncYT[*_Ts, onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, + full_output: Literal[False, 0] = 0, ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToFalse = False, -) -> onp.Array2D[npc.floating]: ... -@overload + printmessg: onp.ToBool = 0, + tfirst: Literal[False, 0] = False, +) -> onp.Array2D[np.float64]: ... +@overload # args=, full_output=False (default), *, tfirst=True def odeint( - func: _ODEFuncInv[_YT, *_Ts], - y0: _YT, - t: onp.ToInt1D, - args: tuple[*_Ts] = ..., - Dfun: _ODEFuncInv[_YT, *_Ts] | None = None, - col_deriv: Literal[0, 1] | bool = 0, - full_output: onp.ToFalse = 0, + func: _FuncTY[*_Ts, onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, + args: tuple[*_Ts], + Dfun: _FuncTY[*_Ts, onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, + full_output: Literal[False, 0] = 0, ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, + printmessg: onp.ToBool = 0, *, - tfirst: onp.ToTrue, -) -> onp.Array2D[npc.floating]: ... -@overload + tfirst: Literal[True, 1], +) -> onp.Array2D[np.float64]: ... +@overload # args=, *, full_output=True, tfirst=False (default) def odeint( - func: _ODEFunc[_YT, *_Ts], - y0: _YT, - t: onp.ToInt1D, - args: tuple[*_Ts] = ..., - Dfun: _ODEFunc[_YT, *_Ts] | None = None, - col_deriv: Literal[0, 1] | bool = 0, + func: _FuncYT[*_Ts, onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, + args: tuple[*_Ts], + Dfun: _FuncYT[*_Ts, onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, *, - full_output: onp.ToTrue, + full_output: Literal[True, 1], ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToFalse = False, -) -> tuple[onp.Array2D[npc.floating], ODEInfoDict]: ... -@overload + printmessg: onp.ToBool = 0, + tfirst: Literal[False, 0] = False, +) -> tuple[onp.Array2D[np.float64], _InfoDict]: ... +@overload # args=, *, full_output=True, tfirst=True def odeint( - func: _ODEFuncInv[_YT, *_Ts], - y0: _YT, - t: onp.ToInt1D, - args: tuple[*_Ts] = ..., - Dfun: _ODEFuncInv[_YT, *_Ts] | None = None, - col_deriv: Literal[0, 1] | bool = 0, + func: _FuncTY[*_Ts, onp.ToFloat1D | float], + y0: onp.ToFloat1D | float, + t: onp.ToFloat1D, + args: tuple[*_Ts], + Dfun: _FuncTY[*_Ts, onp.ToFloat2D] | None = None, + col_deriv: onp.ToBool = 0, *, - full_output: onp.ToTrue, + full_output: Literal[True, 1], ml: int | None = None, mu: int | None = None, rtol: float | None = None, atol: float | None = None, - tcrit: onp.ArrayND[npc.integer | npc.floating] | None = None, + tcrit: onp.ToFloat1D | None = None, h0: float = 0.0, hmax: float = 0.0, hmin: float = 0.0, - ixpr: Literal[0, 1] | bool = 0, + ixpr: onp.ToBool = 0, mxstep: int = 0, mxhnil: int = 0, mxordn: int = 12, mxords: int = 5, - printmessg: Literal[0, 1] | bool = 0, - tfirst: onp.ToTrue, -) -> tuple[onp.Array2D[npc.floating], ODEInfoDict]: ... + printmessg: onp.ToBool = 0, + tfirst: Literal[True, 1], +) -> tuple[onp.Array2D[np.float64], _InfoDict]: ... diff --git a/scipy-stubs/integrate/_typing.pyi b/scipy-stubs/integrate/_typing.pyi index 2b40ba41..9b85c3a1 100644 --- a/scipy-stubs/integrate/_typing.pyi +++ b/scipy-stubs/integrate/_typing.pyi @@ -6,7 +6,7 @@ import numpy as np import optype.numpy as onp import optype.numpy.compat as npc -__all__ = "ODEInfoDict", "QuadInfoDict", "QuadOpts", "QuadWeights" +__all__ = "QuadInfoDict", "QuadOpts", "QuadWeights" _IntLike: TypeAlias = int | npc.integer _FloatLike: TypeAlias = float | npc.floating @@ -48,18 +48,3 @@ class QuadInfoDict(TypedDict): rslst: NotRequired[onp.Array1D[np.float64]] erlst: NotRequired[onp.Array1D[np.float64]] ierlst: NotRequired[onp.Array1D[np.float64]] - -@type_check_only -class ODEInfoDict(TypedDict): - hu: onp.Array1D[np.float64] - tcur: onp.Array1D[np.float64] - tolsf: onp.Array1D[np.float64] - tsw: float - nst: int - nfe: int - nje: int - nqu: onp.Array1D[np.int_] - imxer: int - lenrw: int - leniw: int - mused: onp.Array1D[np.int_] diff --git a/tests/integrate/test_odeint.pyi b/tests/integrate/test_odeint.pyi new file mode 100644 index 00000000..041da4f3 --- /dev/null +++ b/tests/integrate/test_odeint.pyi @@ -0,0 +1,46 @@ +from typing import TypeAlias, assert_type, type_check_only + +import numpy as np +import numpy.typing as npt + +from scipy.integrate import odeint + +# based on the example from the `odeint` docstring +@type_check_only +def pend_0(y: npt.NDArray[np.floating], t: float) -> list[float]: ... +@type_check_only +def pend_ty_0(t: float, y: npt.NDArray[np.floating]) -> list[float]: ... +@type_check_only +def pend(y: npt.NDArray[np.floating], t: float, b: float, c: float) -> list[float]: ... +@type_check_only +def pend_ty(t: float, y: npt.NDArray[np.floating], b: float, c: float) -> list[float]: ... + +y0: list[float] = ... +b: float = ... +c: float = ... +t: npt.NDArray[np.float64] = ... + +_VecF64: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]] +_MatF64: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]] + +### + +assert_type(odeint(pend_0, y0, t), _MatF64) +assert_type(odeint(pend_0, y0, t, full_output=True)[0], _MatF64) +assert_type(odeint(pend_0, y0, t, full_output=True)[1]["hu"], _VecF64) +assert_type(odeint(pend_0, y0, t, full_output=True)[1]["message"], str) + +assert_type(odeint(pend_ty_0, y0, t, tfirst=True), _MatF64) +assert_type(odeint(pend_ty_0, y0, t, full_output=True, tfirst=True)[0], _MatF64) +assert_type(odeint(pend_ty_0, y0, t, full_output=True, tfirst=True)[1]["hu"], _VecF64) +assert_type(odeint(pend_ty_0, y0, t, full_output=True, tfirst=True)[1]["message"], str) + +assert_type(odeint(pend, y0, t, args=(b, c)), _MatF64) +assert_type(odeint(pend, y0, t, full_output=True, args=(b, c))[0], _MatF64) +assert_type(odeint(pend, y0, t, full_output=True, args=(b, c))[1]["hu"], _VecF64) +assert_type(odeint(pend, y0, t, full_output=True, args=(b, c))[1]["message"], str) + +assert_type(odeint(pend_ty, y0, t, args=(b, c), tfirst=True), _MatF64) +assert_type(odeint(pend_ty, y0, t, full_output=True, args=(b, c), tfirst=True)[0], _MatF64) +assert_type(odeint(pend_ty, y0, t, full_output=True, args=(b, c), tfirst=True)[1]["hu"], _VecF64) +assert_type(odeint(pend_ty, y0, t, full_output=True, args=(b, c), tfirst=True)[1]["message"], str)