Merge pull request #11 from jorenham/scipy.integrate

`scipy.integrate`

README.md

@@ -80,8 +80,7 @@ pip install scipy-stubs
 ## Development Progress
 
 According to [basedpyright](https://github.com/DetachHead/basedpyright) (stricter than
-pyright), the "type completeness score" is **38.6%**
-[at the time of writing](https://github.com/jorenham/scipy-stubs/tree/15f8b2c).
+pyright), the "type completeness score" is **39.2%**.
 
 | Module                            | Stubs status    |
 |---------------------------------- |---------------- |
@@ -95,7 +94,7 @@ pyright), the "type completeness score" is **38.6%**
 | `scipy.fft`                       | 2: partial      |
 | `scipy.fft._pocketfft`            | 2: partial      |
 | `scipy.fftpack`                   | 2: partial      |
-| `scipy.integrate`                 | 2: partial      |
+| `scipy.integrate`                 | 3: ready        |
 | `scipy.integrate._ivp`            | 2: partial      |
 | `scipy.interpolate`               | 2: partial      |
 | `scipy.io`                        | 2: partial      |
@@ -126,8 +125,8 @@ pyright), the "type completeness score" is **38.6%**
 
 Status labels:
 
-0. missing (failed stubgen)
-1. skeleton (mostly succesful stubgen)
-2. partial (incomplete/broad annotations)
-3. ready (complete & valid annotations, untested)
-4. done (complete, valid, tested, and production-ready)
+1. missing (failed stubgen)
+2. skeleton (mostly succesful stubgen)
+3. partial (incomplete/broad annotations)
+4. ready (complete & valid annotations, untested)
+5. done (complete, valid, tested, and production-ready)

pyproject.toml

