75 lines
1.8 KiB
Python
75 lines
1.8 KiB
Python
|
import numpy as np
|
||
|
|
from numpy.testing import assert_equal, assert_allclose
|
||
|
|
from scipy.integrate import odeint
|
||
|
|
import scipy.integrate._test_odeint_banded as banded5x5
|
||
|
|
|
||
|
|
|
||
|
|
def rhs(y, t):
|
||
|
|
dydt = np.zeros_like(y)
|
||
|
|
banded5x5.banded5x5(t, y, dydt)
|
||
|
|
return dydt
|
||
|
|
|
||
|
|
|
||
|
|
def jac(y, t):
|
||
|
|
n = len(y)
|
||
|
|
jac = np.zeros((n, n), order='F')
|
||
|
|
banded5x5.banded5x5_jac(t, y, 1, 1, jac)
|
||
|
|
return jac
|
||
|
|
|
||
|
|
|
||
|
|
def bjac(y, t):
|
||
|
|
n = len(y)
|
||
|
|
bjac = np.zeros((4, n), order='F')
|
||
|
|
banded5x5.banded5x5_bjac(t, y, 1, 1, bjac)
|
||
|
|
return bjac
|
||
|
|
|
||
|
|
|
||
|
|
JACTYPE_FULL = 1
|
||
|
|
JACTYPE_BANDED = 4
|
||
|
|
|
||
|
|
|
||
|
|
def check_odeint(jactype):
|
||
|
|
if jactype == JACTYPE_FULL:
|
||
|
|
ml = None
|
||
|
|
mu = None
|
||
|
|
jacobian = jac
|
||
|
|
elif jactype == JACTYPE_BANDED:
|
||
|
|
ml = 2
|
||
|
|
mu = 1
|
||
|
|
jacobian = bjac
|
||
|
|
else:
|
||
|
|
raise ValueError(f"invalid jactype: {jactype!r}")
|
||
|
|
|
||
|
|
y0 = np.arange(1.0, 6.0)
|
||
|
|
# These tolerances must match the tolerances used in banded5x5.f.
|
||
|
|
rtol = 1e-11
|
||
|
|
atol = 1e-13
|
||
|
|
dt = 0.125
|
||
|
|
nsteps = 64
|
||
|
|
t = dt * np.arange(nsteps+1)
|
||
|
|
|
||
|
|
sol, info = odeint(rhs, y0, t,
|
||
|
|
Dfun=jacobian, ml=ml, mu=mu,
|
||
|
|
atol=atol, rtol=rtol, full_output=True)
|
||
|
|
yfinal = sol[-1]
|
||
|
|
odeint_nst = info['nst'][-1]
|
||
|
|
odeint_nfe = info['nfe'][-1]
|
||
|
|
odeint_nje = info['nje'][-1]
|
||
|
|
|
||
|
|
y1 = y0.copy()
|
||
|
|
# Pure Fortran solution. y1 is modified in-place.
|
||
|
|
nst, nfe, nje = banded5x5.banded5x5_solve(y1, nsteps, dt, jactype)
|
||
|
|
|
||
|
|
# It is likely that yfinal and y1 are *exactly* the same, but
|
||
|
|
# we'll be cautious and use assert_allclose.
|
||
|
|
assert_allclose(yfinal, y1, rtol=1e-12)
|
||
|
|
assert_equal((odeint_nst, odeint_nfe, odeint_nje), (nst, nfe, nje))
|
||
|
|
|
||
|
|
|
||
|
|
def test_odeint_full_jac():
|
||
|
|
check_odeint(JACTYPE_FULL)
|
||
|
|
|
||
|
|
|
||
|
|
def test_odeint_banded_jac():
|
||
|
|
check_odeint(JACTYPE_BANDED)
|