def __init__(self, func, times, *, n_states, n_theta, t0=0):
if not callable(func):
raise ValueError("Argument func must be callable.")
if n_states < 1:
raise ValueError("Argument n_states must be at least 1.")
if n_theta <= 0:
raise ValueError("Argument n_theta must be positive.")
# Public
self.func = func
self.t0 = t0
self.times = tuple(times)
self.n_times = len(times)
self.n_states = n_states
self.n_theta = n_theta
self.n_p = n_states + n_theta
# Private
self._augmented_times = np.insert(times, 0, t0).astype(floatX)
self._augmented_func = utils.augment_system(func, self.n_states,
self.n_theta)
self._sens_ic = utils.make_sens_ic(self.n_states, self.n_theta, floatX)
# Cache symbolic sensitivities by the hash of inputs
self._apply_nodes = {}
self._output_sensitivities = {}
def test_gradients():
"""Tests the computation of the sensitivities from the Aesara computation graph"""
# ODE system for which to compute gradients
def ode_func(y, t, p):
return np.exp(-t) - p[0] * y[0]
# Computation of graidients with Aesara
augmented_ode_func = augment_system(ode_func, 1, 1 + 1)
# This is the new system, ODE + Sensitivities, which will be integrated
def augmented_system(Y, t, p):
dydt, ddt_dydp = augmented_ode_func(Y[:1], t, p, Y[1:])
derivatives = np.concatenate([dydt, ddt_dydp])
return derivatives
# Create real sensitivities
y0 = 0.0
t = np.arange(0, 12, 0.25).reshape(-1, 1)
a = 0.472
p = np.array([y0, a])
# Derivatives of the analytic solution with respect to y0 and alpha
# Treat y0 like a parameter and solve analytically. Then differentiate.
# I used CAS to get these derivatives
y0_sensitivity = np.exp(-a * t)
a_sensitivity = (-(np.exp(t * (a - 1)) - 1 + (a - 1) *
(y0 * a - y0 - 1) * t) * np.exp(-a * t) / (a - 1)**2)
sensitivity = np.c_[y0_sensitivity, a_sensitivity]
integrated_solutions = odeint(func=augmented_system,
y0=[y0, 1, 0],
t=t.ravel(),
args=(p, ))
simulated_sensitivity = integrated_solutions[:, 1:]
np.testing.assert_allclose(sensitivity, simulated_sensitivity, rtol=1e-5)