@@ -85,6 +85,7 @@ ignore = [
   "E741",
   "COM812", # ruff format compat
   "ANN", # TODO
+  "PYI054", # numeric-literal-too-long
 ]
 preview = true
 select = [

scipy-stubs/_lib/_ccallback.pyi

@@ -19,7 +19,7 @@ _Function: TypeAlias = CapsuleType | PyCFuncPtr | CFFIFuncP | CData
 _UserData: TypeAlias = CapsuleType | ct.c_void_p | CFFIVoidP
 
 _FuncT_co = TypeVar("_FuncT_co", bound=_Function, covariant=True, default=_Function)
-_DataT_co = TypeVar("_DataT_co", bound=_UserData | None, covariant=True, default=None)
+_DataT_co = TypeVar("_DataT_co", bound=_UserData | None, covariant=True, default=_UserData)
 
 class LowLevelCallable(tuple[CapsuleType, _FuncT_co, _DataT_co], Generic[_FuncT_co, _DataT_co]):
     # NOTE: __getitem__ will raise a ValueError

scipy-stubs/integrate/__init__.pyi

@@ -1,20 +1,39 @@
-from . import dop as dop, lsoda as lsoda, odepack as odepack, quadpack as quadpack, vode as vode
-from ._bvp import solve_bvp as solve_bvp
+from . import _ode, _odepack_py, _quadpack_py, _quadrature, dop, lsoda, odepack, quadpack, vode
+from ._bvp import solve_bvp
 from ._ivp import (
-    BDF as BDF,
-    DOP853 as DOP853,
-    LSODA as LSODA,
-    RK23 as RK23,
-    RK45 as RK45,
-    DenseOutput as DenseOutput,
-    OdeSolution as OdeSolution,
-    OdeSolver as OdeSolver,
-    Radau as Radau,
-    solve_ivp as solve_ivp,
+    BDF,
+    DOP853,
+    LSODA,
+    RK23,
+    RK45,
+    DenseOutput,
+    OdeSolution,
+    OdeSolver,
+    Radau,
+    solve_ivp,
 )
 from ._ode import *
 from ._odepack_py import *
-from ._quad_vec import quad_vec as quad_vec
+from ._quad_vec import quad_vec
 from ._quadpack_py import *
 from ._quadrature import *
-from ._tanhsinh import nsum as nsum
+
+__all__ = ["dop", "lsoda", "odepack", "quadpack", "vode"]
+__all__ += [
+    "BDF",
+    "DOP853",
+    "LSODA",
+    "RK23",
+    "RK45",
+    "DenseOutput",
+    "OdeSolution",
+    "OdeSolver",
+    "Radau",
+    "quad_vec",
+    "solve_bvp",
+    "solve_ivp",
+]
+__all__ += _ode.__all__
+__all__ += _odepack_py.__all__
+__all__ += _quadpack_py.__all__
+__all__ += _quadrature.__all__

scipy-stubs/integrate/_bvp.pyi

@@ -1,50 +1,97 @@
-from scipy._typing import Untyped
-from scipy.optimize import OptimizeResult
+# TODO: Annotate the private functions as well
 
-EPS: Untyped
-TERMINATION_MESSAGES: Untyped
+from collections.abc import Callable, Sequence
+from typing import Any, Final, Generic, Literal, TypeAlias, overload
+from typing_extensions import TypeVar
 
-class BVPResult(OptimizeResult): ...
+import numpy as np
+import numpy.typing as npt
+import optype.numpy as onpt
+from scipy.interpolate import PPoly
 
-def estimate_fun_jac(fun, x, y, p, f0: Untyped | None = None) -> Untyped: ...
-def estimate_bc_jac(bc, ya, yb, p, bc0: Untyped | None = None) -> Untyped: ...
-def compute_jac_indices(n, m, k) -> Untyped: ...
-def stacked_matmul(a, b) -> Untyped: ...
-def construct_global_jac(
-    n,
-    m,
-    k,
-    i_jac,
-    j_jac,
-    h,
-    df_dy,
-    df_dy_middle,
-    df_dp,
-    df_dp_middle,
-    dbc_dya,
-    dbc_dyb,
-    dbc_dp,
-) -> Untyped: ...
-def collocation_fun(fun, y, p, x, h) -> Untyped: ...
-def prepare_sys(n, m, k, fun, bc, fun_jac, bc_jac, x, h) -> Untyped: ...
-def solve_newton(n, m, h, col_fun, bc, jac, y, p, B, bvp_tol, bc_tol) -> Untyped: ...
-def print_iteration_header(): ...
-def print_iteration_progress(iteration, residual, bc_residual, total_nodes, nodes_added): ...
-def estimate_rms_residuals(fun, sol, x, h, p, r_middle, f_middle) -> Untyped: ...
-def create_spline(y, yp, x, h) -> Untyped: ...
-def modify_mesh(x, insert_1, insert_2) -> Untyped: ...
-def wrap_functions(fun, bc, fun_jac, bc_jac, k, a, S, D, dtype) -> Untyped: ...
+_SCT = TypeVar("_SCT", bound=np.generic)
+_SCT_fc = TypeVar("_SCT_fc", bound=np.inexact[Any], default=np.float64 | np.complex128)
+_Array_1d: TypeAlias = onpt.Array[tuple[int], _SCT]
+_Array_2d: TypeAlias = onpt.Array[tuple[int, int], _SCT]
+
+_FunRHS: TypeAlias = Callable[
+    [_Array_1d[np.float64], _Array_2d[_SCT_fc]],
+    npt.NDArray[_SCT_fc],
+]
+_FunRHS_p: TypeAlias = Callable[
+    [_Array_1d[np.float64], _Array_2d[_SCT_fc], _Array_1d[np.float64]],
+    npt.NDArray[_SCT_fc],
+]
+_FunRHS_jac: TypeAlias = Callable[
+    [_Array_1d[np.float64], _Array_2d[_SCT_fc]],
+    npt.NDArray[_SCT_fc],
+]
+_FunRHS_jac_p: TypeAlias = Callable[
+    [_Array_1d[np.float64], _Array_2d[_SCT_fc], _Array_1d[np.float64]],
+    tuple[npt.NDArray[_SCT_fc], npt.NDArray[_SCT_fc]],
+]
+_FunBCR: TypeAlias = Callable[
+    [_Array_1d[_SCT_fc], _Array_1d[_SCT_fc]],
+    npt.NDArray[_SCT_fc],
+]
+_FunBCR_p: TypeAlias = Callable[
+    [_Array_1d[_SCT_fc], _Array_1d[_SCT_fc], _Array_1d[np.float64]],
+    npt.NDArray[_SCT_fc],
+]
+_FunBCR_jac: TypeAlias = Callable[
+    [_Array_1d[_SCT_fc], _Array_1d[_SCT_fc]],
+    tuple[npt.NDArray[_SCT_fc], npt.NDArray[_SCT_fc]]
+]
+_FunBCR_jac_p: TypeAlias = Callable[
+    [_Array_1d[_SCT_fc], _Array_1d[_SCT_fc], _Array_1d[np.float64]],
+    tuple[npt.NDArray[_SCT_fc], npt.NDArray[_SCT_fc], npt.NDArray[_SCT_fc]]
+]
+
+# NOTE: this inherits from `scipy.optimize.OptimizeResult` at runtime.
+# But because `BVPResult` doesn't share all members (and optional attributes
+# still aren't a thing), it was omitted as a base class here.
+class BVPResult(Generic[_SCT_fc]):
+    sol: Final[PPoly]
+    p: Final[onpt.Array[tuple[int], np.float64] | None]
+    x: Final[onpt.Array[tuple[int], np.float64]]
+    rms_residuals: Final[onpt.Array[tuple[int], np.float64]]
+    niter: Final[int]
+    status: Final[Literal[0, 1, 2]]
+    message: Final[str]
+    success: Final[bool]
+    @property
+    def y(self, /) -> onpt.Array[tuple[int, int], _SCT_fc]: ...
+    @property
+    def yp(self) -> onpt.Array[tuple[int, int], _SCT_fc]: ...
+
+# public
+@overload
+def solve_bvp(
+    fun: _FunRHS[_SCT_fc],
+    bc: _FunBCR[_SCT_fc],
+    x: onpt.AnyFloatingArray | Sequence[float],
+    y: onpt.AnyInexactArray | Sequence[Sequence[complex]],
+    p: None = None,
+    S: onpt.AnyFloatingArray | Sequence[Sequence[float]] | None = None,
+    fun_jac: _FunRHS_jac[_SCT_fc] | None = None,
+    bc_jac: _FunBCR_jac[_SCT_fc] | None = None,
+    tol: float = 0.001,
+    max_nodes: int = 1_000,
+    verbose: Literal[0, 1, 2] = 0,
+    bc_tol: float | None = None,
+) -> BVPResult[_SCT_fc]: ...
+@overload
 def solve_bvp(
-    fun,
-    bc,
-    x,
-    y,
-    p: Untyped | None = None,
-    S: Untyped | None = None,
-    fun_jac: Untyped | None = None,
-    bc_jac: Untyped | None = None,
+    fun: _FunRHS_p[_SCT_fc],
+    bc: _FunBCR_p[_SCT_fc],
+    x: onpt.AnyFloatingArray | Sequence[float],
+    y: onpt.AnyInexactArray | Sequence[Sequence[complex]],
+    p: onpt.AnyFloatingArray | Sequence[float],
+    S: onpt.AnyFloatingArray | Sequence[Sequence[float]] | None = None,
+    fun_jac: _FunRHS_jac_p[_SCT_fc] | None = None,
+    bc_jac: _FunBCR_jac_p[_SCT_fc] | None = None,
     tol: float = 0.001,
-    max_nodes: int = 1000,
-    verbose: int = 0,
-    bc_tol: Untyped | None = None,
-) -> Untyped: ...
+    max_nodes: int = 1_000,
+    verbose: Literal[0, 1, 2] = 0,
+    bc_tol: float | None = None,
+) -> BVPResult[_SCT_fc]: ...

