def __new__(cls,
name,
rho,
*args,
steps=None,
constant=False,
ar_order=None,
**kwargs):
rhos = at.atleast_1d(at.as_tensor_variable(floatX(rho)))
ar_order = cls._get_ar_order(rhos=rhos,
constant=constant,
ar_order=ar_order)
steps = get_steps(
steps=steps,
shape=None, # Shape will be checked in `cls.dist`
dims=kwargs.get("dims", None),
observed=kwargs.get("observed", None),
step_shape_offset=ar_order,
)
return super().__new__(cls,
name,
rhos,
*args,
steps=steps,
constant=constant,
ar_order=ar_order,
**kwargs)
def dist(
cls,
rho,
sigma=None,
tau=None,
*,
init_dist=None,
steps=None,
constant=False,
ar_order=None,
**kwargs,
):
_, sigma = get_tau_sigma(tau=tau, sigma=sigma)
sigma = at.as_tensor_variable(floatX(sigma))
rhos = at.atleast_1d(at.as_tensor_variable(floatX(rho)))
if "init" in kwargs:
warnings.warn(
"init parameter is now called init_dist. Using init will raise an error in a future release.",
FutureWarning,
)
init_dist = kwargs.pop("init")
ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
steps = get_steps(steps=steps, shape=kwargs.get("shape", None), step_shape_offset=ar_order)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = at.as_tensor_variable(intX(steps), ndim=0)
if init_dist is not None:
if not isinstance(init_dist, TensorVariable) or not isinstance(
init_dist.owner.op, RandomVariable
):
raise ValueError(
f"Init dist must be a distribution created via the `.dist()` API, "
f"got {type(init_dist)}"
)
check_dist_not_registered(init_dist)
if init_dist.owner.op.ndim_supp > 1:
raise ValueError(
"Init distribution must have a scalar or vector support dimension, ",
f"got ndim_supp={init_dist.owner.op.ndim_supp}.",
)
else:
warnings.warn(
"Initial distribution not specified, defaulting to "
"`Normal.dist(0, 100, shape=...)`. You can specify an init_dist "
"manually to suppress this warning.",
UserWarning,
)
init_dist = Normal.dist(0, 100, shape=(*sigma.shape, ar_order))
# Tell Aeppl to ignore init_dist, as it will be accounted for in the logp term
init_dist = ignore_logprob(init_dist)
return super().dist([rhos, sigma, init_dist, steps, ar_order, constant], **kwargs)
def augment_system(ode_func, n_states, n_theta):
"""
Function to create augmented system.
Take a function which specifies a set of differential equations and return
a compiled function which allows for computation of gradients of the
differential equation's solition with repsect to the parameters.
Uses float64 even if floatX=float32, because the scipy integrator always uses float64.
Parameters
----------
ode_func: function
Differential equation. Returns array-like.
n_states: int
Number of rows of the sensitivity matrix. (n_states)
n_theta: int
Number of ODE parameters
Returns
-------
system: function
Augemted system of differential equations.
"""
# Present state of the system
t_y = at.vector("y", dtype="float64")
t_y.tag.test_value = np.ones((n_states, ), dtype="float64")
# Parameter(s). Should be vector to allow for generaliztion to multiparameter
# systems of ODEs. Is m dimensional because it includes all initial conditions as well as ode parameters
t_p = at.vector("p", dtype="float64")
t_p.tag.test_value = np.ones((n_states + n_theta, ), dtype="float64")
# Time. Allow for non-automonous systems of ODEs to be analyzed
t_t = at.scalar("t", dtype="float64")
t_t.tag.test_value = 2.459
# Present state of the gradients:
# Will always be 0 unless the parameter is the inital condition
# Entry i,j is partial of y[i] wrt to p[j]
dydp_vec = at.vector("dydp", dtype="float64")
dydp_vec.tag.test_value = make_sens_ic(n_states, n_theta, "float64")
dydp = dydp_vec.reshape((n_states, n_states + n_theta))
# Get symbolic representation of the ODEs by passing tensors for y, t and theta
yhat = ode_func(t_y, t_t, t_p[n_states:])
if isinstance(yhat, at.TensorVariable):
t_yhat = at.atleast_1d(yhat)
else:
# Stack the results of the ode_func into a single tensor variable
if not isinstance(yhat, (list, tuple)):
raise TypeError(
f"Unexpected type, {type(yhat)}, returned by ode_func. TensorVariable, list or tuple is expected."
)
t_yhat = at.stack(yhat, axis=0)
if t_yhat.ndim > 1:
raise ValueError(
f"The odefunc returned a {t_yhat.ndim}-dimensional tensor, but 0 or 1 dimensions were expected."
)
# Now compute gradients
J = at.jacobian(t_yhat, t_y)
Jdfdy = at.dot(J, dydp)
grad_f = at.jacobian(t_yhat, t_p)
# This is the time derivative of dydp
ddt_dydp = (Jdfdy + grad_f).flatten()
system = aesara.function(inputs=[t_y, t_t, t_p, dydp_vec],
outputs=[t_yhat, ddt_dydp],
on_unused_input="ignore")
return system