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_double_reflect_ito():
# Check that reflecting twice gives the same answer.
f, g, ts, y0, args = make_example_sde()
b = make_brownian_motion(ts[0], np.zeros(y0.shape), ts[-1], rng)
f2, g2, b2, t2 = time_reflect_ito(*time_reflect_ito(f, g, b, ts))
t = 0.1
assert (np.all(ts == t2))
assert (np.allclose(f(y0, t, args), f2(y0, t, args)))
assert (np.allclose(g(y0, t, args), g2(y0, t, args)))
assert (np.all(b(t) == b2(t)))
def test_mean_and_var():
t0 = 0.0
t1 = 3.0
y0 = np.linspace(0.1, 0.9, D)
num_samples = 300
vals = onp.zeros((num_samples, D))
for i in range(num_samples):
rng = random.PRNGKey(i)
bm = make_brownian_motion(t0, np.zeros(y0.shape), t1, rng)
vals[i, :] = bm(t1)
assert np.allclose(np.mean(vals), 0.0, atol=1e-1, rtol=1e-1)
assert np.allclose(np.var(vals), t1, atol=1e-1, rtol=1e-1)
def make_example_sde(dt=0.1 * 2**-8):
D = 3
ts = np.array([0.1, 0.2])
y0 = np.linspace(0.1, 0.9, D)
def f(y, t, args):
return -np.sqrt(t) - y + 0.1 - np.mean((y + 0.2)**2)
def g(y, t, args):
return y**2 + np.sin(0.1 * y) + np.cos(t)
b = make_brownian_motion(ts[0], np.zeros(y0.shape), ts[1], rng)
return f, g, b, y0, ts, dt
def test_reflect_ito_two_ways():
# Check that reflect_ito = strat_to_ito( reflect_strat ( ito_to_strat )))
f, g, ts, y0, args = make_example_sde()
b = make_brownian_motion(ts[0], np.zeros(y0.shape), ts[-1], rng)
fr, gr, br, tr = time_reflect_ito(f, g, b, ts)
fi, gi = ito_to_stratonovich(f, g)
f2, g2, b3, t3 = time_reflect_stratonovich(fi, gi, b, ts)
f3, g3 = stratonovich_to_ito(f2, g2)
t = 0.1
assert (np.all(tr == t3))
assert (np.allclose(fr(y0, -t, args), f3(y0, -t, args)))
assert (np.allclose(gr(y0, -t, args), g3(y0, -t, args)))
assert (np.all(br(-t) == b3(-t)))