scipy-stubs/integrate/_ivp/__init__.pyi

@@ -1,7 +1,10 @@
-from .base import DenseOutput as DenseOutput, OdeSolver as OdeSolver
-from .bdf import BDF as BDF
-from .common import OdeSolution as OdeSolution
-from .ivp import solve_ivp as solve_ivp
-from .lsoda import LSODA as LSODA
-from .radau import Radau as Radau
-from .rk import DOP853 as DOP853, RK23 as RK23, RK45 as RK45
+from .base import DenseOutput, OdeSolver
+from .bdf import BDF
+from .common import OdeSolution
+from .ivp import solve_ivp
+from .lsoda import LSODA
+from .radau import Radau
+from .rk import DOP853, RK23, RK45
+
+# NOTE: this __all__ does not exist at runtime
+__all__ = ["BDF", "DOP853", "LSODA", "RK23", "RK45", "DenseOutput", "OdeSolution", "OdeSolver", "Radau", "solve_ivp"]

scipy-stubs/integrate/_ivp/base.pyi

@@ -1,36 +1,62 @@
+from collections.abc import Callable, Sequence
+from typing import Any, ClassVar, Final, Generic, Literal, TypeAlias, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+import optype.numpy as onpt
 from scipy._typing import Untyped
 
-def check_arguments(fun, y0, support_complex) -> Untyped: ...
+_VT = TypeVar("_VT", bound=npt.NDArray[np.inexact[Any]], default=npt.NDArray[np.inexact[Any]])
+
+_ArrayLikeReal: TypeAlias = float | Sequence[float] | onpt.AnyFloatingArray | onpt.AnyIntegerArray
+_ArrayLikeComplex: TypeAlias = complex | Sequence[complex] | onpt.AnyComplexFloatingArray
+
+def check_arguments(
+    fun: Callable[[float, npt.NDArray[np.float64]], _ArrayLikeComplex],
+    y0: _ArrayLikeComplex,
+    support_complex: bool,
+) -> (
+    Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.float64]]
+    | Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.complex128]]
+): ...
 
 class OdeSolver:
