def test_fwd_back(): # Run a system forwards then backwards, # and check that we end up in the same place. D = 10 t0 = 0.1 t1 = 2.2 y0 = np.linspace(0.1, 0.9, D) def f(y, t): return -np.sqrt(t) - y + 0.1 - np.mean((y + 0.2)**2) ys = odeint(f, y0, np.array([t0, t1]), atol=1e-8, rtol=1e-8) rys = odeint(f, ys[-1], np.array([t1, t0]), atol=1e-8, rtol=1e-8) assert np.allclose(y0, rys[-1])
def vjp_all(g): vjp_y = g[-1, :] vjp_t0 = 0 time_vjp_list = [] vjp_args = np.zeros(np.size(flat_args)) for i in range(T - 1, 0, -1): # Compute effect of moving measurement time. vjp_cur_t = np.dot(func(yt[i, :], t[i], *func_args), g[i, :]) time_vjp_list.append(vjp_cur_t) vjp_t0 = vjp_t0 - vjp_cur_t # Run augmented system backwards to the previous observation. aug_y0 = np.hstack((yt[i, :], vjp_y, vjp_t0, vjp_args)) aug_ans = odeint(augmented_dynamics, aug_y0, np.stack([t[i], t[i - 1]]), (flat_args, )) _, vjp_y, vjp_t0, vjp_args = unpack(aug_ans[1]) # Add gradient from current output. vjp_y = vjp_y + g[i - 1, :] time_vjp_list.append(vjp_t0) vjp_times = np.hstack(time_vjp_list)[::-1] return None, vjp_y, vjp_times, unravel(vjp_args)
def odeint2(y0, t0, t1, fargs): return odeint(y0, np.array([t0, t1]), fargs, func=f, atol=1e-8, rtol=1e-8)
def twoarg_odeint(y0, args): return odeint(f, y0, np.array([t0, t1]), args=args, atol=1e-8, rtol=1e-8)[1]
def test_odeint_vjp(): D = 3 t0 = 0.1 t1 = 0.2 y0 = np.linspace(0.1, 0.9, D) fargs = (0.1, 0.2) def f(y, t, arg1, arg2): return -np.sqrt(t) - y + arg1 - np.mean((y + arg2)**2) def onearg_odeint(args): return np.sum(odeint(f, *args, atol=1e-8, rtol=1e-8)) numerical_grad = nd(onearg_odeint, (y0, np.array([t0, t1]), fargs)) ys = odeint(f, y0, np.array([t0, t1]), fargs, atol=1e-8, rtol=1e-8) ode_vjp = grad_odeint(ys, f, y0, np.array([t0, t1]), fargs) g = np.ones_like(ys) exact_grad, _ = ravel_pytree(ode_vjp(g)) assert np.allclose(numerical_grad, exact_grad)
init_state, unpack = ravel_pytree((y0, tan_y0)) def augmented_dynamics(augmented_state, t, fargs): # state and senstivity state y, a = unpack(augmented_state) a = ad.instantiate_zeros(y, a) # combined dynamics dy_dt, da_dt = jvp(func, (y, t, fargs), (a, tan_t0, tan_fargs)) # pack back to give dynamics of augmented_state return np.concatenate([dy_dt, da_dt]) # Solve augmented dynamics aug_sol = odeint(init_state, np.array([t0, t1]), fargs, func=augmented_dynamics) yt, at = unpack(aug_sol[1]) # Sensitivities of y(t1) wrt t0 and t1 jvp_t_total = (tan_t1 - tan_t0) * func(yt, t1, fargs) # Combine sensitivities tan_yt = jvp_t_total if at is ad_util.zero else at + jvp_t_total return (np.array([y0, yt]), np.array([tan_y0, tan_yt])) ad.primitive_jvps[odeint.primitive] = jvp_odeint
def odeint_fwrap(y0, ts, fargs): return odeint(y0, ts, func=f, fargs=fargs)
def odeint2(y0, ts, fargs): return odeint(f, y0, ts, fargs, atol=1e-8, rtol=1e-8)
def ode_w_linear_part(func, y0, a0, t0, t1, func_args): # Just a wrapper around odeint for dynamics that are linear in a0, but not in y0. aug_y0, unpack = ravel_pytree((y0, a0)) aug_ans = odeint(func, aug_y0, np.array([t0, t1]), func_args) yt, jvp_all = unpack(aug_ans[1]) return yt, jvp_all
def onearg_odeint(args): return np.sum(odeint(f, *args, atol=1e-8, rtol=1e-8))
def onearg_odeint(y0): return odeint(f, y0, np.array([t0, t1]), atol=1e-8, rtol=1e-8)[1]