Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/scipy/integrate/_ivp/radau.py : 12%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2from scipy.linalg import lu_factor, lu_solve
3from scipy.sparse import csc_matrix, issparse, eye
4from scipy.sparse.linalg import splu
5from scipy.optimize._numdiff import group_columns
6from .common import (validate_max_step, validate_tol, select_initial_step,
7 norm, num_jac, EPS, warn_extraneous,
8 validate_first_step)
9from .base import OdeSolver, DenseOutput
11S6 = 6 ** 0.5
13# Butcher tableau. A is not used directly, see below.
14C = np.array([(4 - S6) / 10, (4 + S6) / 10, 1])
15E = np.array([-13 - 7 * S6, -13 + 7 * S6, -1]) / 3
17# Eigendecomposition of A is done: A = T L T**-1. There is 1 real eigenvalue
18# and a complex conjugate pair. They are written below.
19MU_REAL = 3 + 3 ** (2 / 3) - 3 ** (1 / 3)
20MU_COMPLEX = (3 + 0.5 * (3 ** (1 / 3) - 3 ** (2 / 3))
21 - 0.5j * (3 ** (5 / 6) + 3 ** (7 / 6)))
23# These are transformation matrices.
24T = np.array([
25 [0.09443876248897524, -0.14125529502095421, 0.03002919410514742],
26 [0.25021312296533332, 0.20412935229379994, -0.38294211275726192],
27 [1, 1, 0]])
28TI = np.array([
29 [4.17871859155190428, 0.32768282076106237, 0.52337644549944951],
30 [-4.17871859155190428, -0.32768282076106237, 0.47662355450055044],
31 [0.50287263494578682, -2.57192694985560522, 0.59603920482822492]])
32# These linear combinations are used in the algorithm.
33TI_REAL = TI[0]
34TI_COMPLEX = TI[1] + 1j * TI[2]
36# Interpolator coefficients.
37P = np.array([
38 [13/3 + 7*S6/3, -23/3 - 22*S6/3, 10/3 + 5 * S6],
39 [13/3 - 7*S6/3, -23/3 + 22*S6/3, 10/3 - 5 * S6],
40 [1/3, -8/3, 10/3]])
43NEWTON_MAXITER = 6 # Maximum number of Newton iterations.
44MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size.
45MAX_FACTOR = 10 # Maximum allowed increase in a step size.
48def solve_collocation_system(fun, t, y, h, Z0, scale, tol,
49 LU_real, LU_complex, solve_lu):
50 """Solve the collocation system.
52 Parameters
53 ----------
54 fun : callable
55 Right-hand side of the system.
56 t : float
57 Current time.
58 y : ndarray, shape (n,)
59 Current state.
60 h : float
61 Step to try.
62 Z0 : ndarray, shape (3, n)
63 Initial guess for the solution. It determines new values of `y` at
64 ``t + h * C`` as ``y + Z0``, where ``C`` is the Radau method constants.
65 scale : float
66 Problem tolerance scale, i.e. ``rtol * abs(y) + atol``.
67 tol : float
68 Tolerance to which solve the system. This value is compared with
69 the normalized by `scale` error.
70 LU_real, LU_complex
71 LU decompositions of the system Jacobians.
72 solve_lu : callable
73 Callable which solves a linear system given a LU decomposition. The
74 signature is ``solve_lu(LU, b)``.
76 Returns
77 -------
78 converged : bool
79 Whether iterations converged.
80 n_iter : int
81 Number of completed iterations.
82 Z : ndarray, shape (3, n)
83 Found solution.
84 rate : float
85 The rate of convergence.
86 """
87 n = y.shape[0]
88 M_real = MU_REAL / h
89 M_complex = MU_COMPLEX / h
91 W = TI.dot(Z0)
92 Z = Z0
94 F = np.empty((3, n))
95 ch = h * C
97 dW_norm_old = None
98 dW = np.empty_like(W)
99 converged = False
100 rate = None
101 for k in range(NEWTON_MAXITER):
102 for i in range(3):
103 F[i] = fun(t + ch[i], y + Z[i])
105 if not np.all(np.isfinite(F)):
106 break
108 f_real = F.T.dot(TI_REAL) - M_real * W[0]
109 f_complex = F.T.dot(TI_COMPLEX) - M_complex * (W[1] + 1j * W[2])
111 dW_real = solve_lu(LU_real, f_real)
112 dW_complex = solve_lu(LU_complex, f_complex)
114 dW[0] = dW_real
115 dW[1] = dW_complex.real
116 dW[2] = dW_complex.imag
118 dW_norm = norm(dW / scale)
119 if dW_norm_old is not None:
120 rate = dW_norm / dW_norm_old
122 if (rate is not None and (rate >= 1 or
123 rate ** (NEWTON_MAXITER - k) / (1 - rate) * dW_norm > tol)):
124 break
126 W += dW
127 Z = T.dot(W)
129 if (dW_norm == 0 or
130 rate is not None and rate / (1 - rate) * dW_norm < tol):
131 converged = True
132 break
134 dW_norm_old = dW_norm
136 return converged, k + 1, Z, rate
139def predict_factor(h_abs, h_abs_old, error_norm, error_norm_old):
140 """Predict by which factor to increase/decrease the step size.
142 The algorithm is described in [1]_.
144 Parameters
145 ----------
146 h_abs, h_abs_old : float
147 Current and previous values of the step size, `h_abs_old` can be None
148 (see Notes).
149 error_norm, error_norm_old : float
150 Current and previous values of the error norm, `error_norm_old` can
151 be None (see Notes).
153 Returns
154 -------
155 factor : float
156 Predicted factor.
158 Notes
159 -----
160 If `h_abs_old` and `error_norm_old` are both not None then a two-step
161 algorithm is used, otherwise a one-step algorithm is used.
163 References
164 ----------
165 .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential
166 Equations II: Stiff and Differential-Algebraic Problems", Sec. IV.8.
167 """
168 if error_norm_old is None or h_abs_old is None or error_norm == 0:
169 multiplier = 1
170 else:
171 multiplier = h_abs / h_abs_old * (error_norm_old / error_norm) ** 0.25
173 with np.errstate(divide='ignore'):
174 factor = min(1, multiplier) * error_norm ** -0.25
176 return factor
179class Radau(OdeSolver):
180 """Implicit Runge-Kutta method of Radau IIA family of order 5.
182 The implementation follows [1]_. The error is controlled with a
183 third-order accurate embedded formula. A cubic polynomial which satisfies
184 the collocation conditions is used for the dense output.
186 Parameters
187 ----------
188 fun : callable
189 Right-hand side of the system. The calling signature is ``fun(t, y)``.
190 Here ``t`` is a scalar, and there are two options for the ndarray ``y``:
191 It can either have shape (n,); then ``fun`` must return array_like with
192 shape (n,). Alternatively it can have shape (n, k); then ``fun``
193 must return an array_like with shape (n, k), i.e., each column
194 corresponds to a single column in ``y``. The choice between the two
195 options is determined by `vectorized` argument (see below). The
196 vectorized implementation allows a faster approximation of the Jacobian
197 by finite differences (required for this solver).
198 t0 : float
199 Initial time.
200 y0 : array_like, shape (n,)
201 Initial state.
202 t_bound : float
203 Boundary time - the integration won't continue beyond it. It also
204 determines the direction of the integration.
205 first_step : float or None, optional
206 Initial step size. Default is ``None`` which means that the algorithm
207 should choose.
208 max_step : float, optional
209 Maximum allowed step size. Default is np.inf, i.e., the step size is not
210 bounded and determined solely by the solver.
211 rtol, atol : float and array_like, optional
212 Relative and absolute tolerances. The solver keeps the local error
213 estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a
214 relative accuracy (number of correct digits). But if a component of `y`
215 is approximately below `atol`, the error only needs to fall within
216 the same `atol` threshold, and the number of correct digits is not
217 guaranteed. If components of y have different scales, it might be
218 beneficial to set different `atol` values for different components by
219 passing array_like with shape (n,) for `atol`. Default values are
220 1e-3 for `rtol` and 1e-6 for `atol`.
221 jac : {None, array_like, sparse_matrix, callable}, optional
222 Jacobian matrix of the right-hand side of the system with respect to
223 y, required by this method. The Jacobian matrix has shape (n, n) and
224 its element (i, j) is equal to ``d f_i / d y_j``.
225 There are three ways to define the Jacobian:
227 * If array_like or sparse_matrix, the Jacobian is assumed to
228 be constant.
229 * If callable, the Jacobian is assumed to depend on both
230 t and y; it will be called as ``jac(t, y)`` as necessary.
231 For the 'Radau' and 'BDF' methods, the return value might be a
232 sparse matrix.
233 * If None (default), the Jacobian will be approximated by
234 finite differences.
236 It is generally recommended to provide the Jacobian rather than
237 relying on a finite-difference approximation.
238 jac_sparsity : {None, array_like, sparse matrix}, optional
239 Defines a sparsity structure of the Jacobian matrix for a
240 finite-difference approximation. Its shape must be (n, n). This argument
241 is ignored if `jac` is not `None`. If the Jacobian has only few non-zero
242 elements in *each* row, providing the sparsity structure will greatly
243 speed up the computations [2]_. A zero entry means that a corresponding
244 element in the Jacobian is always zero. If None (default), the Jacobian
245 is assumed to be dense.
246 vectorized : bool, optional
247 Whether `fun` is implemented in a vectorized fashion. Default is False.
249 Attributes
250 ----------
251 n : int
252 Number of equations.
253 status : string
254 Current status of the solver: 'running', 'finished' or 'failed'.
255 t_bound : float
256 Boundary time.
257 direction : float
258 Integration direction: +1 or -1.
259 t : float
260 Current time.
261 y : ndarray
262 Current state.
263 t_old : float
264 Previous time. None if no steps were made yet.
265 step_size : float
266 Size of the last successful step. None if no steps were made yet.
267 nfev : int
268 Number of evaluations of the right-hand side.
269 njev : int
270 Number of evaluations of the Jacobian.
271 nlu : int
272 Number of LU decompositions.
274 References
275 ----------
276 .. [1] E. Hairer, G. Wanner, "Solving Ordinary Differential Equations II:
277 Stiff and Differential-Algebraic Problems", Sec. IV.8.
278 .. [2] A. Curtis, M. J. D. Powell, and J. Reid, "On the estimation of
279 sparse Jacobian matrices", Journal of the Institute of Mathematics
280 and its Applications, 13, pp. 117-120, 1974.
281 """
282 def __init__(self, fun, t0, y0, t_bound, max_step=np.inf,
283 rtol=1e-3, atol=1e-6, jac=None, jac_sparsity=None,
284 vectorized=False, first_step=None, **extraneous):
285 warn_extraneous(extraneous)
286 super(Radau, self).__init__(fun, t0, y0, t_bound, vectorized)
287 self.y_old = None
288 self.max_step = validate_max_step(max_step)
289 self.rtol, self.atol = validate_tol(rtol, atol, self.n)
290 self.f = self.fun(self.t, self.y)
291 # Select initial step assuming the same order which is used to control
292 # the error.
293 if first_step is None:
294 self.h_abs = select_initial_step(
295 self.fun, self.t, self.y, self.f, self.direction,
296 3, self.rtol, self.atol)
297 else:
298 self.h_abs = validate_first_step(first_step, t0, t_bound)
299 self.h_abs_old = None
300 self.error_norm_old = None
302 self.newton_tol = max(10 * EPS / rtol, min(0.03, rtol ** 0.5))
303 self.sol = None
305 self.jac_factor = None
306 self.jac, self.J = self._validate_jac(jac, jac_sparsity)
307 if issparse(self.J):
308 def lu(A):
309 self.nlu += 1
310 return splu(A)
312 def solve_lu(LU, b):
313 return LU.solve(b)
315 I = eye(self.n, format='csc')
316 else:
317 def lu(A):
318 self.nlu += 1
319 return lu_factor(A, overwrite_a=True)
321 def solve_lu(LU, b):
322 return lu_solve(LU, b, overwrite_b=True)
324 I = np.identity(self.n)
326 self.lu = lu
327 self.solve_lu = solve_lu
328 self.I = I
330 self.current_jac = True
331 self.LU_real = None
332 self.LU_complex = None
333 self.Z = None
335 def _validate_jac(self, jac, sparsity):
336 t0 = self.t
337 y0 = self.y
339 if jac is None:
340 if sparsity is not None:
341 if issparse(sparsity):
342 sparsity = csc_matrix(sparsity)
343 groups = group_columns(sparsity)
344 sparsity = (sparsity, groups)
346 def jac_wrapped(t, y, f):
347 self.njev += 1
348 J, self.jac_factor = num_jac(self.fun_vectorized, t, y, f,
349 self.atol, self.jac_factor,
350 sparsity)
351 return J
352 J = jac_wrapped(t0, y0, self.f)
353 elif callable(jac):
354 J = jac(t0, y0)
355 self.njev = 1
356 if issparse(J):
357 J = csc_matrix(J)
359 def jac_wrapped(t, y, _=None):
360 self.njev += 1
361 return csc_matrix(jac(t, y), dtype=float)
363 else:
364 J = np.asarray(J, dtype=float)
366 def jac_wrapped(t, y, _=None):
367 self.njev += 1
368 return np.asarray(jac(t, y), dtype=float)
370 if J.shape != (self.n, self.n):
371 raise ValueError("`jac` is expected to have shape {}, but "
372 "actually has {}."
373 .format((self.n, self.n), J.shape))
374 else:
375 if issparse(jac):
376 J = csc_matrix(jac)
377 else:
378 J = np.asarray(jac, dtype=float)
380 if J.shape != (self.n, self.n):
381 raise ValueError("`jac` is expected to have shape {}, but "
382 "actually has {}."
383 .format((self.n, self.n), J.shape))
384 jac_wrapped = None
386 return jac_wrapped, J
388 def _step_impl(self):
389 t = self.t
390 y = self.y
391 f = self.f
393 max_step = self.max_step
394 atol = self.atol
395 rtol = self.rtol
397 min_step = 10 * np.abs(np.nextafter(t, self.direction * np.inf) - t)
398 if self.h_abs > max_step:
399 h_abs = max_step
400 h_abs_old = None
401 error_norm_old = None
402 elif self.h_abs < min_step:
403 h_abs = min_step
404 h_abs_old = None
405 error_norm_old = None
406 else:
407 h_abs = self.h_abs
408 h_abs_old = self.h_abs_old
409 error_norm_old = self.error_norm_old
411 J = self.J
412 LU_real = self.LU_real
413 LU_complex = self.LU_complex
415 current_jac = self.current_jac
416 jac = self.jac
418 rejected = False
419 step_accepted = False
420 message = None
421 while not step_accepted:
422 if h_abs < min_step:
423 return False, self.TOO_SMALL_STEP
425 h = h_abs * self.direction
426 t_new = t + h
428 if self.direction * (t_new - self.t_bound) > 0:
429 t_new = self.t_bound
431 h = t_new - t
432 h_abs = np.abs(h)
434 if self.sol is None:
435 Z0 = np.zeros((3, y.shape[0]))
436 else:
437 Z0 = self.sol(t + h * C).T - y
439 scale = atol + np.abs(y) * rtol
441 converged = False
442 while not converged:
443 if LU_real is None or LU_complex is None:
444 LU_real = self.lu(MU_REAL / h * self.I - J)
445 LU_complex = self.lu(MU_COMPLEX / h * self.I - J)
447 converged, n_iter, Z, rate = solve_collocation_system(
448 self.fun, t, y, h, Z0, scale, self.newton_tol,
449 LU_real, LU_complex, self.solve_lu)
451 if not converged:
452 if current_jac:
453 break
455 J = self.jac(t, y, f)
456 current_jac = True
457 LU_real = None
458 LU_complex = None
460 if not converged:
461 h_abs *= 0.5
462 LU_real = None
463 LU_complex = None
464 continue
466 y_new = y + Z[-1]
467 ZE = Z.T.dot(E) / h
468 error = self.solve_lu(LU_real, f + ZE)
469 scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol
470 error_norm = norm(error / scale)
471 safety = 0.9 * (2 * NEWTON_MAXITER + 1) / (2 * NEWTON_MAXITER
472 + n_iter)
474 if rejected and error_norm > 1:
475 error = self.solve_lu(LU_real, self.fun(t, y + error) + ZE)
476 error_norm = norm(error / scale)
478 if error_norm > 1:
479 factor = predict_factor(h_abs, h_abs_old,
480 error_norm, error_norm_old)
481 h_abs *= max(MIN_FACTOR, safety * factor)
483 LU_real = None
484 LU_complex = None
485 rejected = True
486 else:
487 step_accepted = True
489 recompute_jac = jac is not None and n_iter > 2 and rate > 1e-3
491 factor = predict_factor(h_abs, h_abs_old, error_norm, error_norm_old)
492 factor = min(MAX_FACTOR, safety * factor)
494 if not recompute_jac and factor < 1.2:
495 factor = 1
496 else:
497 LU_real = None
498 LU_complex = None
500 f_new = self.fun(t_new, y_new)
501 if recompute_jac:
502 J = jac(t_new, y_new, f_new)
503 current_jac = True
504 elif jac is not None:
505 current_jac = False
507 self.h_abs_old = self.h_abs
508 self.error_norm_old = error_norm
510 self.h_abs = h_abs * factor
512 self.y_old = y
514 self.t = t_new
515 self.y = y_new
516 self.f = f_new
518 self.Z = Z
520 self.LU_real = LU_real
521 self.LU_complex = LU_complex
522 self.current_jac = current_jac
523 self.J = J
525 self.t_old = t
526 self.sol = self._compute_dense_output()
528 return step_accepted, message
530 def _compute_dense_output(self):
531 Q = np.dot(self.Z.T, P)
532 return RadauDenseOutput(self.t_old, self.t, self.y_old, Q)
534 def _dense_output_impl(self):
535 return self.sol
538class RadauDenseOutput(DenseOutput):
539 def __init__(self, t_old, t, y_old, Q):
540 super(RadauDenseOutput, self).__init__(t_old, t)
541 self.h = t - t_old
542 self.Q = Q
543 self.order = Q.shape[1] - 1
544 self.y_old = y_old
546 def _call_impl(self, t):
547 x = (t - self.t_old) / self.h
548 if t.ndim == 0:
549 p = np.tile(x, self.order + 1)
550 p = np.cumprod(p)
551 else:
552 p = np.tile(x, (self.order + 1, 1))
553 p = np.cumprod(p, axis=0)
554 # Here we don't multiply by h, not a mistake.
555 y = np.dot(self.Q, p)
556 if y.ndim == 2:
557 y += self.y_old[:, None]
558 else:
559 y += self.y_old
561 return y