-    TOO_SMALL_STEP: str
-    t_old: Untyped
-    t: Untyped
-    t_bound: Untyped
-    vectorized: Untyped
-    fun: Untyped
-    fun_single: Untyped
-    fun_vectorized: Untyped
-    direction: Untyped
-    n: Untyped
-    status: str
+    TOO_SMALL_STEP: ClassVar[str]
+    t: float
+    t_old: float
+    t_bound: float
+    vectorized: bool
+    fun: Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.float64]]
+    fun_single: Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.float64]]
+    fun_vectorized: Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.float64]]
+    direction: float
+    n: int
+    status: Literal["running", "finished", "failed"]
     nfev: int
     njev: int
     nlu: int
-    def __init__(self, fun, t0, y0, t_bound, vectorized, support_complex: bool = False): ...
+    def __init__(
+        self,
+        fun: Callable[[float, npt.NDArray[np.float64]], _ArrayLikeComplex],
+        t0: float,
+        y0: _ArrayLikeComplex,
+        t_bound: float,
+        vectorized: bool,
+        support_complex: bool = False,
+    ) -> None: ...
     @property
-    def step_size(self) -> Untyped: ...
-    def step(self) -> Untyped: ...
-    def dense_output(self) -> Untyped: ...
+    def step_size(self) -> float | None: ...
+    def step(self) -> str | None: ...
+    def dense_output(self) -> ConstantDenseOutput: ...
 
 class DenseOutput:
-    t_old: Untyped
-    t: Untyped
-    t_min: Untyped
-    t_max: Untyped
-    def __init__(self, t_old, t) -> None: ...
-    def __call__(self, t) -> Untyped: ...
+    t_old: Final[float]
+    t: Final[float]
+    t_min: Final[float]
+    t_max: Final[float]
+    def __init__(self, t_old: float, t: float, /) -> None: ...
+    def __call__(self, /, t: _ArrayLikeReal) -> Untyped: ...
 
-class ConstantDenseOutput(DenseOutput):
+class ConstantDenseOutput(DenseOutput, Generic[_VT]):
     value: Untyped
-    def __init__(self, t_old, t, value) -> None: ...
+    def __init__(self, t_old: float, t: float, /, value: _VT) -> None: ...