def test_ito_milstein_vs_euler():
# Check that for a small enough step size, Euler matches Milstein
f, g, b, y0, ts, dt = make_example_sde(dt=0.1 * 2**-8)
ys = ito_integrate(f, g, y0, ts, b, dt)
ys2 = ito_integrate(f, g, y0, ts, b, dt, method='euler_maruyama')
assert np.allclose(ys, ys2, rtol=1e-02, atol=1e-02)
def test_back_and_forth_ito():
# Run a system forwards then backwards,
# and check that we end up in the same place.
f, g, b, y0, ts, dt = make_example_sde(dt=0.0001)
fr, gr, br, tr = time_reflect_ito(f, g, b, ts)
ys = ito_integrate(f, g, y0, ts, b, dt)
rys = ito_integrate(fr, gr, ys[-1], tr, br, dt)[::-1]
assert np.allclose(ys[0], rys[0], rtol=1e-03, atol=1e-03)
def test_strat_to_ito_int():
# Run a system forwards then backwards,
# and check that we end up in the same place.
f, g, b, y0, ts, dt = make_example_sde()
fi, gi = stratonovich_to_ito(f, g)
iys = ito_integrate(fi, gi, y0, ts, b, dt)
sys = stratonovich_integrate(f, g, y0, ts, b, dt)
assert np.allclose(iys[-1], sys[-1], rtol=1e-02, atol=1e-02)
def test_integrating_noise_noop():
# Check that integrating changes in Brownian motion gives the same endpoint as the original noise.
f, g, b, y0, ts, dt = make_example_sde()
# Create an SDE that will only integrate the noise.
y0 = y0 * 0.0
f = lambda y, t, args: np.zeros(y0.shape)
g = lambda y, t, args: np.ones(y0.shape)
sols_i = ito_integrate(f, g, y0, ts, b, dt)
sols_s = stratonovich_integrate(f, g, y0, ts, b, dt)
bs = np.array([b(ti) for ti in ts])
assert np.allclose(bs, sols_i, rtol=1e-03, atol=1e-03)
assert np.allclose(bs, sols_s, rtol=1e-03, atol=1e-03)
def test_ito_int_vjp():
D, ts, y0, args, f, g = make_sde()
flat_args, _ = ravel_pytree(args)
bm = make_brownian_motion(ts[0], np.zeros(y0.shape), ts[1],
random.PRNGKey(0))
dt = 1e-4
eps = 1e-6
method = 'milstein'
def ito_int(argv):
y0, args = argv
ys = ito_integrate(f, g, y0, ts, bm, dt, args, method=method)
return np.sum(ys[1])
ys = ito_integrate(f, g, y0, ts, bm, dt, args, method=method)
v_yt = np.ones_like(y0)
v_argst = np.zeros_like(flat_args)
y0_rec, exact_grad = vjp_ito_integrate(v_yt,
v_argst,
ys[-1],
f,
g,
ts,
bm,
dt,
args,
method=method)
numerical_grad = numerical_gradient(ito_int, (y0, args), eps=eps)
print("states:", y0, y0_rec)
assert np.allclose(y0, y0_rec, rtol=1e-2, atol=1e-02)
flat_grads, unravel = ravel_pytree(exact_grad)
print("numerical grad: ", unravel(numerical_grad))
print(" exact grad: ", exact_grad)
assert np.allclose(numerical_grad, flat_grads, rtol=1e-2, atol=1e-2)
def ito_int(argv):
y0, args = argv
ys = ito_integrate(f, g, y0, ts, bm, dt, args, method=method)
return np.sum(ys[1])
def onearg_int(y0):
return ito_integrate(f, g, y0, ts, b, dt)