Skip to content

Commit 25c6e67

Browse files
jorenhamjorenham
authored
✨ optimize: improved linprog and milp annotations (#427)
2 parents 462a15d + 9e60370 commit 25c6e67

File tree

5 files changed

+340
-124
lines changed

5 files changed

+340
-124
lines changed

‎scipy-stubs/_lib/_util.pyi

Lines changed: 36 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,67 @@
11
import multiprocessing.pool as mpp
22
import types
3-
from collections.abc import Callable, Iterable, Mapping, Sequence
4-
from typing import Concatenate, Final, Generic, NamedTuple, TypeAlias, overload
5-
from typing_extensions import TypeVar, override
3+
from collections.abc import Callable, Iterable, Sequence
4+
from typing import Any, Concatenate, Final, Generic, Literal, NamedTuple, TypeAlias, overload
5+
from typing_extensions import Never, TypeVar, override
66

77
import numpy as np
88
import optype as op
99
import optype.numpy as onp
1010
import optype.numpy.compat as npc
11+
from numpy.random import Generator as Generator # implicit re-export
12+
from optype.numpy.compat import DTypePromotionError as DTypePromotionError # implicit re-export
1113
from scipy._typing import RNG, EnterSelfMixin
1214

13-
_AnyRNG = TypeVar("_AnyRNG", np.random.RandomState, np.random.Generator)
15+
_AnyRNGT = TypeVar("_AnyRNGT", np.random.RandomState, np.random.Generator)
1416

15-
_T = TypeVar("_T", default=object)
16-
_T_co = TypeVar("_T_co", covariant=True, default=object)
17-
_T_contra = TypeVar("_T_contra", contravariant=True, default=object)
1817
_VT = TypeVar("_VT")
1918
_RT = TypeVar("_RT")
19+
20+
_T = TypeVar("_T", default=Any)
21+
_T_co = TypeVar("_T_co", default=Any, covariant=True)
22+
_T_contra = TypeVar("_T_contra", default=Never, contravariant=True)
23+
2024
_AxisT = TypeVar("_AxisT", bound=npc.integer)
2125

2226
###
2327

24-
np_long: Final[type[np.int32 | np.int64]] = ...
25-
np_ulong: Final[type[np.uint32 | np.uint64]] = ...
26-
copy_if_needed: Final[bool | None] = ...
28+
np_long: Final[type[np.int32 | np.int64]] = ... # `np.long` on `numpy>=2`, else `np.int_`
29+
np_ulong: Final[type[np.uint32 | np.uint64]] = ... # `np.ulong` on `numpy>=2`, else `np.uint`
30+
copy_if_needed: Final[Literal[False] | None] = ... # `None` on `numpy>=2`, otherwise `False`
2731

32+
# NOTE: These aliases are implictly exported at runtime
2833
IntNumber: TypeAlias = int | npc.integer
2934
DecimalNumber: TypeAlias = float | npc.floating | npc.integer
30-
3135
_RNG: TypeAlias = np.random.Generator | np.random.RandomState
3236
SeedType: TypeAlias = IntNumber | _RNG | None
33-
# NOTE: This is actually a exported at runtime :(
34-
GeneratorType = TypeVar("GeneratorType", bound=_RNG) # noqa: PYI001
37+
GeneratorType = TypeVar("GeneratorType", bound=_RNG) # noqa: PYI001 # oof
38+
39+
###
3540

3641
class ComplexWarning(RuntimeWarning): ...
3742
class VisibleDeprecationWarning(UserWarning): ...
38-
class DTypePromotionError(TypeError): ...
3943

4044
class AxisError(ValueError, IndexError):
4145
_msg: Final[str | None]
4246
axis: Final[int | None]
43-
ndim: Final[int | None]
44-
47+
ndim: Final[onp.NDim | None]
4548
@overload
4649
def __init__(self, /, axis: str, ndim: None = None, msg_prefix: None = None) -> None: ...
4750
@overload
48-
def __init__(self, /, axis: int, ndim: int, msg_prefix: str | None = None) -> None: ...
51+
def __init__(self, /, axis: int, ndim: onp.NDim, msg_prefix: str | None = None) -> None: ...
4952

5053
class FullArgSpec(NamedTuple):
51-
args: Sequence[str]
54+
args: list[str]
5255
varargs: str | None
5356
varkw: str | None
54-
defaults: tuple[object, ...] | None
55-
kwonlyargs: Sequence[str]
56-
kwonlydefaults: Mapping[str, object] | None
57-
annotations: Mapping[str, type | object | str]
57+
defaults: tuple[Any, ...] | None
58+
kwonlyargs: list[str]
59+
kwonlydefaults: dict[str, Any] | None
60+
annotations: dict[str, Any]
5861

5962
class _FunctionWrapper(Generic[_T_contra, _T_co]):
6063
f: Callable[Concatenate[_T_contra, ...], _T_co]
61-
args: tuple[object, ...]
62-
64+
args: tuple[Any, ...]
6365
@overload
6466
def __init__(self, /, f: Callable[[_T_contra], _T_co], args: tuple[()]) -> None: ...
6567
@overload
@@ -69,8 +71,8 @@ class _FunctionWrapper(Generic[_T_contra, _T_co]):
6971
class MapWrapper(EnterSelfMixin):
7072
pool: int | mpp.Pool | None
7173

72-
def __init__(self, /, pool: Callable[[Callable[[_VT], _RT], Iterable[_VT]], Sequence[_RT]] | int = 1) -> None: ...
73-
def __call__(self, /, func: Callable[[_VT], _RT], iterable: Iterable[_VT]) -> Sequence[_RT]: ...
74+
def __init__(self, /, pool: Callable[[Callable[[_VT], _RT], Iterable[_VT]], Iterable[_RT]] | int = 1) -> None: ...
75+
def __call__(self, /, func: Callable[[_VT], _RT], iterable: Iterable[_VT]) -> Iterable[_RT]: ...
7476
def terminate(self, /) -> None: ...
7577
def join(self, /) -> None: ...
7678
def close(self, /) -> None: ...
@@ -81,23 +83,23 @@ class _RichResult(dict[str, _T]):
8183
def __setattr__(self, name: str, value: _T, /) -> None: ...
8284

8385
#
84-
def float_factorial(n: int) -> float: ...
86+
def float_factorial(n: op.CanIndex) -> float: ... # will be `np.inf` if `n >= 171`
8587

8688
#
8789
def getfullargspec_no_self(func: Callable[..., object]) -> FullArgSpec: ...
8890

8991
#
9092
@overload
91-
def check_random_state(seed: _AnyRNG) -> _AnyRNG: ...
93+
def check_random_state(seed: _AnyRNGT) -> _AnyRNGT: ...
9294
@overload
93-
def check_random_state(seed: int | npc.integer | types.ModuleType | None) -> np.random.RandomState: ...
95+
def check_random_state(seed: onp.ToJustInt | types.ModuleType | None) -> np.random.RandomState: ...
9496

9597
#
9698
@overload
9799
def rng_integers(
98100
gen: RNG | None,
99-
low: onp.ToInt | onp.ToIntND,
100-
high: onp.ToInt | onp.ToIntND | None = None,
101+
low: onp.ToInt,
102+
high: onp.ToInt | None = None,
101103
size: tuple[()] | None = None,
102104
dtype: onp.AnyIntegerDType = "int64",
103105
endpoint: op.CanBool = False,
@@ -114,8 +116,8 @@ def rng_integers(
114116

115117
#
116118
@overload
117-
def normalize_axis_index(axis: int, ndim: int) -> int: ...
119+
def normalize_axis_index(axis: int, ndim: onp.NDim) -> onp.NDim: ...
118120
@overload
119-
def normalize_axis_index(axis: int, ndim: _AxisT) -> _AxisT: ...
121+
def normalize_axis_index(axis: int | _AxisT, ndim: _AxisT) -> _AxisT: ...
120122
@overload
121-
def normalize_axis_index(axis: _AxisT, ndim: int | _AxisT) -> _AxisT: ...
123+
def normalize_axis_index(axis: _AxisT, ndim: onp.NDim | _AxisT) -> _AxisT: ...

‎scipy-stubs/optimize/_linprog.pyi

Lines changed: 186 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,203 @@
1-
from collections.abc import Callable, Mapping, Sequence
2-
from typing import Final, Literal, type_check_only
1+
from collections.abc import Callable, Sequence
2+
from typing import Final, Literal, TypeAlias, TypedDict, overload, type_check_only
3+
from typing_extensions import LiteralString, deprecated
34

45
import numpy as np
56
import optype.numpy as onp
6-
from ._optimize import OptimizeResult
7-
from ._typing import Bound, MethodLinprog
7+
import optype.numpy.compat as npc
8+
from ._optimize import OptimizeResult as _OptimizeResult
9+
from ._typing import Bound, MethodLinprog, MethodLinprogLegacy
810

911
__all__ = ["linprog", "linprog_terse_callback", "linprog_verbose_callback"]
1012

13+
###
14+
15+
_Ignored: TypeAlias = object
16+
_Max3: TypeAlias = Literal[0, 1, 2, 3]
17+
_Max4: TypeAlias = Literal[_Max3, 4]
18+
19+
_Int: TypeAlias = int | np.int32 | np.int64
20+
_Float: TypeAlias = float | np.float64
21+
_Float1D: TypeAlias = onp.Array1D[np.float64]
22+
23+
@type_check_only
24+
class _OptionsCommon(TypedDict, total=False):
25+
maxiter: _Int # default: method-specific
26+
disp: onp.ToBool # default: False
27+
presolve: onp.ToBool # default: True
28+
29+
# highs-ds
1130
@type_check_only
12-
class _OptimizeResult(OptimizeResult):
13-
x: onp.ArrayND[np.float64]
14-
fun: float
15-
slack: onp.ArrayND[np.float64]
16-
con: onp.ArrayND[np.float64]
17-
success: bool
18-
status: Literal[0, 1, 2, 3, 4]
19-
nit: int
20-
message: str
31+
class _OptionsHighsDS(_OptionsCommon, TypedDict, total=False):
32+
time_limit: _Float # default: np.finfo(float).max
33+
dual_feasibility_tolerance: _Float # default: 1e-7
34+
primal_feasibility_tolerance: _Float # default: 1e-7
35+
simplex_dual_edge_weight_strategy: Literal["dantzig", "devex", "steepest", "steepest-devex"] | None # default: None
36+
37+
# highs-ips
38+
@type_check_only
39+
class _OptionsHighsIPM(_OptionsHighsDS, TypedDict, total=False):
40+
ipm_optimality_tolerance: _Float # default: 1e-8
41+
42+
# highs
43+
@type_check_only
44+
class _OptionsHighs(_OptionsHighsIPM, TypedDict, total=False):
45+
min_rel_gap: _Float | None # default: None
46+
47+
@type_check_only
48+
class _OptionsCommonLegacy(_OptionsCommon, TypedDict, total=False):
49+
tol: _Float
50+
autoscale: onp.ToBool # default: False
51+
rr: onp.ToBool # default: True
52+
rr_method: Literal["SVD", "pivot", "ID", "None"] | None # default: None
53+
54+
# interior-point (legacy, see https://github.com/scipy/scipy/issues/15707)
55+
@type_check_only
56+
class _OptionsInteriorPoint(_OptionsCommonLegacy, TypedDict, total=False):
57+
alpha0: _Float # default: 0.99995
58+
beta: _Float # default: 0.1
59+
sparse: onp.ToBool # default: False
60+
lstq: onp.ToBool # default: False
61+
sym_pos: onp.ToBool # default: True
62+
cholsky: onp.ToBool # default: True
63+
pc: onp.ToBool # default: True
64+
ip: onp.ToBool # default: False
65+
perm_spec: Literal["NATURAL", "MMD_ATA", "MMD_AT_PLUS_A", "COLAMD"] | None # default: "MMD_AT_PLUS_A"
66+
67+
# revised simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
68+
@type_check_only
69+
class _OptionsRevisedSimplex(_OptionsCommonLegacy, TypedDict, total=False):
70+
maxupdate: _Int # default: 10
71+
mast: onp.ToBool # default: False
72+
pivot: Literal["mrc", "bland"]
73+
74+
# simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
75+
@type_check_only
76+
class _OptionsSimplex(_OptionsCommonLegacy, TypedDict, total=False):
77+
bland: onp.ToBool # default: False
78+
79+
###
80+
81+
__docformat__: Final = "restructuredtext en" # undocumented
82+
LINPROG_METHODS: Final[Sequence[MethodLinprog | MethodLinprogLegacy]] = ... # undocumented
2183

22-
__docformat__: Final[str] = ...
23-
LINPROG_METHODS: Final[Sequence[MethodLinprog]] = ...
84+
class OptimizeResult(_OptimizeResult):
85+
x: _Float1D # minimizing decision variables w.r.t. the constraints
86+
fun: _Float # optimal objective function value
87+
slack: _Float1D # slack values; nominally positive
88+
con: _Float1D # residuals of equality constraints; nominally zero
89+
status: _Max4
90+
message: LiteralString
91+
nit: int # >=0
92+
success: bool # `success = status == 0`
2493

2594
def linprog_verbose_callback(res: _OptimizeResult) -> None: ...
2695
def linprog_terse_callback(res: _OptimizeResult) -> None: ...
2796

2897
#
98+
@overload # highs (default)
99+
def linprog(
100+
c: onp.ToFloat1D,
101+
A_ub: onp.ToFloat2D | None = None,
102+
b_ub: onp.ToFloat1D | None = None,
103+
A_eq: onp.ToFloat2D | None = None,
104+
b_eq: onp.ToFloat1D | None = None,
105+
bounds: Bound = (0, None),
106+
method: Literal["highs"] = "highs",
107+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
108+
options: _OptionsHighs | None = None,
109+
x0: onp.ToFloat1D | None = None,
110+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
111+
) -> _OptimizeResult: ...
112+
@overload # highs-ds
113+
def linprog(
114+
c: onp.ToFloat1D,
115+
A_ub: onp.ToFloat2D | None = None,
116+
b_ub: onp.ToFloat1D | None = None,
117+
A_eq: onp.ToFloat2D | None = None,
118+
b_eq: onp.ToFloat1D | None = None,
119+
bounds: Bound = (0, None),
120+
*,
121+
method: Literal["highs-ds"],
122+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
123+
options: _OptionsHighsDS | None = None,
124+
x0: onp.ToFloat1D | None = None,
125+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
126+
) -> _OptimizeResult: ...
127+
@overload # highs-ipm
128+
def linprog(
129+
c: onp.ToFloat1D,
130+
A_ub: onp.ToFloat2D | None = None,
131+
b_ub: onp.ToFloat1D | None = None,
132+
A_eq: onp.ToFloat2D | None = None,
133+
b_eq: onp.ToFloat1D | None = None,
134+
bounds: Bound = (0, None),
135+
*,
136+
method: Literal["highs-ipm"],
137+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
138+
options: _OptionsHighsIPM | None = None,
139+
x0: onp.ToFloat1D | None = None,
140+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
141+
) -> _OptimizeResult: ...
142+
@overload # interior-point (legacy, see https://github.com/scipy/scipy/issues/15707)
143+
@deprecated("`method='interior-point'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
144+
def linprog(
145+
c: onp.ToFloat1D,
146+
A_ub: onp.ToFloat2D | None = None,
147+
b_ub: onp.ToFloat1D | None = None,
148+
A_eq: onp.ToFloat2D | None = None,
149+
b_eq: onp.ToFloat1D | None = None,
150+
bounds: Bound = (0, None),
151+
*,
152+
method: Literal["interior-point"],
153+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
154+
options: _OptionsInteriorPoint | None = None,
155+
x0: onp.ToFloat1D | None = None,
156+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
157+
) -> _OptimizeResult: ...
158+
@overload # revised simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
159+
@deprecated("`method='revised simplex'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
160+
def linprog(
161+
c: onp.ToFloat1D,
162+
A_ub: onp.ToFloat2D | None = None,
163+
b_ub: onp.ToFloat1D | None = None,
164+
A_eq: onp.ToFloat2D | None = None,
165+
b_eq: onp.ToFloat1D | None = None,
166+
bounds: Bound = (0, None),
167+
*,
168+
method: Literal["revised simplex"],
169+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
170+
options: _OptionsRevisedSimplex | None = None,
171+
x0: onp.ToFloat1D | None = None,
172+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
173+
) -> _OptimizeResult: ...
174+
@overload # simplex (legacy, see https://github.com/scipy/scipy/issues/15707)
175+
@deprecated("`method='simplex'` is deprecated and will be removed in SciPy 1.16.0. Please use one of the HIGHS solvers.")
176+
def linprog(
177+
c: onp.ToFloat1D,
178+
A_ub: onp.ToFloat2D | None = None,
179+
b_ub: onp.ToFloat1D | None = None,
180+
A_eq: onp.ToFloat2D | None = None,
181+
b_eq: onp.ToFloat1D | None = None,
182+
bounds: Bound = (0, None),
183+
*,
184+
method: Literal["simplex"],
185+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
186+
options: _OptionsSimplex | None = None,
187+
x0: onp.ToFloat1D | None = None,
188+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
189+
) -> _OptimizeResult: ...
190+
@overload # any "highs"
29191
def linprog(
30-
c: onp.ToScalar | onp.ToArrayND,
31-
A_ub: onp.ToScalar | onp.ToArrayND | None = None,
32-
b_ub: onp.ToScalar | onp.ToArrayND | None = None,
33-
A_eq: onp.ToScalar | onp.ToArrayND | None = None,
34-
b_eq: onp.ToScalar | onp.ToArrayND | None = None,
192+
c: onp.ToFloat1D,
193+
A_ub: onp.ToFloat2D | None = None,
194+
b_ub: onp.ToFloat1D | None = None,
195+
A_eq: onp.ToFloat2D | None = None,
196+
b_eq: onp.ToFloat1D | None = None,
35197
bounds: Bound = (0, None),
36198
method: MethodLinprog = "highs",
37-
callback: Callable[[_OptimizeResult], None] | None = None,
38-
options: Mapping[str, object] | None = None,
39-
x0: onp.ToScalar | onp.ToArrayND | None = None,
40-
integrality: onp.ToScalar | onp.ToArrayND | None = None,
199+
callback: Callable[[_OptimizeResult], _Ignored] | None = None,
200+
options: _OptionsHighs | None = None,
201+
x0: onp.ToFloat1D | None = None,
202+
integrality: _Max3 | Sequence[_Max3] | onp.CanArrayND[npc.integer] | None = None,
41203
) -> _OptimizeResult: ...

0 commit comments

Comments
 (0)