Skip to content

Commit 846f790

Browse files
authored
optimize: Generic InverseJacobian and related types (#429)
2 parents ddb04a2 + bfdfed3 commit 846f790

File tree

1 file changed

+134
-104
lines changed

1 file changed

+134
-104
lines changed

scipy-stubs/optimize/_nonlin.pyi

Lines changed: 134 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
import abc
22
from collections.abc import Callable
3-
from typing import Any, Final, Literal, Protocol, TypeAlias, TypedDict, type_check_only
4-
from typing_extensions import Unpack, override
3+
from typing import Final, Generic, Literal, Protocol, TypeAlias, TypedDict, overload, type_check_only
4+
from typing_extensions import TypeVar, Unpack, override
55

66
import numpy as np
7-
import numpy.typing as npt
87
import optype as op
98
import optype.numpy as onp
10-
from scipy.sparse import sparray, spmatrix
9+
from scipy._typing import Falsy, Truthy
10+
from scipy.sparse._base import _spbase
1111
from scipy.sparse.linalg import LinearOperator
1212

1313
__all__ = [
@@ -24,17 +24,11 @@ __all__ = [
2424
"newton_krylov",
2525
]
2626

27-
_Float: TypeAlias = np.floating[Any]
28-
_Complex: TypeAlias = np.complexfloating[Any, Any]
29-
_Inexact: TypeAlias = _Float | _Complex
30-
31-
_FloatND: TypeAlias = onp.ArrayND[_Float]
27+
_Floating: TypeAlias = np.float32 | np.float64
28+
_Inexact: TypeAlias = _Floating | np.complex64 | np.complex128
3229
_Inexact1D: TypeAlias = onp.Array1D[_Inexact]
33-
_Inexact2D: TypeAlias = onp.Array2D[_Inexact]
3430
_InexactND: TypeAlias = onp.ArrayND[_Inexact]
3531

36-
_SparseArray: TypeAlias = sparray | spmatrix
37-
3832
_JacobianMethod: TypeAlias = Literal[
3933
"krylov",
4034
"broyden1",
@@ -53,36 +47,47 @@ _Callback: TypeAlias = (
5347
)
5448
_ResidFunc: TypeAlias = Callable[[onp.ArrayND[np.float64]], onp.ToFloat] | Callable[[onp.ArrayND[np.complex128]], onp.ToFloat]
5549

50+
_InexactT = TypeVar("_InexactT", bound=_Inexact, default=_Inexact)
51+
_InexactT_co = TypeVar("_InexactT_co", bound=_Inexact, default=_Inexact, covariant=True)
52+
5653
_JacobianLike: TypeAlias = (
57-
Jacobian
58-
| type[Jacobian]
59-
| onp.ArrayND
60-
| _SparseArray
61-
| _SupportsJacobian
62-
| Callable[[onp.ArrayND[np.float64]], _FloatND | _SparseArray]
63-
| Callable[[onp.ArrayND[np.complex128]], _InexactND | _SparseArray]
64-
| str
54+
Jacobian[_InexactT]
55+
| type[Jacobian[_InexactT]]
56+
| onp.ArrayND[_InexactT]
57+
| _spbase[_InexactT]
58+
| _SupportsJacobian[_InexactT]
59+
| Callable[[onp.Array1D[np.float64]], onp.ArrayND[_InexactT] | _spbase[_InexactT]]
60+
| Callable[[onp.Array1D[np.complex128]], onp.ArrayND[_InexactT] | _spbase[_InexactT]]
6561
)
6662

6763
@type_check_only
68-
class _SupportsJacobian(Protocol):
64+
class _SupportsJacobian(Protocol[_InexactT_co]):
6965
@property
70-
def shape(self, /) -> tuple[int, ...]: ...
66+
def shape(self, /) -> tuple[int, int]: ...
7167
@property
72-
def dtype(self, /) -> np.dtype[np.generic]: ...
73-
def solve(self, /, v: onp.ToComplexND, tol: onp.ToFloat = 0) -> _Inexact2D: ...
68+
def dtype(self, /) -> np.dtype[_InexactT_co]: ...
69+
def solve(self, v: _InexactND, /, tol: float = 0) -> onp.ToComplex2D: ...
7470

7571
@type_check_only
76-
class _JacobianKwargs(TypedDict, total=False):
77-
solve: Callable[[_InexactND], _Inexact2D] | Callable[[_InexactND, onp.ToFloat], _Inexact2D]
78-
rsolve: Callable[[_InexactND], _Inexact2D] | Callable[[_InexactND, onp.ToFloat], _Inexact2D]
79-
matvec: Callable[[_InexactND], _Inexact1D] | Callable[[_InexactND, onp.ToFloat], _Inexact1D]
80-
rmatvec: Callable[[_InexactND], _Inexact1D] | Callable[[_InexactND, onp.ToFloat], _Inexact1D]
81-
matmat: Callable[[_InexactND], _Inexact2D]
82-
update: Callable[[_InexactND, _InexactND], None]
83-
todense: Callable[[], _Inexact2D]
84-
shape: tuple[int, ...]
85-
dtype: np.dtype[np.generic]
72+
class _JacobianKwargs(TypedDict, Generic[_InexactT_co], total=False):
73+
solve: Callable[[_InexactND], onp.Array2D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array2D[_InexactT_co]]
74+
rsolve: Callable[[_InexactND], onp.Array2D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array2D[_InexactT_co]]
75+
matvec: Callable[[_InexactND], onp.Array1D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array2D[_InexactT_co]]
76+
rmatvec: Callable[[_InexactND], onp.Array1D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array1D[_InexactT_co]]
77+
matmat: Callable[[_InexactND], onp.Array2D[_InexactT_co]]
78+
update: Callable[[_InexactND, onp.ArrayND[_InexactT_co]], None]
79+
todense: Callable[[], onp.Array2D[_InexactT_co]]
80+
shape: tuple[int, int]
81+
dtype: np.dtype[_InexactT_co]
82+
83+
#
84+
@type_check_only
85+
class _NonlinInfoDict(TypedDict):
86+
fun: float | np.float64
87+
nit: int
88+
status: int
89+
success: bool
90+
message: str
8691

8792
###
8893

@@ -111,35 +116,33 @@ class TerminationCondition:
111116
) -> None: ...
112117
def check(self, /, f: _InexactND, x: _InexactND, dx: _InexactND) -> int: ...
113118

114-
class Jacobian:
119+
class Jacobian(Generic[_InexactT_co]): # undocumented
115120
shape: Final[tuple[int, int]]
116-
dtype: Final[np.dtype[np.generic]]
121+
dtype: np.dtype[_InexactT_co]
117122
func: Final[_ResidFunc]
118123

119-
def __init__(self, /, **kw: Unpack[_JacobianKwargs]) -> None: ...
124+
def __init__(self, /, **kw: Unpack[_JacobianKwargs[_InexactT_co]]) -> None: ...
120125
#
121126
@abc.abstractmethod
122-
def solve(self, /, v: _InexactND, tol: float = 0) -> _Inexact2D: ...
127+
def solve(self, /, v: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
123128
# `x` and `F` are 1-d
124-
def setup(self, /, x: _InexactND, F: _InexactND, func: _ResidFunc) -> None: ...
125-
def update(self, /, x: _InexactND, F: _InexactND) -> None: ... # does nothing
129+
def setup(self: Jacobian[_InexactT], /, x: _InexactND, F: onp.ArrayND[_InexactT], func: _ResidFunc) -> None: ...
130+
def update(self: Jacobian[_InexactT], /, x: _InexactND, F: onp.ArrayND[_InexactT]) -> None: ... # does nothing
126131
def aspreconditioner(self, /) -> InverseJacobian: ...
127132

128-
class InverseJacobian: # undocumented
129-
jacobian: Final[Jacobian]
130-
matvec: Final[Callable[[_InexactND], _Inexact1D] | Callable[[_InexactND, onp.ToFloat], _Inexact1D]]
131-
rmatvec: Final[Callable[[_InexactND], _Inexact1D] | Callable[[_InexactND, onp.ToFloat], _Inexact1D]]
133+
class InverseJacobian(Generic[_InexactT_co]):
134+
jacobian: Jacobian[_InexactT_co]
135+
matvec: Callable[[_InexactND], onp.Array1D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array1D[_InexactT_co]]
136+
rmatvec: Callable[[_InexactND], onp.Array1D[_InexactT_co]] | Callable[[_InexactND, onp.ToFloat], onp.Array1D[_InexactT_co]]
132137

133138
@property
134-
def shape(self, /) -> tuple[int, ...]: ...
139+
def shape(self, /) -> tuple[int, int]: ...
135140
@property
136-
def dtype(self, /) -> np.dtype[np.generic]: ...
141+
def dtype(self, /) -> np.dtype[_InexactT_co]: ...
137142
#
138-
def __init__(self, /, jacobian: Jacobian) -> None: ...
139-
140-
def asjacobian(J: _JacobianLike) -> Jacobian: ... # undocumented
143+
def __init__(self, /, jacobian: Jacobian[_InexactT_co]) -> None: ...
141144

142-
class GenericBroyden(Jacobian, metaclass=abc.ABCMeta):
145+
class GenericBroyden(Jacobian[_InexactT_co], Generic[_InexactT_co], metaclass=abc.ABCMeta):
143146
alpha: Final[float | None]
144147
last_x: _Inexact1D
145148
last_f: float
@@ -149,97 +152,97 @@ class GenericBroyden(Jacobian, metaclass=abc.ABCMeta):
149152
@override
150153
def update(self, /, x0: _InexactND, f0: _InexactND) -> None: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
151154

152-
class LowRankMatrix:
155+
class LowRankMatrix(Generic[_InexactT_co]):
156+
dtype: np.dtype[_InexactT_co]
153157
alpha: Final[float]
154158
n: Final[int]
155159
cs: Final[list[_InexactND]]
156160
ds: Final[list[_InexactND]]
157-
dtype: Final[np.dtype[np.inexact[Any]]]
158161
collapsed: _InexactND | None
159162

160-
def __init__(self, /, alpha: float, n: int, dtype: np.dtype[np.inexact[Any]]) -> None: ...
161-
def __array__(self, /, dtype: npt.DTypeLike | None = None, copy: bool | None = None) -> _Inexact2D: ...
162-
def solve(self, /, v: _InexactND, tol: float = 0) -> _Inexact2D: ...
163-
def rsolve(self, /, v: _InexactND, tol: float = 0) -> _Inexact2D: ...
164-
def matvec(self, /, v: _InexactND) -> _Inexact1D: ...
165-
def rmatvec(self, /, v: _InexactND) -> _Inexact1D: ...
163+
def __init__(self, /, alpha: float, n: int, dtype: np.dtype[_InexactT_co]) -> None: ...
164+
def __array__(self, /, dtype: None = None, copy: None = None) -> onp.Array2D[_InexactT_co]: ...
165+
def solve(self, /, v: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
166+
def rsolve(self, /, v: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
167+
def matvec(self, /, v: _InexactND) -> onp.Array1D[_InexactT_co]: ...
168+
def rmatvec(self, /, v: _InexactND) -> onp.Array1D[_InexactT_co]: ...
166169
def append(self, /, c: _InexactND, d: _InexactND) -> None: ...
167170
def collapse(self, /) -> None: ...
168-
def restart_reduce(self, /, rank: float) -> None: ...
169-
def simple_reduce(self, /, rank: float) -> None: ...
170-
def svd_reduce(self, /, max_rank: float, to_retain: int | None = None) -> None: ...
171+
def restart_reduce(self, /, rank: int) -> None: ...
172+
def simple_reduce(self, /, rank: int) -> None: ...
173+
def svd_reduce(self, /, max_rank: int, to_retain: int | None = None) -> None: ...
171174

172-
class BroydenFirst(GenericBroyden):
173-
max_rank: Final[float]
174-
Gm: LowRankMatrix | None
175+
class BroydenFirst(GenericBroyden[_InexactT_co], Generic[_InexactT_co]):
176+
max_rank: Final[int]
177+
Gm: LowRankMatrix[_InexactT_co] | None
175178

176179
def __init__(
177180
self,
178181
/,
179182
alpha: float | None = None,
180183
reduction_method: _ReductionMethod = "restart",
181-
max_rank: float | None = None,
184+
max_rank: int | None = None,
182185
) -> None: ...
183186
@override
184-
def solve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
185-
def rsolve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ...
186-
def matvec(self, /, f: _InexactND) -> _Inexact1D: ...
187-
def rmatvec(self, /, f: _InexactND) -> _Inexact1D: ...
188-
def todense(self, /) -> _Inexact2D: ...
187+
def solve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
188+
def rsolve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
189+
def matvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
190+
def rmatvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
191+
def todense(self, /) -> onp.Array2D[_InexactT_co]: ...
189192

190-
class BroydenSecond(BroydenFirst): ...
193+
class BroydenSecond(BroydenFirst[_InexactT_co], Generic[_InexactT_co]): ...
191194

192-
class Anderson(GenericBroyden):
195+
class Anderson(GenericBroyden[_InexactT_co], Generic[_InexactT_co]):
193196
w0: Final[float]
194197
M: Final[float]
195-
dx: list[_InexactND]
196-
df: list[_InexactND]
197-
gamma: _InexactND | None
198+
dx: list[onp.Array1D[_InexactT_co]]
199+
df: list[onp.Array1D[_InexactT_co]]
200+
gamma: onp.ArrayND[_InexactT_co] | None
198201

199202
def __init__(self, /, alpha: float | None = None, w0: float = 0.01, M: float = 5) -> None: ...
200203
@override
201-
def solve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
202-
def matvec(self, /, f: _InexactND) -> _Inexact1D: ...
204+
def solve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
205+
def matvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
203206

204-
class DiagBroyden(GenericBroyden):
205-
d: _Inexact1D
207+
class DiagBroyden(GenericBroyden[_InexactT_co], Generic[_InexactT_co]):
208+
d: onp.Array1D[_InexactT_co]
206209

207210
def __init__(self, /, alpha: float | None = None) -> None: ...
208211
@override
209-
def solve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
210-
def rsolve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ...
211-
def matvec(self, /, f: _InexactND) -> _Inexact1D: ...
212-
def rmatvec(self, /, f: _InexactND) -> _Inexact1D: ...
213-
def todense(self, /) -> _Inexact2D: ...
212+
def solve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
213+
def rsolve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
214+
def matvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
215+
def rmatvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
216+
def todense(self, /) -> onp.Array2D[_InexactT_co]: ...
214217

215-
class LinearMixing(GenericBroyden):
218+
class LinearMixing(GenericBroyden[_InexactT_co], Generic[_InexactT_co]):
216219
def __init__(self, /, alpha: float | None = None) -> None: ...
217220
@override
218-
def solve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
219-
def rsolve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ...
220-
def matvec(self, /, f: _InexactND) -> _Inexact1D: ...
221-
def rmatvec(self, /, f: _InexactND) -> _Inexact1D: ...
222-
def todense(self, /) -> _Inexact2D: ...
221+
def solve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
222+
def rsolve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
223+
def matvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
224+
def rmatvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
225+
def todense(self, /) -> onp.Array2D[_InexactT_co]: ...
223226

224-
class ExcitingMixing(GenericBroyden):
227+
class ExcitingMixing(GenericBroyden[_InexactT_co], Generic[_InexactT_co]):
225228
alphamax: Final[float]
226-
beta: _Inexact1D | None
229+
beta: onp.Array1D[_InexactT_co] | None
227230

228231
def __init__(self, /, alpha: float | None = None, alphamax: float = 1.0) -> None: ...
229232
@override
230-
def solve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
231-
def rsolve(self, /, f: _InexactND, tol: float = 0) -> _Inexact2D: ...
232-
def matvec(self, /, f: _InexactND) -> _Inexact1D: ...
233-
def rmatvec(self, /, f: _InexactND) -> _Inexact1D: ...
234-
def todense(self, /) -> _Inexact2D: ...
233+
def solve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
234+
def rsolve(self, /, f: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ...
235+
def matvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
236+
def rmatvec(self, /, f: _InexactND) -> onp.Array1D[_InexactT_co]: ...
237+
def todense(self, /) -> onp.Array2D[_InexactT_co]: ...
235238

236-
class KrylovJacobian(Jacobian):
239+
class KrylovJacobian(Jacobian[_InexactT_co], Generic[_InexactT_co]):
237240
rdiff: Final[float]
238241
method: Final[_KrylovMethod]
239242
method_kw: Final[dict[str, object]]
240243
preconditioner: LinearOperator | InverseJacobian | None
241-
x0: _InexactND
242-
f0: _InexactND
244+
x0: _Inexact1D
245+
f0: _Inexact1D
243246
op: LinearOperator
244247

245248
def __init__(
@@ -252,17 +255,25 @@ class KrylovJacobian(Jacobian):
252255
outer_k: int = 10,
253256
**kw: object,
254257
) -> None: ...
255-
def matvec(self, /, v: _InexactND) -> _Inexact2D: ...
258+
def matvec(self, /, v: _InexactND) -> onp.Array2D[_InexactT_co]: ...
256259
@override
257-
def solve(self, /, rhs: _InexactND, tol: float = 0) -> _Inexact2D: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
260+
def solve(self, /, rhs: _InexactND, tol: float = 0) -> onp.Array2D[_InexactT_co]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
258261
@override
259262
def update(self, /, x: _InexactND, f: _InexactND) -> None: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
260263
@override
261264
def setup(self, /, x: _InexactND, f: _InexactND, func: _ResidFunc) -> None: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
262265

266+
# undocumented
267+
@overload
268+
def asjacobian(J: _JacobianLike[_InexactT]) -> Jacobian[_InexactT]: ...
269+
@overload
270+
def asjacobian(J: _JacobianMethod) -> Jacobian: ...
271+
272+
#
263273
def maxnorm(x: onp.ToComplexND) -> float | np.float64: ... # undocumented
264274

265-
# TOOD: overload on full_output
275+
#
276+
@overload
266277
def nonlin_solve(
267278
F: _ResidFunc,
268279
x0: onp.ToComplexND,
@@ -277,9 +288,28 @@ def nonlin_solve(
277288
tol_norm: onp.ToFloat | None = None,
278289
line_search: _LineSearch = "armijo",
279290
callback: _Callback | None = None,
280-
full_output: op.CanBool = False,
291+
full_output: Falsy = False,
281292
raise_exception: op.CanBool = True,
282293
) -> _InexactND: ...
294+
@overload
295+
def nonlin_solve(
296+
F: _ResidFunc,
297+
x0: onp.ToComplexND,
298+
jacobian: _JacobianMethod | _JacobianLike = "krylov",
299+
iter: onp.ToInt | None = None,
300+
verbose: op.CanBool = False,
301+
maxiter: onp.ToInt | None = None,
302+
f_tol: onp.ToFloat | None = None,
303+
f_rtol: onp.ToFloat | None = None,
304+
x_tol: onp.ToFloat | None = None,
305+
x_rtol: onp.ToFloat | None = None,
306+
tol_norm: onp.ToFloat | None = None,
307+
line_search: _LineSearch = "armijo",
308+
callback: _Callback | None = None,
309+
*,
310+
full_output: Truthy,
311+
raise_exception: op.CanBool = True,
312+
) -> tuple[_InexactND, _NonlinInfoDict]: ...
283313

284314
#
285315
def broyden1(

0 commit comments

Comments
 (0)