def test_adjoint(self):
"""
Test against dopri5
"""
tf.compat.v1.set_random_seed(0)
f, y0, t_points, _ = problems.construct_problem(TEST_DEVICE)
y0 = tf.cast(y0, tf.float64)
t_points = tf.cast(t_points, tf.float64)
func = lambda y0, t_points: tfdiffeq.odeint(f, y0, t_points, method='dopri5')
with tf.GradientTape() as tape:
tape.watch(t_points)
ys = func(y0, t_points)
reg_t_grad, reg_a_grad, reg_b_grad = tape.gradient(ys, [t_points, f.a, f.b])
f, y0, t_points, _ = problems.construct_problem(TEST_DEVICE)
y0 = tf.cast(y0, tf.float64)
t_points = tf.cast(t_points, tf.float64)
y0 = (y0,)
func = lambda y0, t_points: tfdiffeq.odeint_adjoint(f, y0, t_points, method='dopri5')
with tf.GradientTape() as tape:
tape.watch(t_points)
ys = func(y0, t_points)
grads = tape.gradient(ys, [t_points, f.a, f.b])
adj_t_grad, adj_a_grad, adj_b_grad = grads
self.assertLess(max_abs(reg_t_grad - adj_t_grad), 1.2e-7)
self.assertLess(max_abs(reg_a_grad - adj_a_grad), 1.2e-7)
self.assertLess(max_abs(reg_b_grad - adj_b_grad), 1.2e-7)
def test_dopri5(self):
for ode in problems.PROBLEMS.keys():
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE,
ode=ode)
y = tfdiffeq.odeint(f, y0, t_points, method='dopri5')
with self.subTest(ode=ode):
self.assertLess(rel_error(sol, y), error_tol)
def test_adams_gradient(self):
f, y0, t_points, sol = construct_problem(TEST_DEVICE)
tuple_f = lambda t, y: (f(t, y[0]), f(t, y[1]))
for i in range(2):
func = lambda y0, t_points: tfdiffeq.odeint(tuple_f, (y0, y0), t_points, method='adams')[i]
self.assertTrue(gradcheck(func, (y0, t_points)))
def test_adams(self):
f, y0, t_points, sol = construct_problem(TEST_DEVICE)
tuple_f = lambda t, y: (f(t, y[0]), f(t, y[1]))
tuple_y0 = (y0, y0)
tuple_y = tfdiffeq.odeint(tuple_f, tuple_y0, t_points, method='adams')
max_error0 = tf.reduce_max(sol - tuple_y[0])
max_error1 = tf.reduce_max(sol - tuple_y[1])
self.assertLess(max_error0, eps)
self.assertLess(max_error1, eps)
def test_adaptive_heun(self):
for ode in problems.PROBLEMS.keys():
if ode == 'sine':
# Sine test never finishes.
continue
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE,
ode=ode)
y = tfdiffeq.odeint(f, y0, t_points, method='adaptive_heun')
with self.subTest(ode=ode):
self.assertLess(rel_error(sol, y), error_tol)
def test_adjoint(self):
for ode in problems.PROBLEMS.keys():
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE,
reverse=True)
y0 = tf.cast(y0, tf.float64)
t_points = tf.cast(t_points, tf.float64)
sol = tf.cast(sol, tf.float64)
y = tfdiffeq.odeint_adjoint(f, y0, t_points, method='dopri5')
with self.subTest(ode=ode):
self.assertLess(rel_error(sol, y), error_tol)
def test_adams(self):
f, y0, t_points, _ = problems.construct_problem(TEST_DEVICE)
func = lambda y0, t_points: tfdiffeq.odeint(f, y0, t_points, method='adams')
self.assertTrue(gradcheck(func, (y0, t_points)))
def test_rk4(self):
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE,
reverse=True)
y = tfdiffeq.odeint(f, y0, t_points, method='rk4')
self.assertLess(rel_error(sol, y), error_tol)
def test_explicit_adams(self):
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE)
y = tfdiffeq.odeint(f, y0, t_points, method='explicit_adams')
self.assertLess(rel_error(sol, y), error_tol)
def test_dopri5(self):
f, y0, t_points, sol = problems.construct_problem(TEST_DEVICE,
reverse=True)
y = tfdiffeq.odeint(f, y0, t_points[0:1], method='dopri5')
self.assertLess(max_abs(sol[0] - y), error_tol)
