573 lines
19 KiB
Python
573 lines
19 KiB
Python
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
|