✨ optimize: improved linprog and milp annotations (#427)

scipy-stubs/_lib/_util.pyi

@@ -1,65 +1,67 @@
 import multiprocessing.pool as mpp
 import types
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from typing import Concatenate, Final, Generic, NamedTuple, TypeAlias, overload
-from typing_extensions import TypeVar, override
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, Concatenate, Final, Generic, Literal, NamedTuple, TypeAlias, overload
+from typing_extensions import Never, TypeVar, override
 
 import numpy as np
 import optype as op
 import optype.numpy as onp
 import optype.numpy.compat as npc
+from numpy.random import Generator as Generator  # implicit re-export
+from optype.numpy.compat import DTypePromotionError as DTypePromotionError  # implicit re-export
 from scipy._typing import RNG, EnterSelfMixin
 
-_AnyRNG = TypeVar("_AnyRNG", np.random.RandomState, np.random.Generator)
+_AnyRNGT = TypeVar("_AnyRNGT", np.random.RandomState, np.random.Generator)
 
-_T = TypeVar("_T", default=object)
-_T_co = TypeVar("_T_co", covariant=True, default=object)
-_T_contra = TypeVar("_T_contra", contravariant=True, default=object)
 _VT = TypeVar("_VT")
 _RT = TypeVar("_RT")
+
+_T = TypeVar("_T", default=Any)
+_T_co = TypeVar("_T_co", default=Any, covariant=True)
+_T_contra = TypeVar("_T_contra", default=Never, contravariant=True)
+
 _AxisT = TypeVar("_AxisT", bound=npc.integer)
 
 ###
 
-np_long: Final[type[np.int32 | np.int64]] = ...
-np_ulong: Final[type[np.uint32 | np.uint64]] = ...
-copy_if_needed: Final[bool | None] = ...
+np_long: Final[type[np.int32 | np.int64]] = ...  # `np.long` on `numpy>=2`, else `np.int_`
+np_ulong: Final[type[np.uint32 | np.uint64]] = ...  # `np.ulong` on `numpy>=2`, else `np.uint`
+copy_if_needed: Final[Literal[False] | None] = ...  # `None` on `numpy>=2`, otherwise `False`
 
+# NOTE: These aliases are implictly exported at runtime
 IntNumber: TypeAlias = int | npc.integer
 DecimalNumber: TypeAlias = float | npc.floating | npc.integer
-
 _RNG: TypeAlias = np.random.Generator | np.random.RandomState
 SeedType: TypeAlias = IntNumber | _RNG | None
-# NOTE: This is actually a exported at runtime :(
-GeneratorType = TypeVar("GeneratorType", bound=_RNG)  # noqa: PYI001
+GeneratorType = TypeVar("GeneratorType", bound=_RNG)  # noqa: PYI001  # oof
+
+###
 
 class ComplexWarning(RuntimeWarning): ...
 class VisibleDeprecationWarning(UserWarning): ...
-class DTypePromotionError(TypeError): ...
 
 class AxisError(ValueError, IndexError):
     _msg: Final[str | None]
     axis: Final[int | None]
-    ndim: Final[int | None]
-
+    ndim: Final[onp.NDim | None]
     @overload
     def __init__(self, /, axis: str, ndim: None = None, msg_prefix: None = None) -> None: ...
     @overload
-    def __init__(self, /, axis: int, ndim: int, msg_prefix: str | None = None) -> None: ...
+    def __init__(self, /, axis: int, ndim: onp.NDim, msg_prefix: str | None = None) -> None: ...
 
 class FullArgSpec(NamedTuple):
-    args: Sequence[str]
+    args: list[str]
     varargs: str | None
     varkw: str | None
-    defaults: tuple[object, ...] | None
-    kwonlyargs: Sequence[str]
-    kwonlydefaults: Mapping[str, object] | None
-    annotations: Mapping[str, type | object | str]
+    defaults: tuple[Any, ...] | None
+    kwonlyargs: list[str]
+    kwonlydefaults: dict[str, Any] | None
+    annotations: dict[str, Any]
 
 class _FunctionWrapper(Generic[_T_contra, _T_co]):
     f: Callable[Concatenate[_T_contra, ...], _T_co]
-    args: tuple[object, ...]
-
+    args: tuple[Any, ...]
     @overload
     def __init__(self, /, f: Callable[[_T_contra], _T_co], args: tuple[()]) -> None: ...
     @overload
@@ -69,8 +71,8 @@ class _FunctionWrapper(Generic[_T_contra, _T_co]):
 class MapWrapper(EnterSelfMixin):
     pool: int | mpp.Pool | None
 
-    def __init__(self, /, pool: Callable[[Callable[[_VT], _RT], Iterable[_VT]], Sequence[_RT]] | int = 1) -> None: ...
-    def __call__(self, /, func: Callable[[_VT], _RT], iterable: Iterable[_VT]) -> Sequence[_RT]: ...
+    def __init__(self, /, pool: Callable[[Callable[[_VT], _RT], Iterable[_VT]], Iterable[_RT]] | int = 1) -> None: ...
+    def __call__(self, /, func: Callable[[_VT], _RT], iterable: Iterable[_VT]) -> Iterable[_RT]: ...
     def terminate(self, /) -> None: ...
     def join(self, /) -> None: ...
     def close(self, /) -> None: ...
@@ -81,23 +83,23 @@ class _RichResult(dict[str, _T]):
     def __setattr__(self, name: str, value: _T, /) -> None: ...
 
 #
-def float_factorial(n: int) -> float: ...
+def float_factorial(n: op.CanIndex) -> float: ...  # will be `np.inf` if `n >= 171`
 
 #
 def getfullargspec_no_self(func: Callable[..., object]) -> FullArgSpec: ...
 
 #
 @overload
-def check_random_state(seed: _AnyRNG) -> _AnyRNG: ...
+def check_random_state(seed: _AnyRNGT) -> _AnyRNGT: ...
 @overload
-def check_random_state(seed: int | npc.integer | types.ModuleType | None) -> np.random.RandomState: ...
+def check_random_state(seed: onp.ToJustInt | types.ModuleType | None) -> np.random.RandomState: ...
 
 #
 @overload
 def rng_integers(
     gen: RNG | None,
-    low: onp.ToInt | onp.ToIntND,
-    high: onp.ToInt | onp.ToIntND | None = None,
+    low: onp.ToInt,
+    high: onp.ToInt | None = None,
     size: tuple[()] | None = None,
     dtype: onp.AnyIntegerDType = "int64",
     endpoint: op.CanBool = False,
@@ -114,8 +116,8 @@ def rng_integers(
 
 #
 @overload
-def normalize_axis_index(axis: int, ndim: int) -> int: ...
+def normalize_axis_index(axis: int, ndim: onp.NDim) -> onp.NDim: ...
 @overload
-def normalize_axis_index(axis: int, ndim: _AxisT) -> _AxisT: ...
+def normalize_axis_index(axis: int | _AxisT, ndim: _AxisT) -> _AxisT: ...
 @overload
-def normalize_axis_index(axis: _AxisT, ndim: int | _AxisT) -> _AxisT: ...
+def normalize_axis_index(axis: _AxisT, ndim: onp.NDim | _AxisT) -> _AxisT: ...

scipy-stubs/optimize/_linprog.pyi

@@ -1,41 +1,203 @@
-from collections.abc import Callable, Mapping, Sequence
-from typing import Final, Literal, type_check_only
+from collections.abc import Callable, Sequence
+from typing import Final, Literal, TypeAlias, TypedDict, overload, type_check_only
+from typing_extensions import LiteralString, deprecated
 
 import numpy as np
 import optype.numpy as onp
-from ._optimize import OptimizeResult
-from ._typing import Bound, MethodLinprog
+import optype.numpy.compat as npc
+from ._optimize import OptimizeResult as _OptimizeResult
+from ._typing import Bound, MethodLinprog, MethodLinprogLegacy
 
 __all__ = ["linprog", "linprog_terse_callback", "linprog_verbose_callback"]
 
+###
+
+_Ignored: TypeAlias = object
+_Max3: TypeAlias = Literal[0, 1, 2, 3]
+_Max4: TypeAlias = Literal[_Max3, 4]
+
+_Int: TypeAlias = int | np.int32 | np.int64
+_Float: TypeAlias = float | np.float64
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+
+@type_check_only
+class _OptionsCommon(TypedDict, total=False):
+    maxiter: _Int  # default: method-specific
+    disp: onp.ToBool  # default: False
+    presolve: onp.ToBool  # default: True
+
+# highs-ds
 @type_check_only
-class _OptimizeResult(OptimizeResult):
-    x: onp.ArrayND[np.float64]
-    fun: float
-    slack: onp.ArrayND[np.float64]
-    con: onp.ArrayND[np.float64]
-    success: bool
-    status: Literal[0, 1, 2, 3, 4]
-    nit: int
-    message: str
+class _OptionsHighsDS(_OptionsCommon, TypedDict, total=False):
+    time_limit: _Float  # default: np.finfo(float).max
+    dual_feasibility_tolerance: _Float  # default: 1e-7
+    primal_feasibility_tolerance: _Float  # default: 1e-7
+    simplex_dual_edge_weight_strategy: Literal["dantzig", "devex", "steepest", "steepest-devex"] | None  # default: None
+
+# highs-ips
+@type_check_only
+class _OptionsHighsIPM(_OptionsHighsDS, TypedDict, total=False):
+    ipm_optimality_tolerance: _Float  # default: 1e-8
+
+# highs
+@type_check_only
+class _OptionsHighs(_OptionsHighsIPM, TypedDict, total=False):
+    min_rel_gap: _Float | None  # default: None
+
+@type_check_only
+class _OptionsCommonLegacy(_OptionsCommon, TypedDict, total=False):
+    tol: _Float
+    autoscale: onp.ToBool  # default: False
+    rr: onp.ToBool  # default: True
+    rr_method: Literal["SVD", "pivot", "ID", "None"] | None  # default: None
+
+# interior-point (legacy, see https://github.com/scipy/scipy/issues/15707)
+@type_check_only
+class _OptionsInteriorPoint(_OptionsCommonLegacy, TypedDict, total=False):
+    alpha0: _Float  # default: 0.99995
+    beta: _Float  # default: 0.1
+    sparse: onp.ToBool  # default: False
+    lstq: onp.ToBool  # default: False
+    sym_pos: onp.ToBool  # default: True
+    cholsky: onp.ToBool  # default: True
+    pc: onp.ToBool  # default: True
+    ip: onp.ToBool  # default: False
+    perm_spec: Literal["NATURAL", "MMD_ATA", "MMD_AT_PLUS_A", "COLAMD"] | None  # default: "MMD_AT_PLUS_A"
+
+# revised simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
+@type_check_only
+class _OptionsRevisedSimplex(_OptionsCommonLegacy, TypedDict, total=False):
+    maxupdate: _Int  # default: 10
+    mast: onp.ToBool  # default: False
+    pivot: Literal["mrc", "bland"]
+
+# simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
+@type_check_only
+class _OptionsSimplex(_OptionsCommonLegacy, TypedDict, total=False):
+    bland: onp.ToBool  # default: False
+
+###
+
+__docformat__: Final = "restructuredtext en"  # undocumented
+LINPROG_METHODS: Final[Sequence[MethodLinprog | MethodLinprogLegacy]] = ...  # undocumented
 
-__docformat__: Final[str] = ...
-LINPROG_METHODS: Final[Sequence[MethodLinprog]] = ...
+class OptimizeResult(_OptimizeResult):
+    x: _Float1D  # minimizing decision variables w.r.t. the constraints
+    fun: _Float  # optimal objective function value
+    slack: _Float1D  # slack values; nominally positive
+    con: _Float1D  # residuals of equality constraints; nominally zero
+    status: _Max4
+    message: LiteralString
+    nit: int  # >=0
+    success: bool  # `success = status == 0`
 
 def linprog_verbose_callback(res: _OptimizeResult) -> None: ...
 def linprog_terse_callback(res: _OptimizeResult) -> None: ...
 
 #
+@overload  # highs (default)
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    method: Literal["highs"] = "highs",
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsHighs | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # highs-ds
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    *,
+    method: Literal["highs-ds"],
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsHighsDS | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # highs-ipm
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    *,
+    method: Literal["highs-ipm"],
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsHighsIPM | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # interior-point (legacy, see https://github.com/scipy/scipy/issues/15707)
+@deprecated("`method='interior-point'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    *,
+    method: Literal["interior-point"],
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsInteriorPoint | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # revised simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
+@deprecated("`method='revised simplex'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    *,
+    method: Literal["revised simplex"],
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsRevisedSimplex | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
+@deprecated("`method='simplex'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
+def linprog(
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
+    bounds: Bound = (0, None),
+    *,
+    method: Literal["simplex"],
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsSimplex | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
+) -> _OptimizeResult: ...
+@overload  # any "highs"
 def linprog(
-    c: onp.ToScalar | onp.ToArrayND,
-    A_ub: onp.ToScalar | onp.ToArrayND | None = None,
-    b_ub: onp.ToScalar | onp.ToArrayND | None = None,
-    A_eq: onp.ToScalar | onp.ToArrayND | None = None,
-    b_eq: onp.ToScalar | onp.ToArrayND | None = None,
+    c: onp.ToFloat1D,
+    A_ub: onp.ToFloat2D | None = None,
+    b_ub: onp.ToFloat1D | None = None,
+    A_eq: onp.ToFloat2D | None = None,
+    b_eq: onp.ToFloat1D | None = None,
     bounds: Bound = (0, None),
     method: MethodLinprog = "highs",
-    callback: Callable[[_OptimizeResult], None] | None = None,
-    options: Mapping[str, object] | None = None,
-    x0: onp.ToScalar | onp.ToArrayND | None = None,
-    integrality: onp.ToScalar | onp.ToArrayND | None = None,
+    callback: Callable[[_OptimizeResult], _Ignored] | None = None,
+    options: _OptionsHighs | None = None,
+    x0: onp.ToFloat1D | None = None,
+    integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
 ) -> _OptimizeResult: ...