def test_back_and_forth_strat():
# 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()
fr, gr, br, tr = time_reflect_stratonovich(f, g, b, ts)
ys = stratonovich_integrate(f, g, y0, ts, b, dt)
rys = stratonovich_integrate(fr, gr, ys[-1], tr, br, dt)[::-1]
assert np.allclose(ys, rys, rtol=5e-03, atol=5e-03)
def test_strat_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
def strat_int(argv):
y0, args = argv
ys = stratonovich_integrate(f, g, y0, ts, bm, dt, args)
return np.sum(ys[1])
ys = stratonovich_integrate(f, g, y0, ts, bm, dt, args)
v_yt = np.ones_like(y0)
v_argst = np.zeros_like(flat_args)
y0_rec, exact_grad = vjp_strat_integrate(v_yt, v_argst, ys[-1], f, g, ts,
bm, dt, args)
numerical_grad = numerical_gradient(strat_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 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 strat_int(argv):
y0, args = argv
ys = stratonovich_integrate(f, g, y0, ts, bm, dt, args)
return np.sum(ys[1])