def test_integration_and_representation():
for ffmt in D.available_float_fmt():
D.set_float_fmt(ffmt)
print("Testing {} float format".format(D.float_fmt()))
de_mat = D.array([[0.0, 1.0], [-1.0, 0.0]])
@de.rhs_prettifier("""[vx, -x+t]""")
def rhs(t, state, k, **kwargs):
return de_mat @ state + D.array([0.0, t])
def analytic_soln(t, initial_conditions):
c1 = initial_conditions[0]
c2 = initial_conditions[1] - 1
return D.stack([
c2 * D.sin(D.to_float(D.asarray(t))) +
c1 * D.cos(D.to_float(D.asarray(t))) + D.asarray(t),
c2 * D.cos(D.to_float(D.asarray(t))) -
c1 * D.sin(D.to_float(D.asarray(t))) + 1
])
def kbinterrupt_cb(ode_sys):
if ode_sys[-1][0] > D.pi:
raise KeyboardInterrupt("Test Interruption and Catching")
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=True,
t=(0, 2 * D.pi),
dt=0.01,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5,
constants=dict(k=1.0))
a.integrate()
try:
print(str(a))
print(repr(a))
assert (D.max(D.abs(a.sol(a.t[0]) - y_init)) <=
8 * D.epsilon()**0.5)
assert (D.max(
D.abs(a.sol(a.t[-1]) - analytic_soln(a.t[-1], y_init))) <=
8 * D.epsilon()**0.5)
assert (D.max(D.abs(a.sol(a.t).T - analytic_soln(a.t, y_init))) <=
8 * D.epsilon()**0.5)
except:
raise
def test_matrix_inv():
A = D.array([
[-1.0, 3 / 2],
[1.0, -1.0],
], dtype=D.float64)
Ainv = D.matrix_inv(A)
assert (D.max(D.abs(D.to_float(Ainv @ A - D.eye(2)))) <= 8 * D.epsilon())
def test_float_formats_typical_shape(ffmt, integrator,
use_richardson_extrapolation, device):
if use_richardson_extrapolation and integrator.__implicit__:
pytest.skip(
"Richardson Extrapolation is too slow with implicit methods")
D.set_float_fmt(ffmt)
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(False) # Enable if a test fails
device = torch.device(device)
print("Testing {} float format".format(D.float_fmt()))
from .common import set_up_basic_system
de_mat, rhs, analytic_soln, y_init, dt, _ = set_up_basic_system(
integrator, hook_jacobian=True)
y_init = D.array([1., 0.])
if D.backend() == 'torch':
y_init = y_init.to(device)
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=False,
t=(0, D.pi / 4),
dt=D.pi / 64,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5)
method = integrator
method_tolerance = a.atol * 10 + D.epsilon()
if use_richardson_extrapolation:
method = de.integrators.generate_richardson_integrator(method)
method_tolerance = method_tolerance * 5
with de.utilities.BlockTimer(section_label="Integrator Tests") as sttimer:
a.set_method(method)
print("Testing {} with dt = {:.4e}".format(a.integrator, a.dt))
a.integrate(eta=True)
print("Average step-size:",
D.mean(D.abs(D.array(a.t[1:]) - D.array(a.t[:-1]))))
max_diff = D.max(D.abs(analytic_soln(a.t[-1], y_init) - a.y[-1]))
if a.integrator.adaptive:
assert max_diff <= method_tolerance, "{} Failed with max_diff from analytical solution = {}".format(
a.integrator, max_diff)
if a.integrator.__implicit__:
assert rhs.analytic_jacobian_called and a.njev > 0, "Analytic jacobian was called as part of integration"
a.reset()
print("")
print("{} backend test passed successfully!".format(D.backend()))
def test_integration_and_representation():
for ffmt in D.available_float_fmt():
D.set_float_fmt(ffmt)
print("Testing {} float format".format(D.float_fmt()))
de_mat = D.array([[0.0, 1.0],[-1.0, 0.0]])
@de.rhs_prettifier("""[vx, -x+t]""")
def rhs(t, state, k, **kwargs):
return de_mat @ state + D.array([0.0, t])
def analytic_soln(t, initial_conditions):
c1 = initial_conditions[0]
c2 = initial_conditions[1] - 1
return D.stack([
c2 * D.sin(D.to_float(D.asarray(t))) + c1 * D.cos(D.to_float(D.asarray(t))) + D.asarray(t),
c2 * D.cos(D.to_float(D.asarray(t))) - c1 * D.sin(D.to_float(D.asarray(t))) + 1
])
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs, y0=y_init, dense_output=True, t=(0, 2*D.pi), dt=0.01, rtol=D.epsilon()**0.5, atol=D.epsilon()**0.5, constants=dict(k=1.0))
assert(a.integration_status() == "Integration has not been run.")
a.integrate()
assert(a.integration_status() == "Integration completed successfully.")
try:
print(str(a))
print(repr(a))
assert(D.max(D.abs(a.sol(a.t[0]) - y_init)) <= 8*D.epsilon()**0.5)
assert(D.max(D.abs(a.sol(a.t[-1]) - analytic_soln(a.t[-1], y_init))) <= 8*D.epsilon()**0.5)
assert(D.max(D.abs(a.sol(a.t).T - analytic_soln(a.t, y_init))) <= 8*D.epsilon()**0.5)
except:
raise
for i in a:
assert(D.max(D.abs(i.y - analytic_soln(i.t, y_init))) <= 8*D.epsilon()**0.5)
assert(len(a.y) == len(a))
assert(len(a.t) == len(a))
def test_matrix_inv_bigger():
for diag_size in range(2, 101):
np.random.seed(15)
for trial in range(3):
A = np.random.normal(size=(diag_size, diag_size))
while np.abs(np.linalg.det(D.to_float(A))) <= 1e-5:
A = np.random.normal(size=(diag_size, diag_size), std=250.0)
A = D.array(D.cast_to_float_fmt(A))
Ainv = D.matrix_inv(A)
assert (
D.max(D.abs(D.to_float(Ainv @ A - D.eye(diag_size)))) <=
4 * D.epsilon()**0.5
), "Matrix inversion failed for diagonal with size: " + str(
diag_size)
def test_integration_and_representation(ffmt):
D.set_float_fmt(ffmt)
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(True)
print("Testing {} float format".format(D.float_fmt()))
from . import common
(de_mat, rhs, analytic_soln, y_init, dt, a) = common.set_up_basic_system()
assert (a.integration_status == "Integration has not been run.")
a.integrate()
assert (a.integration_status == "Integration completed successfully.")
print(str(a))
print(repr(a))
assert (D.max(D.abs(a.sol(a.t[0]) - y_init)) <= 8 * D.epsilon()**0.5)
assert (D.max(D.abs(a.sol(a.t[-1]) - analytic_soln(a.t[-1], y_init))) <=
8 * D.epsilon()**0.5)
assert (D.max(D.abs(a.sol(a.t).T - analytic_soln(a.t, y_init))) <=
8 * D.epsilon()**0.5)
for i in a:
assert (D.max(D.abs(i.y - analytic_soln(i.t, y_init))) <=
8 * D.epsilon()**0.5)
assert (len(a.y) == len(a))
assert (len(a.t) == len(a))
def test_brentsrootvec(ffmt, tol):
print("Set dtype to:", ffmt)
D.set_float_fmt(ffmt)
if tol is not None:
tol = tol * D.epsilon()
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(True)
if ffmt == 'gdual_vdouble':
pytest.skip("Root-finding is ill-conceived with vectorised gduals")
for _ in range(10):
slope_list = D.array(
np.copysign(np.random.uniform(0.9, 1.1, size=25),
np.random.uniform(-1, 1, size=25)))
intercept_list = slope_list
gt_root_list = -intercept_list / slope_list
fun_list = [(lambda m, b: lambda x: m * x + b)(m, b)
for m, b in zip(slope_list, intercept_list)]
assert (all(
map((lambda i: D.to_numpy(D.to_float(D.abs(i))) <= 32 * D.epsilon(
)), map((lambda x: x[0](x[1])), zip(fun_list, gt_root_list)))))
root_list, success = de.utilities.optimizer.brentsrootvec(
fun_list, [D.min(gt_root_list) - 1.,
D.max(gt_root_list) + 1.],
tol,
verbose=True)
assert (np.all(D.to_numpy(success)))
assert (np.allclose(D.to_numpy(D.to_float(gt_root_list)),
D.to_numpy(D.to_float(root_list)),
32 * D.epsilon(), 32 * D.epsilon()))
assert (all(
map((lambda i: D.to_numpy(D.to_float(D.abs(i))) <= 32 * D.epsilon(
)), map((lambda x: x[0](x[1])), zip(fun_list, root_list)))))
def test_brentsrootvec():
for fmt in D.available_float_fmt():
print("Set dtype to:", fmt)
D.set_float_fmt(fmt)
if fmt == 'gdual_vdouble':
continue
for _ in range(10):
slope_list = D.array(
np.copysign(np.random.uniform(0.9, 1.1, size=25),
np.random.uniform(-1, 1, size=25)))
intercept_list = slope_list
gt_root_list = -intercept_list / slope_list
fun_list = [(lambda m, b: lambda x: m * x + b)(m, b)
for m, b in zip(slope_list, intercept_list)]
assert (all(
map((lambda i: D.to_numpy(D.to_float(D.abs(i))) <= 32 * D.
epsilon()),
map((lambda x: x[0](x[1])), zip(fun_list, gt_root_list)))))
root_list, success = de.utilities.optimizer.brentsrootvec(
fun_list, [D.min(gt_root_list) - 1.,
D.max(gt_root_list) + 1.],
4 * D.epsilon(),
verbose=True)
assert (np.all(D.to_numpy(success)))
assert (np.allclose(D.to_numpy(D.to_float(gt_root_list)),
D.to_numpy(D.to_float(root_list)),
32 * D.epsilon(), 32 * D.epsilon()))
assert (all(
map((lambda i: D.to_numpy(D.to_float(D.abs(i))) <= 32 * D.
epsilon()),
map((lambda x: x[0](x[1])), zip(fun_list, root_list)))))
def test_torch_max_with_axis():
a = np.array([[1.0, 2.0, 3.0], [0.0, 3.0, -1.0]])
a_torch = D.array(a)
assert (np.all(np.max(a, axis=1) == D.max(a_torch, axis=1).cpu().numpy()))
def do(self, A, b):
x = D.solve_linear_system(A, b)
assert (D.max(D.abs(D.to_float(A @ x - b))) <= 256 * D.epsilon())
def do(self, A):
Q, R = D.qr(A)
assert (D.max(D.abs(D.to_float(Q @ R - A))) <= 256 * D.epsilon())
def test_event_detection_numerous_events(ffmt,
integrator,
use_richardson_extrapolation,
device,
dt,
dense_output=False):
if use_richardson_extrapolation and integrator.__implicit__:
pytest.skip(
"Richardson Extrapolation is too slow with implicit methods")
D.set_float_fmt(ffmt)
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(False) # Enable if a test fails
device = torch.device(device)
print("Testing event detection for float format {}".format(D.float_fmt()))
from .common import set_up_basic_system
de_mat, rhs, analytic_soln, y_init, _, _ = set_up_basic_system(
integrator, hook_jacobian=True)
if D.backend() == 'torch':
y_init = y_init.to(device)
event_times = D.linspace(0, D.pi / 4, 32)
class ev_proto:
def __init__(self, ev_time, component):
self.ev_time = ev_time
self.component = component
def __call__(self, t, y, **csts):
return y[self.component] - analytic_soln(self.ev_time,
y_init)[self.component]
def __repr__(self):
return "<ev_proto({}, {})>".format(self.ev_time, self.component)
events = [ev_proto(ev_t, 0) for ev_t in event_times]
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=dense_output,
t=(0, D.pi / 4),
dt=dt,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.75)
method = integrator
if use_richardson_extrapolation:
method = de.integrators.generate_richardson_integrator(method)
with de.utilities.BlockTimer(section_label="Integrator Tests") as sttimer:
a.set_method(method)
print("Testing {} with dt = {:.4e}".format(a.integrator, a.dt))
a.integrate(eta=True, events=events)
print(a)
print(a.events)
print(len(events) - len(a.events))
assert (len(events) - 3 <= len(a.events) <= len(events))
for ev_detected in a.events:
assert (D.max(
D.abs(
ev_detected.event(ev_detected.t, ev_detected.y, **
a.constants))) <= 4 * D.epsilon())
def test_dense_output(ffmt, use_richardson_extrapolation):
D.set_float_fmt(ffmt)
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(True)
print("Testing {} float format".format(D.float_fmt()))
from . import common
(de_mat, rhs, analytic_soln, y_init, dt, a) = common.set_up_basic_system()
assert (a.integration_status == "Integration has not been run.")
if use_richardson_extrapolation:
a.method = de.integrators.generate_richardson_integrator(a.method)
a.rtol = a.atol = D.epsilon()**0.75
a.integrate()
assert (a.integration_status == "Integration completed successfully.")
assert (D.max(D.abs(a[0].y - analytic_soln(a[0].t, y_init))) <=
4 * D.epsilon())
assert (D.max(D.abs(a[0].t)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[-1].y - analytic_soln(a[-1].t, y_init))) <=
10 * D.epsilon()**0.5)
assert (D.max(D.abs(a[a[0].t].y - analytic_soln(a[0].t, y_init))) <=
4 * D.epsilon())
assert (D.max(D.abs(a[a[0].t].t)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[a[-1].t].y - analytic_soln(a[-1].t, y_init))) <=
10 * D.epsilon()**0.5)
assert (D.max(D.abs(D.stack(a[a[0].t:a[-1].t].y) - D.stack(a.y))) <=
4 * D.epsilon())
assert (D.max(D.abs(D.stack(a[:a[-1].t].y) - D.stack(a.y))) <=
4 * D.epsilon())
assert (D.max(D.abs(D.stack(a[a[0].t:a[-1].t:2].y) - D.stack(a.y[::2]))) <=
4 * D.epsilon())
assert (D.max(D.abs(D.stack(a[a[0].t::2].y) - D.stack(a.y[::2]))) <=
4 * D.epsilon())
assert (D.max(D.abs(D.stack(a[:a[-1].t:2].y) - D.stack(a.y[::2]))) <=
4 * D.epsilon())
np.random.seed(42)
sample_points = D.array(np.random.uniform(a.t[0], a.t[-1], 1024))
assert (D.max(
D.abs(a.sol(sample_points) -
analytic_soln(sample_points, y_init).T)).item() <= D.epsilon()**
0.5)
def test_dense_output():
for ffmt in D.available_float_fmt():
D.set_float_fmt(ffmt)
print("Testing {} float format".format(D.float_fmt()))
de_mat = D.array([[0.0, 1.0], [-1.0, 0.0]])
@de.rhs_prettifier("""[vx, -x+t]""")
def rhs(t, state, k, **kwargs):
return de_mat @ state + D.array([0.0, t])
def analytic_soln(t, initial_conditions):
c1 = initial_conditions[0]
c2 = initial_conditions[1] - 1
return D.stack([
c2 * D.sin(D.to_float(D.asarray(t))) +
c1 * D.cos(D.to_float(D.asarray(t))) + D.asarray(t),
c2 * D.cos(D.to_float(D.asarray(t))) -
c1 * D.sin(D.to_float(D.asarray(t))) + 1
])
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=True,
t=(0, 2 * D.pi),
dt=0.01,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5,
constants=dict(k=1.0))
assert (a.integration_status() == "Integration has not been run.")
a.integrate()
assert (
a.integration_status() == "Integration completed successfully.")
assert (D.max(D.abs(a[0].y - analytic_soln(a[0].t, y_init))) <=
4 * D.epsilon())
assert (D.max(D.abs(a[0].t)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[-1].y - analytic_soln(a[-1].t, y_init))) <=
10 * D.epsilon()**0.5)
assert (D.max(D.abs(a[a[0].t].y - analytic_soln(a[0].t, y_init))) <=
4 * D.epsilon())
assert (D.max(D.abs(a[a[0].t].t)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[a[-1].t].y - analytic_soln(a[-1].t, y_init))) <=
10 * D.epsilon()**0.5)
assert (D.max(D.abs(a[a[0].t:a[-1].t].y - a.y)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[:a[-1].t].y - a.y)) <= 4 * D.epsilon())
assert (D.max(D.abs(a[a[0].t:a[-1].t:2].y - a.y[::2])) <=
4 * D.epsilon())
assert (D.max(D.abs(a[a[0].t::2].y - a.y[::2])) <= 4 * D.epsilon())
assert (D.max(D.abs(a[:a[-1].t:2].y - a.y[::2])) <= 4 * D.epsilon())
@de.rhs_prettifier(
equ_repr="[vx, -k*x/m]",
md_repr=r"""
$$
\frac{dx}{dt} = \begin{bmatrix}
0 & 1 \\
-\frac{k}{m} & 0
\end{bmatrix} \cdot \begin{bmatrix}x \\ v_x\end{bmatrix}
$$
"""
)
def rhs(t, state, k, m, **kwargs):
return D.array([[0.0, 1.0], [-k/m, 0.0]])@state
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs, y0=y_init, dense_output=True, t=(0, 2*D.pi), dt=0.01, rtol=1e-9, atol=1e-9, constants=dict(k=1.0, m=1.0))
print(a)
a.integrate()
print(a)
print("If the integration was successful and correct, a[0].y and a[-1].y should be near identical.")
print("a[0].y = {}".format(a[0].y))
print("a[-1].y = {}".format(a[-1].y))
print("Maximum difference from initial state after one oscillation cycle: {}".format(D.max(D.abs(a[0].y-a[-1].y))))
def test_float_formats():
for ffmt in D.available_float_fmt():
D.set_float_fmt(ffmt)
print("Testing {} float format".format(D.float_fmt()))
de_mat = D.array([[0.0, 1.0], [-1.0, 0.0]])
@de.rhs_prettifier("""[vx, -x+t]""")
def rhs(t, state, **kwargs):
return de_mat @ state + D.array([0.0, t])
def analytic_soln(t, initial_conditions):
c1 = initial_conditions[0]
c2 = initial_conditions[1] - 1
return D.array([
c2 * D.sin(t) + c1 * D.cos(t) + t,
c2 * D.cos(t) - c1 * D.sin(t) + 1
])
def kbinterrupt_cb(ode_sys):
if ode_sys[-1][0] > D.pi:
raise KeyboardInterrupt("Test Interruption and Catching")
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=True,
t=(0, 2 * D.pi),
dt=0.01,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5)
with de.utilities.BlockTimer(
section_label="Integrator Tests") as sttimer:
for i in sorted(set(de.available_methods(False).values()),
key=lambda x: x.__name__):
if "Heun-Euler" in i.__name__ and D.float_fmt(
) == "gdual_real128":
print(
"skipping {} due to ridiculous timestep requirements.".
format(i))
continue
try:
a.set_method(i)
print("Testing {}".format(a.integrator))
try:
a.integrate(callback=kbinterrupt_cb, eta=True)
except KeyboardInterrupt as e:
pass
try:
a.integrate(eta=True)
except:
raise
max_diff = D.max(
D.abs(analytic_soln(a.t[-1], a.y[0]) - a.y[-1]))
if a.method.__adaptive__ and max_diff >= a.atol * 10 + D.epsilon(
):
print(
"{} Failed with max_diff from analytical solution = {}"
.format(a.integrator, max_diff))
raise RuntimeError(
"Failed to meet tolerances for adaptive integrator {}"
.format(str(i)))
else:
print(
"{} Succeeded with max_diff from analytical solution = {}"
.format(a.integrator, max_diff))
a.reset()
except Exception as e:
print(e)
raise RuntimeError(
"Test failed for integration method: {}".format(
a.integrator))
print("")
print("{} backend test passed successfully!".format(D.backend()))
def test_float_formats_atypical_shape(ffmt, integrator,
use_richardson_extrapolation, device):
if use_richardson_extrapolation and integrator.__implicit__:
pytest.skip(
"Richardson Extrapolation is too slow with implicit methods")
D.set_float_fmt(ffmt)
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(False) # Enable if a test fails
device = torch.device(device)
print("Testing {} float format".format(D.float_fmt()))
from .common import set_up_basic_system
de_mat, _, analytic_soln, y_init, dt, _ = set_up_basic_system(integrator)
@de.rhs_prettifier("""[vx, -x+t]""")
def rhs(t, state, **kwargs):
nonlocal de_mat
extra = D.array([0.0, t])
if D.backend() == 'torch':
de_mat = de_mat.to(state.device)
extra = extra.to(state.device)
return D.sum(de_mat[:, :, None, None, None] * state,
axis=1) + extra[:, None, None, None]
y_init = D.array([[[[1., 0.]] * 1] * 1] * 3).T
print(rhs(0.0, y_init).shape)
if D.backend() == 'torch':
y_init = y_init.contiguous().to(device)
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=False,
t=(0, D.pi / 4),
dt=D.pi / 64,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5)
method = integrator
method_tolerance = a.atol * 10 + D.epsilon()
if use_richardson_extrapolation:
method = de.integrators.generate_richardson_integrator(method)
method_tolerance = method_tolerance * 5
with de.utilities.BlockTimer(section_label="Integrator Tests") as sttimer:
a.set_method(method)
print("Testing {} with dt = {:.4e}".format(a.integrator, a.dt))
a.integrate(eta=True)
max_diff = D.max(D.abs(analytic_soln(a.t[-1], y_init) - a.y[-1]))
if a.integrator.adaptive:
assert max_diff <= method_tolerance, "{} Failed with max_diff from analytical solution = {}".format(
a.integrator, max_diff)
a.reset()
print("")
print("{} backend test passed successfully!".format(D.backend()))
def do(self, A):
Ainv = D.matrix_inv(A)
assert (D.max(D.abs(D.to_float(Ainv @ A - D.eye(A.shape[0])))) <=
256 * D.epsilon())