1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

from __future__ import division, print_function, absolute_import 

import numpy as np 

from scipy.linalg import lu_factor, lu_solve 

from scipy.sparse import csc_matrix, issparse, eye 

from scipy.sparse.linalg import splu 

from scipy.optimize._numdiff import group_columns 

from .common import (validate_max_step, validate_tol, select_initial_step, 

norm, num_jac, EPS, warn_extraneous) 

from .base import OdeSolver, DenseOutput 

 

S6 = 6 ** 0.5 

 

# Butcher tableau. A is not used directly, see below. 

C = np.array([(4 - S6) / 10, (4 + S6) / 10, 1]) 

E = np.array([-13 - 7 * S6, -13 + 7 * S6, -1]) / 3 

 

# Eigendecomposition of A is done: A = T L T**-1. There is 1 real eigenvalue 

# and a complex conjugate pair. They are written below. 

MU_REAL = 3 + 3 ** (2 / 3) - 3 ** (1 / 3) 

MU_COMPLEX = (3 + 0.5 * (3 ** (1 / 3) - 3 ** (2 / 3)) 

- 0.5j * (3 ** (5 / 6) + 3 ** (7 / 6))) 

 

# These are transformation matrices. 

T = np.array([ 

[0.09443876248897524, -0.14125529502095421, 0.03002919410514742], 

[0.25021312296533332, 0.20412935229379994, -0.38294211275726192], 

[1, 1, 0]]) 

TI = np.array([ 

[4.17871859155190428, 0.32768282076106237, 0.52337644549944951], 

[-4.17871859155190428, -0.32768282076106237, 0.47662355450055044], 

[0.50287263494578682, -2.57192694985560522, 0.59603920482822492]]) 

# These linear combinations are used in the algorithm. 

TI_REAL = TI[0] 

TI_COMPLEX = TI[1] + 1j * TI[2] 

 

# Interpolator coefficients. 

P = np.array([ 

[13/3 + 7*S6/3, -23/3 - 22*S6/3, 10/3 + 5 * S6], 

[13/3 - 7*S6/3, -23/3 + 22*S6/3, 10/3 - 5 * S6], 

[1/3, -8/3, 10/3]]) 

 

 

NEWTON_MAXITER = 6 # Maximum number of Newton iterations. 

MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size. 

MAX_FACTOR = 10 # Maximum allowed increase in a step size. 

 

 

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

LU_real, LU_complex, solve_lu): 

"""Solve the collocation system. 

 

Parameters 

---------- 

fun : callable 

Right-hand side of the system. 

t : float 

Current time. 

y : ndarray, shape (n,) 

Current state. 

h : float 

Step to try. 

Z0 : ndarray, shape (3, n) 

Initial guess for the solution. It determines new values of `y` at 

``t + h * C`` as ``y + Z0``, where ``C`` is the Radau method constants. 

scale : float 

Problem tolerance scale, i.e. ``rtol * abs(y) + atol``. 

tol : float 

Tolerance to which solve the system. This value is compared with 

the normalized by `scale` error. 

LU_real, LU_complex 

LU decompositions of the system Jacobians. 

solve_lu : callable 

Callable which solves a linear system given a LU decomposition. The 

signature is ``solve_lu(LU, b)``. 

 

Returns 

------- 

converged : bool 

Whether iterations converged. 

n_iter : int 

Number of completed iterations. 

Z : ndarray, shape (3, n) 

Found solution. 

rate : float 

The rate of convergence. 

""" 

n = y.shape[0] 

M_real = MU_REAL / h 

M_complex = MU_COMPLEX / h 

 

W = TI.dot(Z0) 

Z = Z0 

 

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

ch = h * C 

 

dW_norm_old = None 

dW = np.empty_like(W) 

converged = False 

for k in range(NEWTON_MAXITER): 

for i in range(3): 

F[i] = fun(t + ch[i], y + Z[i]) 

 

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

break 

 

f_real = F.T.dot(TI_REAL) - M_real * W[0] 

f_complex = F.T.dot(TI_COMPLEX) - M_complex * (W[1] + 1j * W[2]) 

 

dW_real = solve_lu(LU_real, f_real) 

dW_complex = solve_lu(LU_complex, f_complex) 

 

dW[0] = dW_real 

dW[1] = dW_complex.real 

