def test_keyboard_interrupt_caught():
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])
y_init = D.array([1., 0.])
def kb_callback(ode_sys):
if ode_sys.t[-1] > D.pi:
raise KeyboardInterrupt()
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 pytest.raises(KeyboardInterrupt):
a.integrate(callback=kb_callback)
assert (a.integration_status ==
"A KeyboardInterrupt exception was raised during integration.")
def test_callback_called():
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])
y_init = D.array([1., 0.])
callback_called = False
def callback(ode_sys):
nonlocal callback_called
if not callback_called and ode_sys.t[-1] > D.pi:
callback_called = True
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)
a.integrate(callback=callback)
assert (callback_called)
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_brentsroot():
for fmt in D.available_float_fmt():
print("Set dtype to:", fmt)
D.set_float_fmt(fmt)
for _ in range(10):
ac_prod = D.array(np.random.uniform(0.9, 1.1))
a = D.array(np.random.uniform(-1, 1))
a = D.to_float(-1 * (a <= 0) + 1 * (a > 0))
c = ac_prod / a
b = D.sqrt(0.01 + 4 * ac_prod)
gt_root = -b / (2 * a) - 0.1 / (2 * a)
ub = -b / (2 * a)
lb = -b / (2 * a) - 1.0 / (2 * a)
fun = lambda x: a * x**2 + b * x + c
assert (D.to_numpy(D.to_float(D.abs(fun(gt_root)))) <=
32 * D.epsilon())
root, success = de.utilities.optimizer.brentsroot(fun, [lb, ub],
4 * D.epsilon(),
verbose=True)
assert (success)
assert (np.allclose(D.to_numpy(D.to_float(gt_root)),
D.to_numpy(D.to_float(root)), 32 * D.epsilon(),
32 * D.epsilon()))
assert (D.to_numpy(D.to_float(D.abs(fun(root)))) <=
32 * D.epsilon())
def test_brentsroot_wrong_order():
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(True)
a = D.array(1.0)
b = D.array(1.0)
gt_root = -b / a
lb, ub = -b / a - 1, -b / a + 1
fun = lambda x: a * x + b
assert (D.to_numpy(D.to_float(D.abs(fun(gt_root)))) <= 32 * D.epsilon())
root, success = de.utilities.optimizer.brentsroot(fun, [ub, lb],
4 * D.epsilon(),
verbose=True)
assert (success)
assert (np.allclose(D.to_numpy(D.to_float(gt_root)),
D.to_numpy(D.to_float(root)), 32 * D.epsilon(),
32 * D.epsilon()))
def test_logical_xor_out():
a = D.array([True, False, False, True], dtype=D.bool)
b = D.array([False, False, True, True], dtype=D.bool)
ref = D.array([True, False, True, False], dtype=D.bool)
out = D.zeros_like(a, dtype=D.bool)
D.logical_xor(a, b, out=out)
assert (D.all(out == ref))
def test_dt_dir_fix(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()))
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])
y_init = D.array([1., 0.])
a = de.OdeSystem(rhs,
y0=y_init,
dense_output=False,
t=(0, 2 * D.pi),
dt=-0.01,
rtol=D.epsilon()**0.5,
atol=D.epsilon()**0.5,
constants=dict(k=1.0))
def test_integration_and_nearestfloat_no_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=False, 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.abs(a.t[-2] - a[2*D.pi].t) <= D.abs(a.dt))
def test_non_callable_rhs():
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(de_mat, y0=y_init, dense_output=False, t=(0, 2*D.pi), dt=0.01, rtol=D.epsilon()**0.5, atol=D.epsilon()**0.5, constants=dict(k=1.0))
a.tf = 0.0
def test_brentsroot_same_sign():
if D.backend() == 'torch':
import torch
torch.set_printoptions(precision=17)
torch.autograd.set_detect_anomaly(True)
ac_prod = D.array(np.random.uniform(0.9, 1.1))
a = D.array(1.0)
b = D.array(1.0)
gt_root = -b / a
lb, ub = -b / a - 1, -b / a - 2
fun = lambda x: a * x + b
assert (D.to_numpy(D.to_float(D.abs(fun(gt_root)))) <= 32 * D.epsilon())
root, success = de.utilities.optimizer.brentsroot(fun, [lb, ub],
4 * D.epsilon(),
verbose=True)
assert (np.isinf(root))
assert (not success)
def test_newtonraphson_pytorch_jacobian(ffmt, tol):
print("Set dtype to:", ffmt)
D.set_float_fmt(ffmt)
np.random.seed(21)
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(False)
if ffmt == 'gdual_vdouble':
pytest.skip("Root-finding is ill-conceived with vectorised gduals")
for _ in range(10):
ac_prod = D.array(np.random.uniform(0.9, 1.1))
a = D.array(np.random.uniform(-1, 1))
a = D.to_float(-1 * (a <= 0) + 1 * (a > 0))
c = ac_prod / a
b = D.sqrt(0.01 + 4 * ac_prod)
gt_root1 = -b / (2 * a) - 0.1 / (2 * a)
gt_root2 = -b / (2 * a) + 0.1 / (2 * a)
ub = -b / (2 * a) - 0.2 / (2 * a)
lb = -b / (2 * a) - 0.4 / (2 * a)
x0 = D.array(np.random.uniform(ub, lb))
fun = lambda x: a * x**2 + b * x + c
assert (D.to_numpy(D.to_float(D.abs(fun(gt_root1)))) <=
32 * D.epsilon())
assert (D.to_numpy(D.to_float(D.abs(fun(gt_root2)))) <=
32 * D.epsilon())
root, (success, num_iter,
prec) = de.utilities.optimizer.newtonraphson(fun,
x0,
tol=tol,
verbose=True)
if tol is None:
tol = D.epsilon()
conv_root1 = np.allclose(D.to_numpy(D.to_float(gt_root1)),
D.to_numpy(D.to_float(root)), 128 * tol,
32 * tol)
conv_root2 = np.allclose(D.to_numpy(D.to_float(gt_root2)),
D.to_numpy(D.to_float(root)), 128 * tol,
32 * tol)
print(conv_root1, conv_root2, root, gt_root1, gt_root2, x0,
root - gt_root1, root - gt_root2, num_iter, prec)
assert (success)
assert (conv_root1 or conv_root2)
assert (D.to_numpy(D.to_float(D.abs(fun(root)))) <= 32 * tol)
def test_torch_append():
a = np.array([1.0, 2.0, 3.0])
b = np.array([1.0])
a_torch = D.array(a)
b_torch = D.array(b)
assert (np.all(
np.append(a, b) == D.append(a_torch, b_torch).cpu().numpy()))
def test_logical_not_out_where():
a = D.array([True, False, False, True], dtype=D.bool)
ref = D.array([False, False, False, False], dtype=D.bool)
out = D.zeros_like(a, dtype=D.bool)
where = D.array([True, False, False, True], dtype=D.bool)
D.logical_not(a, out=out, where=where)
assert (D.all(out[where] == ref[where]))
def test_torch_append_axis_none():
a = np.array([[1.0, 2.0, 3.0]])
b = np.array([[1.0]])
a_torch = D.array(a)
b_torch = D.array(b)
assert (np.all(
np.append(a, b, axis=None) == D.append(a_torch, b_torch,
axis=None).cpu().numpy()))
def test_torch_concatenate():
a = np.array([1.0, 2.0, 3.0])
b = np.array([1.0])
a_torch = D.array(a)
b_torch = D.array(b)
assert (np.all(
np.concatenate([a, b, a, b]) == D.concatenate(
[a_torch, b_torch, a_torch, b_torch]).cpu().numpy()))
def test_torch_concatenate_with_none_axis():
a = np.array([[1.0, 2.0, 3.0]])
b = np.array([[1.0]])
a_torch = D.array(a)
b_torch = D.array(b)
assert (np.all(
np.concatenate([a, b, a, b], axis=None) == D.concatenate(
[a_torch, b_torch, a_torch, b_torch], axis=None).cpu().numpy()))
def test_contract_first_ndims_case_2(ffmt):
arr1 = D.array([[2.0, 1.0], [1.0, 0.0]])
arr2 = D.array([[1.0, 1.0], [-1.0, 1.0]])
arr4 = D.contract_first_ndims(arr1, arr2, 2)
true_arr4 = D.array(2.)
assert (D.norm(arr4 - true_arr4) <= 2 * D.epsilon())
def test_contract_first_ndims_case_1(ffmt):
arr1 = D.array([[2.0, 1.0], [1.0, 0.0]])
arr2 = D.array([[1.0, 1.0], [-1.0, 1.0]])
arr3 = D.contract_first_ndims(arr1, arr2, 1)
true_arr3 = D.array([1.0, 1.0])
assert (D.norm(arr3 - true_arr3) <= 2 * D.epsilon())
def test_dense_right_interval_vec():
denseoutput = DenseOutput(None, None)
inputs = D.array([0, 1, 0, 1, 1, 1])
interpolator = CubicHermiteInterp(*inputs)
denseoutput.add_interpolant(1, interpolator)
denseoutput.add_interpolant(2, interpolator)
assert (D.all(
denseoutput.find_interval_vec([0.5, 0.99999, 1.00001, 1.5]) == D.array(
[0, 0, 1, 1], dtype=D.int64)))
def test_newtonraphson_dims(ffmt, tol, dim):
print("Set dtype to:", ffmt)
D.set_float_fmt(ffmt)
np.random.seed(30)
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(False)
if ffmt == 'gdual_vdouble':
pytest.skip("Root-finding is ill-conceived with vectorised gduals")
shift = D.array(np.random.uniform(1, 10, size=(dim, )))
exponent = D.array(np.random.uniform(1, 5, size=(dim, )))
gt_root1 = shift**(1 / exponent)
gt_root2 = -shift**(1 / exponent)
def fun(x):
return x**exponent - shift
def jac(x):
return D.diag(exponent * D.reshape(x, (-1, ))**(exponent - 1))
x0 = D.array(np.random.uniform(1, 3, size=(dim, )))
print(gt_root1, gt_root2)
print(x0)
print(fun(x0))
print(jac(x0))
root, (success, num_iter,
prec) = de.utilities.optimizer.newtonraphson(fun,
x0,
jac=jac,
tol=tol,
verbose=True)
if tol is None:
tol = D.epsilon()
assert (success)
conv_root1 = D.stack([
D.array(np.allclose(D.to_numpy(D.to_float(r1)),
D.to_numpy(D.to_float(r)), 128 * tol, 32 * tol),
dtype=D.bool) for r, r1 in zip(root, gt_root1)
])
conv_root2 = D.stack([
D.array(np.allclose(D.to_numpy(D.to_float(r2)),
D.to_numpy(D.to_float(r)), 128 * tol, 32 * tol),
dtype=D.bool) for r, r2 in zip(root, gt_root2)
])
assert (D.all(conv_root1 | conv_root2))
def test_gdual_double_solve_linear(self):
D.set_float_fmt('gdual_double')
A = D.array([
[D.gdual_double(-1.0, 'a11', 5),
D.gdual_double(3 / 2, 'a12', 5)],
[D.gdual_double(1.0, 'a21', 5),
D.gdual_double(-1.0, 'a22', 5)],
])
b = D.array([[D.gdual_double(1.0, 'b1', 5)],
[D.gdual_double(1.0, 'b2', 5)]])
self.do(A, b)
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_event_detection():
for ffmt in D.available_float_fmt():
if ffmt == 'float16':
continue
D.set_float_fmt(ffmt)
print("Testing event detection for 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
])
y_init = D.array([1., 0.])
def time_event(t, y, **kwargs):
return t - D.pi/8
time_event.is_terminal = True
time_event.direction = 0
a = de.OdeSystem(rhs, y0=y_init, dense_output=True, t=(0, D.pi/4), 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__):
try:
a.set_method(i)
print("Testing {}".format(a.integrator))
a.integrate(eta=True, events=time_event)
if D.abs(a.t[-1] - D.pi/8) > 10*D.epsilon():
print("Event detection with integrator {} failed with t[-1] = {}".format(a.integrator, a.t[-1]))
raise RuntimeError("Failed to detect event for integrator {}".format(str(i)))
else:
print("Event detection with integrator {} succeeded with t[-1] = {}".format(a.integrator, a.t[-1]))
a.reset()
except Exception as e:
raise e
raise RuntimeError("Test failed for integration method: {}".format(a.integrator))
print("")
print("{} backend test passed successfully!".format(D.backend()))
def test_gdual_vdouble_solve_linear(self):
D.set_float_fmt('gdual_vdouble')
A = D.array([
[
D.gdual_vdouble([-1.0, 1 / 2], 'a11', 5),
D.gdual_vdouble([3 / 2, 3 / 2], 'a12', 5)
],
[
D.gdual_vdouble([1.0, 1.0], 'a21', 5),
D.gdual_vdouble([-1.0, -1.0], 'a22', 5)
],
])
b = D.array([[D.gdual_vdouble([1.0, -1.0], 'b1', 5)],
[D.gdual_vdouble([1.0, 1.0], 'b2', 5)]])
self.do(A, b)
def integ(self, potential, events=None):
"""
Create and solve initial value problem for the particle. Takes the
potential (which we expect to be a function of t and r) and a boundary
time.
Args:
potential: method, with args (t, r). Passed to _dQ_dt. e.g. a
`trap.potential`
events: (list of) callable: event method(s) to be passed to Desolver
integrator
"""
# Initialise Desolver integrator
integ = de.OdeSystem(self._dQ_dt,
y0=D.array(self.Q(0)),
dense_output=True,
t=(self._t_0, self._t_end),
dt=self._dt,
constants={'potential': potential})
integ.set_method('SymplecticEulerSolver')
integ.integrate(events=events)
# Look for early completion by comparing last element of evaluation
# times to desired end time. Use isclose because integ can be off by
# fp error
integ_end_time = integ.t[-1]
if not np.isclose(integ_end_time, self._t_end):
self._terminate(integ_end_time)
# Post-processing of resulting trajectory
self._result = [integ[float(t)][1] for t in self.result_times]
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 integ(self, potential, events=None):
"""
Create and solve initial value problem for the particle. Takes the
potential (which we expect to be a function of t and r) and a boundary
time.
Args:
potential: method, with args (t, r). Passed to _dQ_dt. e.g. a
`trap.potential`
events: (list of) callable: event method(s) to be passed to Desolver
integrator
"""
# Ensure that time values are defined
conditions = [self._t_0 is None, self._t_end is None, self._dt is None]
if True in conditions:
raise RuntimeError('Particle time values undefined')
# Check if sample points are defined
if self._points is None:
raise NotImplemented('Cannot yet handle dense output') # FIXME
# Construct dQ_dt, which is to be stepped by the integrator
integ = de.OdeSystem(self._dQ_dt,
y0=D.array(self.Q(0)),
dense_output=False,
t=(self._t_0, self._t_end),
dt=self._dt,
constants={'potential': potential})
integ.set_method('SymplecticEulerSolver')
integ.integrate(events=events)
# Look for early completion by comparing last element of evaluation
# times to desired end time
integ_end_time = integ.t[-1]
if integ_end_time != self._t_end:
self._terminate(integ_end_time)
self._result = [integ[float(t)][1] for t in self.result_times]
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]
def test_jacobian_wrapper_calls_estimate(ffmt):
D.set_float_fmt(ffmt)
rhs = lambda x: D.exp(-x)
jac_rhs = de.utilities.JacobianWrapper(rhs, richardson_iter=0, adaptive=False, rtol=D.epsilon() ** 0.5, atol=D.epsilon() ** 0.5)
x = D.array(0.0)
assert (D.allclose(D.to_float(jac_rhs.estimate(x)), D.to_float(jac_rhs(x)), rtol=4 * D.epsilon() ** 0.5, atol=4 * D.epsilon() ** 0.5))
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
])
