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])