def frozen_random_walk(name, num_steps=100, num_frozen=10):
# last random value is repeated frozen-1 times
num_random = min(max(0, num_steps - num_frozen), num_steps)
num_frozen = num_steps - num_random
rw = numpyro.sample(name, dist.GaussianRandomWalk(num_steps=num_random))
rw = np.concatenate((rw, np.repeat(rw[-1], num_frozen)))
return rw
def model(returns):
step_size = numpyro.sample("sigma", dist.Exponential(50.0))
s = numpyro.sample(
"s", dist.GaussianRandomWalk(scale=step_size, num_steps=jnp.shape(returns)[0])
)
nu = numpyro.sample("nu", dist.Exponential(0.1))
return numpyro.sample(
"r", dist.StudentT(df=nu, loc=0.0, scale=jnp.exp(s)), obs=returns
)
def model(returns):
step_size = numpyro.sample('sigma', dist.Exponential(50.))
s = numpyro.sample(
's',
dist.GaussianRandomWalk(scale=step_size,
num_steps=np.shape(returns)[0]))
nu = numpyro.sample('nu', dist.Exponential(.1))
return numpyro.sample('r',
dist.StudentT(df=nu, loc=0., scale=np.exp(-2 * s)),
obs=returns)
def ExponentialRandomWalk(loc=1., scale=1e-2, drift=0., num_steps=100):
'''
Return distrubtion of exponentiated Gaussian random walk
Variables are x_0, ..., x_{T-1}
Dynamics in log-space are random walk with drift:
log(x_0) := log(loc)
log(x_t) := log(x_{t-1}) + drift + eps_t, eps_t ~ N(0, scale)
==> Dynamics in non-log space are:
x_0 := loc
x_t := x_{t-1} * exp(drift + eps_t), eps_t ~ N(0, scale)
'''
log_loc = np.log(loc) + drift * (np.arange(num_steps) + 0.)
return dist.TransformedDistribution(
dist.GaussianRandomWalk(scale=scale, num_steps=num_steps), [
dist.transforms.AffineTransform(loc=log_loc, scale=1.),
dist.transforms.ExpTransform()
])
def __call__(self,
T=50,
N=1e5,
T_future=0,
E_duration_est=4.0,
I_duration_est=2.0,
R0_est=3.0,
beta_shape=1.,
sigma_shape=100.,
gamma_shape=100.,
det_prob_est=0.3,
det_prob_conc=50.,
confirmed_dispersion=0.3,
death_dispersion=0.3,
rw_scale=2e-1,
forecast_rw_scale=0.,
drift_scale=None,
num_frozen=0,
rw_use_last=1,
confirmed=None,
death=None):
'''
Stochastic SEIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
# Sample dispersion parameters around specified values
rw = numpyro.sample("rw",
dist.GaussianRandomWalk(scale=100, num_steps=T))
D0 = numpyro.sample("D0", dist.Normal(.000002 * N, 1000))
# Sample parameters
if death is None:
death = None
death0 = None
else:
death = clean_daily_obs(death)
death0 = death[0]
# First observation
z_hat = np.exp(np.cumsum(rw) + np.log(D0))
z0 = observe_normal("z0", z_hat[0], .95, death_dispersion, obs=death0)
y0 = observe_normal("y0", z_hat[0], .95, death_dispersion, obs=death0)
z = observe_normal("z", z_hat, .95, death_dispersion, obs=death)
y = observe_normal("y", z_hat, .95, death_dispersion, obs=death)
z = np.append(z0, z)
y = np.append(y0, y)
if (T_future > 0):
z_hat_future = np.exp(
np.cumsum(np.repeat(rw[-1], T_future)) + np.log(D0))
z_future = observe_normal("z_future",
z_hat_future,
.95,
death_dispersion,
obs=death)
y_future = observe_normal("y_future",
z_hat_future,
.95,
death_dispersion,
obs=death)
z = np.append(z, z_future)
y = np.append(y, y_future)
return None, None, y, z, None, None
