You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
219 lines
6.5 KiB
Python
219 lines
6.5 KiB
Python
import itertools
|
|
import numpy as np
|
|
from numpy.testing import assert_allclose
|
|
from scipy.integrate import ode
|
|
|
|
|
|
def _band_count(a):
|
|
"""Returns ml and mu, the lower and upper band sizes of a."""
|
|
nrows, ncols = a.shape
|
|
ml = 0
|
|
for k in range(-nrows+1, 0):
|
|
if np.diag(a, k).any():
|
|
ml = -k
|
|
break
|
|
mu = 0
|
|
for k in range(nrows-1, 0, -1):
|
|
if np.diag(a, k).any():
|
|
mu = k
|
|
break
|
|
return ml, mu
|
|
|
|
|
|
def _linear_func(t, y, a):
|
|
"""Linear system dy/dt = a * y"""
|
|
return a.dot(y)
|
|
|
|
|
|
def _linear_jac(t, y, a):
|
|
"""Jacobian of a * y is a."""
|
|
return a
|
|
|
|
|
|
def _linear_banded_jac(t, y, a):
|
|
"""Banded Jacobian."""
|
|
ml, mu = _band_count(a)
|
|
bjac = [np.r_[[0] * k, np.diag(a, k)] for k in range(mu, 0, -1)]
|
|
bjac.append(np.diag(a))
|
|
for k in range(-1, -ml-1, -1):
|
|
bjac.append(np.r_[np.diag(a, k), [0] * (-k)])
|
|
return bjac
|
|
|
|
|
|
def _solve_linear_sys(a, y0, tend=1, dt=0.1,
|
|
solver=None, method='bdf', use_jac=True,
|
|
with_jacobian=False, banded=False):
|
|
"""Use scipy.integrate.ode to solve a linear system of ODEs.
|
|
|
|
a : square ndarray
|
|
Matrix of the linear system to be solved.
|
|
y0 : ndarray
|
|
Initial condition
|
|
tend : float
|
|
Stop time.
|
|
dt : float
|
|
Step size of the output.
|
|
solver : str
|
|
If not None, this must be "vode", "lsoda" or "zvode".
|
|
method : str
|
|
Either "bdf" or "adams".
|
|
use_jac : bool
|
|
Determines if the jacobian function is passed to ode().
|
|
with_jacobian : bool
|
|
Passed to ode.set_integrator().
|
|
banded : bool
|
|
Determines whether a banded or full jacobian is used.
|
|
If `banded` is True, `lband` and `uband` are determined by the
|
|
values in `a`.
|
|
"""
|
|
if banded:
|
|
lband, uband = _band_count(a)
|
|
else:
|
|
lband = None
|
|
uband = None
|
|
|
|
if use_jac:
|
|
if banded:
|
|
r = ode(_linear_func, _linear_banded_jac)
|
|
else:
|
|
r = ode(_linear_func, _linear_jac)
|
|
else:
|
|
r = ode(_linear_func)
|
|
|
|
if solver is None:
|
|
if np.iscomplexobj(a):
|
|
solver = "zvode"
|
|
else:
|
|
solver = "vode"
|
|
|
|
r.set_integrator(solver,
|
|
with_jacobian=with_jacobian,
|
|
method=method,
|
|
lband=lband, uband=uband,
|
|
rtol=1e-9, atol=1e-10,
|
|
)
|
|
t0 = 0
|
|
r.set_initial_value(y0, t0)
|
|
r.set_f_params(a)
|
|
r.set_jac_params(a)
|
|
|
|
t = [t0]
|
|
y = [y0]
|
|
while r.successful() and r.t < tend:
|
|
r.integrate(r.t + dt)
|
|
t.append(r.t)
|
|
y.append(r.y)
|
|
|
|
t = np.array(t)
|
|
y = np.array(y)
|
|
return t, y
|
|
|
|
|
|
def _analytical_solution(a, y0, t):
|
|
"""
|
|
Analytical solution to the linear differential equations dy/dt = a*y.
|
|
|
|
The solution is only valid if `a` is diagonalizable.
|
|
|
|
Returns a 2-D array with shape (len(t), len(y0)).
|
|
"""
|
|
lam, v = np.linalg.eig(a)
|
|
c = np.linalg.solve(v, y0)
|
|
e = c * np.exp(lam * t.reshape(-1, 1))
|
|
sol = e.dot(v.T)
|
|
return sol
|
|
|
|
|
|
def test_banded_ode_solvers():
|
|
# Test the "lsoda", "vode" and "zvode" solvers of the `ode` class
|
|
# with a system that has a banded Jacobian matrix.
|
|
|
|
t_exact = np.linspace(0, 1.0, 5)
|
|
|
|
# --- Real arrays for testing the "lsoda" and "vode" solvers ---
|
|
|
|
# lband = 2, uband = 1:
|
|
a_real = np.array([[-0.6, 0.1, 0.0, 0.0, 0.0],
|
|
[0.2, -0.5, 0.9, 0.0, 0.0],
|
|
[0.1, 0.1, -0.4, 0.1, 0.0],
|
|
[0.0, 0.3, -0.1, -0.9, -0.3],
|
|
[0.0, 0.0, 0.1, 0.1, -0.7]])
|
|
|
|
# lband = 0, uband = 1:
|
|
a_real_upper = np.triu(a_real)
|
|
|
|
# lband = 2, uband = 0:
|
|
a_real_lower = np.tril(a_real)
|
|
|
|
# lband = 0, uband = 0:
|
|
a_real_diag = np.triu(a_real_lower)
|
|
|
|
real_matrices = [a_real, a_real_upper, a_real_lower, a_real_diag]
|
|
real_solutions = []
|
|
|
|
for a in real_matrices:
|
|
y0 = np.arange(1, a.shape[0] + 1)
|
|
y_exact = _analytical_solution(a, y0, t_exact)
|
|
real_solutions.append((y0, t_exact, y_exact))
|
|
|
|
def check_real(idx, solver, meth, use_jac, with_jac, banded):
|
|
a = real_matrices[idx]
|
|
y0, t_exact, y_exact = real_solutions[idx]
|
|
t, y = _solve_linear_sys(a, y0,
|
|
tend=t_exact[-1],
|
|
dt=t_exact[1] - t_exact[0],
|
|
solver=solver,
|
|
method=meth,
|
|
use_jac=use_jac,
|
|
with_jacobian=with_jac,
|
|
banded=banded)
|
|
assert_allclose(t, t_exact)
|
|
assert_allclose(y, y_exact)
|
|
|
|
for idx in range(len(real_matrices)):
|
|
p = [['vode', 'lsoda'], # solver
|
|
['bdf', 'adams'], # method
|
|
[False, True], # use_jac
|
|
[False, True], # with_jacobian
|
|
[False, True]] # banded
|
|
for solver, meth, use_jac, with_jac, banded in itertools.product(*p):
|
|
check_real(idx, solver, meth, use_jac, with_jac, banded)
|
|
|
|
# --- Complex arrays for testing the "zvode" solver ---
|
|
|
|
# complex, lband = 2, uband = 1:
|
|
a_complex = a_real - 0.5j * a_real
|
|
|
|
# complex, lband = 0, uband = 0:
|
|
a_complex_diag = np.diag(np.diag(a_complex))
|
|
|
|
complex_matrices = [a_complex, a_complex_diag]
|
|
complex_solutions = []
|
|
|
|
for a in complex_matrices:
|
|
y0 = np.arange(1, a.shape[0] + 1) + 1j
|
|
y_exact = _analytical_solution(a, y0, t_exact)
|
|
complex_solutions.append((y0, t_exact, y_exact))
|
|
|
|
def check_complex(idx, solver, meth, use_jac, with_jac, banded):
|
|
a = complex_matrices[idx]
|
|
y0, t_exact, y_exact = complex_solutions[idx]
|
|
t, y = _solve_linear_sys(a, y0,
|
|
tend=t_exact[-1],
|
|
dt=t_exact[1] - t_exact[0],
|
|
solver=solver,
|
|
method=meth,
|
|
use_jac=use_jac,
|
|
with_jacobian=with_jac,
|
|
banded=banded)
|
|
assert_allclose(t, t_exact)
|
|
assert_allclose(y, y_exact)
|
|
|
|
for idx in range(len(complex_matrices)):
|
|
p = [['bdf', 'adams'], # method
|
|
[False, True], # use_jac
|
|
[False, True], # with_jacobian
|
|
[False, True]] # banded
|
|
for meth, use_jac, with_jac, banded in itertools.product(*p):
|
|
check_complex(idx, "zvode", meth, use_jac, with_jac, banded)
|