dW[2] = dW_complex.imag 

 

dW_norm = norm(dW / scale) 

if dW_norm_old is not None: 

rate = dW_norm / dW_norm_old 

else: 

rate = None 

 

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

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

break 

 

W += dW 

Z = T.dot(W) 

 

if (dW_norm == 0 or 

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

converged = True 

break 

 

dW_norm_old = dW_norm 

 

return converged, k + 1, Z, rate 

 

 

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

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

 

The algorithm is described in [1]_. 

 

Parameters 

---------- 

h_abs, h_abs_old : float 

Current and previous values of the step size, `h_abs_old` can be None 

(see Notes). 

error_norm, error_norm_old : float 

Current and previous values of the error norm, `error_norm_old` can 

be None (see Notes). 

 

Returns 

------- 

factor : float 

Predicted factor. 

 

Notes 

----- 

If `h_abs_old` and `error_norm_old` are both not None then a two-step 

algorithm is used, otherwise a one-step algorithm is used. 

 

References 

---------- 

.. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential 

Equations II: Stiff and Differential-Algebraic Problems", Sec. IV.8. 

""" 

if error_norm_old is None or h_abs_old is None or error_norm == 0: 

multiplier = 1 

else: 

multiplier = h_abs / h_abs_old * (error_norm_old / error_norm) ** 0.25 

 

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

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

 

return factor 

 

 

class Radau(OdeSolver): 

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

 

The implementation follows [1]_. The error is controlled with a 

third-order accurate embedded formula. A cubic polynomial which satisfies 

the collocation conditions is used for the dense output. 

 

Parameters 

---------- 

fun : callable 

Right-hand side of the system. The calling signature is ``fun(t, y)``. 

Here ``t`` is a scalar, and there are two options for the ndarray ``y``: 

It can either have shape (n,); then ``fun`` must return array_like with 

shape (n,). Alternatively it can have shape (n, k); then ``fun`` 

must return an array_like with shape (n, k), i.e. each column 

corresponds to a single column in ``y``. The choice between the two 

options is determined by `vectorized` argument (see below). The 

vectorized implementation allows a faster approximation of the Jacobian 

by finite differences (required for this solver). 

t0 : float 

Initial time. 

y0 : array_like, shape (n,) 

Initial state. 

t_bound : float 

Boundary time - the integration won't continue beyond it. It also 

determines the direction of the integration. 

max_step : float, optional 

Maximum allowed step size. Default is np.inf, i.e. the step size is not 

bounded and determined solely by the solver. 

rtol, atol : float and array_like, optional 

Relative and absolute tolerances. The solver keeps the local error 

estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a 

relative accuracy (number of correct digits). But if a component of `y` 

is approximately below `atol`, the error only needs to fall within 

the same `atol` threshold, and the number of correct digits is not 

guaranteed. If components of y have different scales, it might be 

beneficial to set different `atol` values for different components by 

passing array_like with shape (n,) for `atol`. Default values are 

1e-3 for `rtol` and 1e-6 for `atol`. 

jac : {None, array_like, sparse_matrix, callable}, optional 

Jacobian matrix of the right-hand side of the system with respect to 

y, required by this method. The Jacobian matrix has shape (n, n) and 

its element (i, j) is equal to ``d f_i / d y_j``. 

There are three ways to define the Jacobian: 

 

* If array_like or sparse_matrix, the Jacobian is assumed to 

be constant. 

* If callable, the Jacobian is assumed to depend on both 

t and y; it will be called as ``jac(t, y)`` as necessary. 

For the 'Radau' and 'BDF' methods, the return value might be a 

sparse matrix. 

* If None (default), the Jacobian will be approximated by 

finite differences. 

 

It is generally recommended to provide the Jacobian rather than 

relying on a finite-difference approximation. 

jac_sparsity : {None, array_like, sparse matrix}, optional 

Defines a sparsity structure of the Jacobian matrix for a 

finite-difference approximation. Its shape must be (n, n). This argument 

is ignored if `jac` is not `None`. If the Jacobian has only few non-zero 

elements in *each* row, providing the sparsity structure will greatly 

speed up the computations [2]_. A zero entry means that a corresponding 

element in the Jacobian is always zero. If None (default), the Jacobian 

is assumed to be dense. 

vectorized : bool, optional 

Whether `fun` is implemented in a vectorized fashion. Default is False. 

 

Attributes 

---------- 

n : int 

Number of equations. 

status : string 

Current status of the solver: 'running', 'finished' or 'failed'. 

t_bound : float 

Boundary time. 

direction : float 

Integration direction: +1 or -1. 

t : float 

Current time. 

y : ndarray 

Current state. 

t_old : float 

Previous time. None if no steps were made yet. 

step_size : float 

Size of the last successful step. None if no steps were made yet. 

nfev : int 

Number of evaluations of the right-hand side. 

njev : int 

Number of evaluations of the Jacobian. 

nlu : int 

Number of LU decompositions. 

 

References 

---------- 

.. [1] E. Hairer, G. Wanner, "Solving Ordinary Differential Equations II: 

Stiff and Differential-Algebraic Problems", Sec. IV.8. 

.. [2] A. Curtis, M. J. D. Powell, and J. Reid, "On the estimation of 

sparse Jacobian matrices", Journal of the Institute of Mathematics 

and its Applications, 13, pp. 117-120, 1974. 

""" 

def __init__(self, fun, t0, y0, t_bound, max_step=np.inf, 

rtol=1e-3, atol=1e-6, jac=None, jac_sparsity=None, 

vectorized=False, **extraneous): 

warn_extraneous(extraneous) 

super(Radau, self).__init__(fun, t0, y0, t_bound, vectorized) 

self.y_old = None 

self.max_step = validate_max_step(max_step) 

self.rtol, self.atol = validate_tol(rtol, atol, self.n) 

self.f = self.fun(self.t, self.y) 

# Select initial step assuming the same order which is used to control 

# the error. 

self.h_abs = select_initial_step( 

self.fun, self.t, self.y, self.f, self.direction, 

3, self.rtol, self.atol) 

self.h_abs_old = None 

self.error_norm_old = None 

 

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

self.sol = None 

 

self.jac_factor = None 

self.jac, self.J = self._validate_jac(jac, jac_sparsity) 

if issparse(self.J): 

def lu(A): 

self.nlu += 1 

return splu(A) 

 

def solve_lu(LU, b): 

return LU.solve(b) 

 

I = eye(self.n, format='csc') 

else: 

def lu(A): 

self.nlu += 1 

return lu_factor(A, overwrite_a=True) 

 

def solve_lu(LU, b): 

return lu_solve(LU, b, overwrite_b=True) 

 

I = np.identity(self.n) 

 

self.lu = lu 

self.solve_lu = solve_lu 

self.I = I 

 

self.current_jac = True 

self.LU_real = None 

self.LU_complex = None 

self.Z = None 

 

def _validate_jac(self, jac, sparsity): 

t0 = self.t 

y0 = self.y 

 

if jac is None: 

if sparsity is not None: 

if issparse(sparsity): 

sparsity = csc_matrix(sparsity) 

groups = group_columns(sparsity) 

sparsity = (sparsity, groups) 

 

def jac_wrapped(t, y, f): 

self.njev += 1 

J, self.jac_factor = num_jac(self.fun_vectorized, t, y, f, 

self.atol, self.jac_factor, 

sparsity) 

return J 

J = jac_wrapped(t0, y0, self.f) 

elif callable(jac): 

J = jac(t0, y0) 

self.njev = 1 

if issparse(J): 

J = csc_matrix(J) 

 

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

self.njev += 1 

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

 

else: 

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

 

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

self.njev += 1 

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

 

if J.shape != (self.n, self.n): 

raise ValueError("`jac` is expected to have shape {}, but " 

"actually has {}." 

.format((self.n, self.n), J.shape)) 

else: 

if issparse(jac): 

J = csc_matrix(jac) 

else: 

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

 

if J.shape != (self.n, self.n): 

raise ValueError("`jac` is expected to have shape {}, but " 

"actually has {}." 

.format((self.n, self.n), J.shape)) 

jac_wrapped = None 

 

return jac_wrapped, J 

 

def _step_impl(self): 

t = self.t 

y = self.y 

f = self.f 

 

max_step = self.max_step 

atol = self.atol 

rtol = self.rtol 

 

min_step = 10 * np.abs(np.nextafter(t, self.direction * np.inf) - t) 

if self.h_abs > max_step: 

h_abs = max_step 

h_abs_old = None 

error_norm_old = None 

elif self.h_abs < min_step: 

h_abs = min_step 

h_abs_old = None 

error_norm_old = None 

else: 

h_abs = self.h_abs 

h_abs_old = self.h_abs_old 

error_norm_old = self.error_norm_old 

 

J = self.J 

LU_real = self.LU_real 

LU_complex = self.LU_complex 

 

current_jac = self.current_jac 

jac = self.jac 

 

rejected = False 

step_accepted = False 

message = None 

while not step_accepted: 

if h_abs < min_step: 

return False, self.TOO_SMALL_STEP 

 

h = h_abs * self.direction 

t_new = t + h 

 

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

t_new = self.t_bound 

 

h = t_new - t 

h_abs = np.abs(h) 

 

if self.sol is None: 

Z0 = np.zeros((3, y.shape[0])) 

else: 

Z0 = self.sol(t + h * C).T - y 

 

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

 

converged = False 

while not converged: 

if LU_real is None or LU_complex is None: 

LU_real = self.lu(MU_REAL / h * self.I - J) 

LU_complex = self.lu(MU_COMPLEX / h * self.I - J) 

 

converged, n_iter, Z, rate = solve_collocation_system( 

self.fun, t, y, h, Z0, scale, self.newton_tol, 

LU_real, LU_complex, self.solve_lu) 

 

if not converged: 

if current_jac: 

break 

 

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

current_jac = True 

LU_real = None 

LU_complex = None 

 

if not converged: 

h_abs *= 0.5 

LU_real = None 

LU_complex = None 

continue 

 

y_new = y + Z[-1] 

ZE = Z.T.dot(E) / h 

error = self.solve_lu(LU_real, f + ZE) 

scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol 

error_norm = norm(error / scale) 

safety = 0.9 * (2 * NEWTON_MAXITER + 1) / (2 * NEWTON_MAXITER 

+ n_iter) 

 

if rejected and error_norm > 1: 

error = self.solve_lu(LU_real, self.fun(t, y + error) + ZE) 

error_norm = norm(error / scale) 

 

if error_norm > 1: 

factor = predict_factor(h_abs, h_abs_old, 

error_norm, error_norm_old) 

h_abs *= max(MIN_FACTOR, safety * factor) 

 

LU_real = None 

LU_complex = None 

rejected = True 

else: 

step_accepted = True 

 

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

 

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

factor = min(MAX_FACTOR, safety * factor) 

 

if not recompute_jac and factor < 1.2: 

factor = 1 

else: 

LU_real = None 

LU_complex = None 

 

f_new = self.fun(t_new, y_new) 

if recompute_jac: 

J = jac(t_new, y_new, f_new) 

current_jac = True 

elif jac is not None: 

current_jac = False 

 

self.h_abs_old = self.h_abs 

self.error_norm_old = error_norm 

 

self.h_abs = h_abs * factor 

 

self.y_old = y 

 

self.t = t_new 

self.y = y_new 

self.f = f_new 

 

self.Z = Z 

 

self.LU_real = LU_real 

self.LU_complex = LU_complex 

self.current_jac = current_jac 

self.J = J 

 

self.t_old = t 

self.sol = self._compute_dense_output() 

 

return step_accepted, message 

 

def _compute_dense_output(self): 

Q = np.dot(self.Z.T, P) 

return RadauDenseOutput(self.t_old, self.t, self.y_old, Q) 

 

def _dense_output_impl(self): 

return self.sol 

 

 

class RadauDenseOutput(DenseOutput): 

def __init__(self, t_old, t, y_old, Q): 

super(RadauDenseOutput, self).__init__(t_old, t) 

self.h = t - t_old 

self.Q = Q 

self.order = Q.shape[1] - 1 

self.y_old = y_old 

 

def _call_impl(self, t): 

x = (t - self.t_old) / self.h 

if t.ndim == 0: 

p = np.tile(x, self.order + 1) 

p = np.cumprod(p) 

else: 

p = np.tile(x, (self.order + 1, 1)) 

p = np.cumprod(p, axis=0) 

# Here we don't multiply by h, not a mistake. 

y = np.dot(self.Q, p) 

if y.ndim == 2: 

y += self.y_old[:, None] 

else: 

y += self.y_old 

 

return y