File size: 19,676 Bytes
7885a28 |
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 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 |
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,
validate_first_step)
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 : ndarray, shape (n)
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
rate = None
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
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 time derivative of the state ``y``
at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a
scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must
return an array of the same shape as ``y``. See `vectorized` for more
information.
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.
first_step : float or None, optional
Initial step size. Default is ``None`` which means that the algorithm
should choose.
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)``. HHere `rtol` controls a
relative accuracy (number of correct digits), while `atol` controls
absolute accuracy (number of correct decimal places). To achieve the
desired `rtol`, set `atol` to be smaller than the smallest value that
can be expected from ``rtol * abs(y)`` so that `rtol` dominates the
allowable error. If `atol` is larger than ``rtol * abs(y)`` the
number of correct digits is not guaranteed. Conversely, to achieve the
desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller
than `atol`. 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` can be called in a vectorized fashion. Default is False.
If ``vectorized`` is False, `fun` will always be called with ``y`` of
shape ``(n,)``, where ``n = len(y0)``.
If ``vectorized`` is True, `fun` may be called with ``y`` of shape
``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave
such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of
the returned array is the time derivative of the state corresponding
with a column of ``y``).
Setting ``vectorized=True`` allows for faster finite difference
approximation of the Jacobian by this method, but may result in slower
execution overall in some circumstances (e.g. small ``len(y0)``).
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, first_step=None, **extraneous):
warn_extraneous(extraneous)
super().__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.
if first_step is None:
self.h_abs = select_initial_step(
self.fun, self.t, self.y, t_bound, max_step, self.f, self.direction,
3, self.rtol, self.atol)
else:
self.h_abs = validate_first_step(first_step, t0, t_bound)
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(f"`jac` is expected to have shape {(self.n, self.n)},"
f" but actually has {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(f"`jac` is expected to have shape {(self.n, self.n)},"
f" but actually has {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().__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
|