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_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 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 test_gradients_simple_oscillator(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) print("Testing {} float format".format(D.float_fmt())) import torch torch.set_printoptions(precision=17) device = torch.device(device) torch.autograd.set_detect_anomaly(False) # Enable if a test fails def rhs(t, state, k, m, **kwargs): return D.array([[0.0, 1.0], [-k / m, 0.0]], device=device) @ state csts = dict(k=1.0, m=1.0) T = 2 * D.pi * D.sqrt(D.array(csts['m'] / csts['k'])).to(device) dt = max(0.5 * (D.epsilon()**0.5)**(1.0 / (max(2, integrator.order - 1))), 5e-2) def true_solution_sho(t, initial_state, k, m): w2 = D.array(k / m).to(device) w = D.sqrt(w2) A = D.sqrt(initial_state[0]**2 + initial_state[1]**2 / w2) phi = D.atan2(-initial_state[1], w * initial_state[0]) return D.stack([A * D.cos(w * t + phi), -w * A * D.sin(w * t + phi)]).T method = integrator if use_richardson_extrapolation: method = de.integrators.generate_richardson_integrator(method) with de.utilities.BlockTimer(section_label="Integrator Tests"): y_init = D.array([1., 1.], requires_grad=True).to(device) a = de.OdeSystem(rhs, y_init, t=(0, T), dt=T * dt, rtol=D.epsilon()**0.5, atol=D.epsilon()**0.5, constants=csts) a.set_method(method) print("Testing {} with dt = {:.4e}".format(a.integrator, a.dt)) a.integrate(eta=True) Jy = D.jacobian(a.y[-1], a.y[0]) true_Jy = D.jacobian(true_solution_sho(a.t[-1], a.y[0], **csts), a.y[0]) print(a.integrator.adaptive, D.mean(D.abs(D.stack(a.t[1:]) - D.stack(a.t[:-1]))), D.norm(true_Jy - Jy), D.epsilon()**0.5) if a.integrator.adaptive: assert (D.allclose(true_Jy, Jy, rtol=4 * a.rtol**0.75, atol=4 * a.atol**0.75)) print("{} method test succeeded!".format(a.integrator)) print("") print("{} backend test passed successfully!".format(D.backend()))
def test_gradients_simple_decay(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 integrator.__symplectic__: pytest.skip( "Exponential decay system is not in the form compatible with symplectic integrators" ) print("Testing {} float format".format(D.float_fmt())) import torch torch.set_printoptions(precision=17) device = torch.device(device) torch.autograd.set_detect_anomaly(False) # Enable if a test fails def rhs(t, state, k, **kwargs): return -k * state def rhs_jac(t, state, k, **kwargs): return -k rhs.jac = rhs_jac y_init = D.array(5.0, requires_grad=True) csts = dict(k=D.array(1.0, device=device)) dt = max(0.5 * (D.epsilon()**0.5)**(1.0 / (max(2, integrator.order - 1))), 5e-2) def true_solution_decay(t, initial_state, k): return initial_state * D.exp(-k * t) method = integrator if use_richardson_extrapolation: method = de.integrators.generate_richardson_integrator(method) with de.utilities.BlockTimer(section_label="Integrator Tests"): y_init = D.ones((1, ), requires_grad=True).to(device) y_init = y_init * D.e a = de.OdeSystem(rhs, y_init, t=(0, 1.0), dt=dt, rtol=D.epsilon()**0.5, atol=D.epsilon()**0.5, constants=csts) a.set_method(method) print("Testing {} with dt = {:.4e}".format(a.integrator, a.dt)) a.integrate(eta=True) Jy = D.jacobian(a.y[-1], a.y[0]) true_Jy = D.jacobian(true_solution_decay(a.t[-1], a.y[0], **csts), a.y[0]) print(a.y[-1], true_solution_decay(a.t[-1], a.y[0], **csts), D.abs(a.y[-1] - true_solution_decay(a.t[-1], a.y[0], **csts))) print(a.integrator.adaptive, D.mean(D.abs(D.stack(a.t[1:]) - D.stack(a.t[:-1]))), D.norm(true_Jy - Jy), 32 * a.rtol) print(a.integrator.adaptive, true_Jy, Jy) if a.integrator.adaptive: assert (D.allclose(true_Jy, Jy, rtol=32 * a.rtol, atol=32 * a.atol)) print("{} method test succeeded!".format(a.integrator)) print("") print("{} backend test passed successfully!".format(D.backend()))
def true_solution_sho(t, initial_state, k, m): w2 = D.array(k / m).to(device) w = D.sqrt(w2) A = D.sqrt(initial_state[0]**2 + initial_state[1]**2 / w2) phi = D.atan2(-initial_state[1], w * initial_state[0]) return D.stack([A * D.cos(w * t + phi), -w * A * D.sin(w * t + phi)]).T