1
1
import abc
2
2
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
5
5
6
6
import numpy as np
7
- import numpy .typing as npt
8
7
import optype as op
9
8
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
11
11
from scipy .sparse .linalg import LinearOperator
12
12
13
13
__all__ = [
@@ -24,17 +24,11 @@ __all__ = [
24
24
"newton_krylov" ,
25
25
]
26
26
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
32
29
_Inexact1D : TypeAlias = onp .Array1D [_Inexact ]
33
- _Inexact2D : TypeAlias = onp .Array2D [_Inexact ]
34
30
_InexactND : TypeAlias = onp .ArrayND [_Inexact ]
35
31
36
- _SparseArray : TypeAlias = sparray | spmatrix
37
-
38
32
_JacobianMethod : TypeAlias = Literal [
39
33
"krylov" ,
40
34
"broyden1" ,
@@ -53,36 +47,47 @@ _Callback: TypeAlias = (
53
47
)
54
48
_ResidFunc : TypeAlias = Callable [[onp .ArrayND [np .float64 ]], onp .ToFloat ] | Callable [[onp .ArrayND [np .complex128 ]], onp .ToFloat ]
55
49
50
+ _InexactT = TypeVar ("_InexactT" , bound = _Inexact , default = _Inexact )
51
+ _InexactT_co = TypeVar ("_InexactT_co" , bound = _Inexact , default = _Inexact , covariant = True )
52
+
56
53
_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 ]]
65
61
)
66
62
67
63
@type_check_only
68
- class _SupportsJacobian (Protocol ):
64
+ class _SupportsJacobian (Protocol [ _InexactT_co ] ):
69
65
@property
70
- def shape (self , / ) -> tuple [int , ... ]: ...
66
+ def shape (self , / ) -> tuple [int , int ]: ...
71
67
@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 : ...
74
70
75
71
@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
86
91
87
92
###
88
93
@@ -111,35 +116,33 @@ class TerminationCondition:
111
116
) -> None : ...
112
117
def check (self , / , f : _InexactND , x : _InexactND , dx : _InexactND ) -> int : ...
113
118
114
- class Jacobian :
119
+ class Jacobian ( Generic [ _InexactT_co ]): # undocumented
115
120
shape : Final [tuple [int , int ]]
116
- dtype : Final [ np .dtype [np . generic ] ]
121
+ dtype : np .dtype [_InexactT_co ]
117
122
func : Final [_ResidFunc ]
118
123
119
- def __init__ (self , / , ** kw : Unpack [_JacobianKwargs ]) -> None : ...
124
+ def __init__ (self , / , ** kw : Unpack [_JacobianKwargs [ _InexactT_co ] ]) -> None : ...
120
125
#
121
126
@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 ] : ...
123
128
# `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
126
131
def aspreconditioner (self , / ) -> InverseJacobian : ...
127
132
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 ]]
132
137
133
138
@property
134
- def shape (self , / ) -> tuple [int , ... ]: ...
139
+ def shape (self , / ) -> tuple [int , int ]: ...
135
140
@property
136
- def dtype (self , / ) -> np .dtype [np . generic ]: ...
141
+ def dtype (self , / ) -> np .dtype [_InexactT_co ]: ...
137
142
#
138
- def __init__ (self , / , jacobian : Jacobian ) -> None : ...
139
-
140
- def asjacobian (J : _JacobianLike ) -> Jacobian : ... # undocumented
143
+ def __init__ (self , / , jacobian : Jacobian [_InexactT_co ]) -> None : ...
141
144
142
- class GenericBroyden (Jacobian , metaclass = abc .ABCMeta ):
145
+ class GenericBroyden (Jacobian [ _InexactT_co ], Generic [ _InexactT_co ] , metaclass = abc .ABCMeta ):
143
146
alpha : Final [float | None ]
144
147
last_x : _Inexact1D
145
148
last_f : float
@@ -149,97 +152,97 @@ class GenericBroyden(Jacobian, metaclass=abc.ABCMeta):
149
152
@override
150
153
def update (self , / , x0 : _InexactND , f0 : _InexactND ) -> None : ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
151
154
152
- class LowRankMatrix :
155
+ class LowRankMatrix (Generic [_InexactT_co ]):
156
+ dtype : np .dtype [_InexactT_co ]
153
157
alpha : Final [float ]
154
158
n : Final [int ]
155
159
cs : Final [list [_InexactND ]]
156
160
ds : Final [list [_InexactND ]]
157
- dtype : Final [np .dtype [np .inexact [Any ]]]
158
161
collapsed : _InexactND | None
159
162
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 ] : ...
166
169
def append (self , / , c : _InexactND , d : _InexactND ) -> None : ...
167
170
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 : ...
171
174
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
175
178
176
179
def __init__ (
177
180
self ,
178
181
/ ,
179
182
alpha : float | None = None ,
180
183
reduction_method : _ReductionMethod = "restart" ,
181
- max_rank : float | None = None ,
184
+ max_rank : int | None = None ,
182
185
) -> None : ...
183
186
@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 ] : ...
189
192
190
- class BroydenSecond (BroydenFirst ): ...
193
+ class BroydenSecond (BroydenFirst [ _InexactT_co ], Generic [ _InexactT_co ] ): ...
191
194
192
- class Anderson (GenericBroyden ):
195
+ class Anderson (GenericBroyden [ _InexactT_co ], Generic [ _InexactT_co ] ):
193
196
w0 : Final [float ]
194
197
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
198
201
199
202
def __init__ (self , / , alpha : float | None = None , w0 : float = 0.01 , M : float = 5 ) -> None : ...
200
203
@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 ] : ...
203
206
204
- class DiagBroyden (GenericBroyden ):
205
- d : _Inexact1D
207
+ class DiagBroyden (GenericBroyden [ _InexactT_co ], Generic [ _InexactT_co ] ):
208
+ d : onp . Array1D [ _InexactT_co ]
206
209
207
210
def __init__ (self , / , alpha : float | None = None ) -> None : ...
208
211
@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 ] : ...
214
217
215
- class LinearMixing (GenericBroyden ):
218
+ class LinearMixing (GenericBroyden [ _InexactT_co ], Generic [ _InexactT_co ] ):
216
219
def __init__ (self , / , alpha : float | None = None ) -> None : ...
217
220
@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 ] : ...
223
226
224
- class ExcitingMixing (GenericBroyden ):
227
+ class ExcitingMixing (GenericBroyden [ _InexactT_co ], Generic [ _InexactT_co ] ):
225
228
alphamax : Final [float ]
226
- beta : _Inexact1D | None
229
+ beta : onp . Array1D [ _InexactT_co ] | None
227
230
228
231
def __init__ (self , / , alpha : float | None = None , alphamax : float = 1.0 ) -> None : ...
229
232
@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 ] : ...
235
238
236
- class KrylovJacobian (Jacobian ):
239
+ class KrylovJacobian (Jacobian [ _InexactT_co ], Generic [ _InexactT_co ] ):
237
240
rdiff : Final [float ]
238
241
method : Final [_KrylovMethod ]
239
242
method_kw : Final [dict [str , object ]]
240
243
preconditioner : LinearOperator | InverseJacobian | None
241
- x0 : _InexactND
242
- f0 : _InexactND
244
+ x0 : _Inexact1D
245
+ f0 : _Inexact1D
243
246
op : LinearOperator
244
247
245
248
def __init__ (
@@ -252,17 +255,25 @@ class KrylovJacobian(Jacobian):
252
255
outer_k : int = 10 ,
253
256
** kw : object ,
254
257
) -> None : ...
255
- def matvec (self , / , v : _InexactND ) -> _Inexact2D : ...
258
+ def matvec (self , / , v : _InexactND ) -> onp . Array2D [ _InexactT_co ] : ...
256
259
@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]
258
261
@override
259
262
def update (self , / , x : _InexactND , f : _InexactND ) -> None : ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
260
263
@override
261
264
def setup (self , / , x : _InexactND , f : _InexactND , func : _ResidFunc ) -> None : ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
262
265
266
+ # undocumented
267
+ @overload
268
+ def asjacobian (J : _JacobianLike [_InexactT ]) -> Jacobian [_InexactT ]: ...
269
+ @overload
270
+ def asjacobian (J : _JacobianMethod ) -> Jacobian : ...
271
+
272
+ #
263
273
def maxnorm (x : onp .ToComplexND ) -> float | np .float64 : ... # undocumented
264
274
265
- # TOOD: overload on full_output
275
+ #
276
+ @overload
266
277
def nonlin_solve (
267
278
F : _ResidFunc ,
268
279
x0 : onp .ToComplexND ,
@@ -277,9 +288,28 @@ def nonlin_solve(
277
288
tol_norm : onp .ToFloat | None = None ,
278
289
line_search : _LineSearch = "armijo" ,
279
290
callback : _Callback | None = None ,
280
- full_output : op . CanBool = False ,
291
+ full_output : Falsy = False ,
281
292
raise_exception : op .CanBool = True ,
282
293
) -> _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 ]: ...
283
313
284
314
#
285
315
def broyden1 (
0 commit comments