Hide keyboard shortcuts

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 

10 

11S6 = 6 ** 0.5 

12 

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 

16 

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))) 

22 

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] 

35 

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]]) 

41 

42 

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. 

46 

47 

48def solve_collocation_system(fun, t, y, h, Z0, scale, tol, 

49 LU_real, LU_complex, solve_lu): 

50 """Solve the collocation system. 

51 

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)``. 

75 

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 

90 

91 W = TI.dot(Z0) 

92 Z = Z0 

93 

94 F = np.empty((3, n)) 

95 ch = h * C 

96 

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]) 

104 

105 if not np.all(np.isfinite(F)): 

106 break 

107 

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]) 

110 

111 dW_real = solve_lu(LU_real, f_real) 

112 dW_complex = solve_lu(LU_complex, f_complex) 

113 

114 dW[0] = dW_real 

115 dW[1] = dW_complex.real 

116 dW[2] = dW_complex.imag 

117 

118 dW_norm = norm(dW / scale) 

119 if dW_norm_old is not None: 

120 rate = dW_norm / dW_norm_old 

121 

122 if (rate is not None and (rate >= 1 or 

123 rate ** (NEWTON_MAXITER - k) / (1 - rate) * dW_norm > tol)): 

124 break 

125 

126 W += dW 

127 Z = T.dot(W) 

128 

129 if (dW_norm == 0 or 

130 rate is not None and rate / (1 - rate) * dW_norm < tol): 

131 converged = True 

132 break 

133 

134 dW_norm_old = dW_norm 

135 

136 return converged, k + 1, Z, rate 

137 

138 

139def predict_factor(h_abs, h_abs_old, error_norm, error_norm_old): 

140 """Predict by which factor to increase/decrease the step size. 

141 

142 The algorithm is described in [1]_. 

143 

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). 

152 

153 Returns 

154 ------- 

155 factor : float 

156 Predicted factor. 

157 

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. 

162 

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 

172 

173 with np.errstate(divide='ignore'): 

174 factor = min(1, multiplier) * error_norm ** -0.25 

175 

176 return factor 

177 

178 

179class Radau(OdeSolver): 

180 """Implicit Runge-Kutta method of Radau IIA family of order 5. 

181 

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. 

185 

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: 

226 

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. 

235 

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. 

248 

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. 

273 

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 

301 

302 self.newton_tol = max(10 * EPS / rtol, min(0.03, rtol ** 0.5)) 

303 self.sol = None 

304 

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) 

311 

312 def solve_lu(LU, b): 

313 return LU.solve(b) 

314 

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) 

320 

321 def solve_lu(LU, b): 

322 return lu_solve(LU, b, overwrite_b=True) 

323 

324 I = np.identity(self.n) 

325 

326 self.lu = lu 

327 self.solve_lu = solve_lu 

328 self.I = I 

329 

330 self.current_jac = True 

331 self.LU_real = None 

332 self.LU_complex = None 

333 self.Z = None 

334 

335 def _validate_jac(self, jac, sparsity): 

336 t0 = self.t 

337 y0 = self.y 

338 

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) 

345 

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) 

358 

359 def jac_wrapped(t, y, _=None): 

360 self.njev += 1 

361 return csc_matrix(jac(t, y), dtype=float) 

362 

363 else: 

364 J = np.asarray(J, dtype=float) 

365 

366 def jac_wrapped(t, y, _=None): 

367 self.njev += 1 

368 return np.asarray(jac(t, y), dtype=float) 

369 

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) 

379 

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 

385 

386 return jac_wrapped, J 

387 

388 def _step_impl(self): 

389 t = self.t 

390 y = self.y 

391 f = self.f 

392 

393 max_step = self.max_step 

394 atol = self.atol 

395 rtol = self.rtol 

396 

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 

410 

411 J = self.J 

412 LU_real = self.LU_real 

413 LU_complex = self.LU_complex 

414 

415 current_jac = self.current_jac 

416 jac = self.jac 

417 

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 

424 

425 h = h_abs * self.direction 

426 t_new = t + h 

427 

428 if self.direction * (t_new - self.t_bound) > 0: 

429 t_new = self.t_bound 

430 

431 h = t_new - t 

432 h_abs = np.abs(h) 

433 

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 

438 

439 scale = atol + np.abs(y) * rtol 

440 

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) 

446 

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) 

450 

451 if not converged: 

452 if current_jac: 

453 break 

454 

455 J = self.jac(t, y, f) 

456 current_jac = True 

457 LU_real = None 

458 LU_complex = None 

459 

460 if not converged: 

461 h_abs *= 0.5 

462 LU_real = None 

463 LU_complex = None 

464 continue 

465 

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) 

473 

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) 

477 

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) 

482 

483 LU_real = None 

484 LU_complex = None 

485 rejected = True 

486 else: 

487 step_accepted = True 

488 

489 recompute_jac = jac is not None and n_iter > 2 and rate > 1e-3 

490 

491 factor = predict_factor(h_abs, h_abs_old, error_norm, error_norm_old) 

492 factor = min(MAX_FACTOR, safety * factor) 

493 

494 if not recompute_jac and factor < 1.2: 

495 factor = 1 

496 else: 

497 LU_real = None 

498 LU_complex = None 

499 

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 

506 

507 self.h_abs_old = self.h_abs 

508 self.error_norm_old = error_norm 

509 

510 self.h_abs = h_abs * factor 

511 

512 self.y_old = y 

513 

514 self.t = t_new 

515 self.y = y_new 

516 self.f = f_new 

517 

518 self.Z = Z 

519 

520 self.LU_real = LU_real 

521 self.LU_complex = LU_complex 

522 self.current_jac = current_jac 

523 self.J = J 

524 

525 self.t_old = t 

526 self.sol = self._compute_dense_output() 

527 

528 return step_accepted, message 

529 

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) 

533 

534 def _dense_output_impl(self): 

535 return self.sol 

536 

537 

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 

545 

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 

560 

561 return y