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=5,
gamma_shape=8,
det_prob_est=0.3,
det_prob_conc=50,
det_noise_scale=0.15,
rw_scale=2e-1,
forecast_rw_scale=0,
drift_scale=None,
num_frozen=0,
confirmed=None,
death=None):
'''
Stochastic SEIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(0, 0.02 * N))
E0 = numpyro.sample("E0", dist.Uniform(0, 0.02 * N))
H0 = numpyro.sample("H0", dist.Uniform(0, 0.02 * N))
D0 = numpyro.sample("D0", dist.Uniform(0, 0.02 * N))
# Sample parameters
sigma = numpyro.sample(
"sigma", dist.Gamma(sigma_shape, sigma_shape * E_duration_est))
gamma = numpyro.sample(
"gamma", dist.Gamma(gamma_shape, gamma_shape * I_duration_est))
beta0 = numpyro.sample(
"beta0",
dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est))
det_prob0 = numpyro.sample(
"det_prob0",
dist.Beta(det_prob_est * det_prob_conc,
(1 - det_prob_est) * det_prob_conc))
det_prob_d = numpyro.sample("det_prob_d",
dist.Beta(.9 * 100, (1 - .9) * 100))
death_prob = numpyro.sample("death_prob",
dist.Beta(.01 * 100, (1 - .01) * 100))
death_rate = numpyro.sample("death_rate", dist.Gamma(10, 10 * 10))
if drift_scale is not None:
drift = numpyro.sample("drift",
dist.Normal(loc=0, scale=drift_scale))
else:
drift = 0
x0 = SEIRDModel.seed(N=N, I=I0, E=E0, H=H0, D=D0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
confirmed0, confirmed = (None, None) if confirmed is None else (
confirmed[0], confirmed[1:])
death0, death = (None, None) if death is None else (death[0],
death[1:])
# First observation
with numpyro.handlers.scale(scale_factor=0.5):
y0 = observe("y0",
x0[6],
det_prob0,
det_noise_scale,
obs=confirmed0)
with numpyro.handlers.scale(scale_factor=2.0):
z0 = observe("z0", x0[5], det_prob_d, det_noise_scale, obs=death0)
params = (beta0, sigma, gamma, rw_scale, drift, det_prob0,
det_noise_scale, death_prob, death_rate, det_prob_d)
beta, det_prob, x, y, z = self.dynamics(T,
params,
x0,
num_frozen=num_frozen,
confirmed=confirmed,
death=death)
x = np.vstack((x0, x))
y = np.append(y0, y)
z = np.append(z0, z)
if T_future > 0:
params = (beta[-1], sigma, gamma, forecast_rw_scale, drift,
det_prob[-1], det_noise_scale, death_prob, death_rate,
det_prob_d)
beta_f, det_rate_rw_f, x_f, y_f, z_f = self.dynamics(
T_future + 1, params, x[-1, :], suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
z = np.append(z, z_f)
return beta, x, y, z, det_prob, death_prob
def __call__(self,
T = 50,
N = 1e5,
T_future = 0,
E_duration_est = 5.5,
I_duration_est = 3.0,
R0_est = 3.0,
beta_shape = 1,
sigma_shape = .2,
gamma_shape = .2,
det_prob_est = 0.15,
det_prob_conc = 50,
det_noise_scale = 0.15,
rw_scale = 2e-1,
forecast_rw_scale = 0,
drift_scale = None,
#confirmed=None,
hosp=None):
'''
Stochastic SEIHR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(10, 0.05*N))
E0 = numpyro.sample("E0", dist.Uniform(10, 0.05*N))
R0 = numpyro.sample("R0", dist.Uniform(0, 0.01*N))
H0 = numpyro.sample("H0", dist.Uniform(0, 0.01*N))
# Sample parameters
sigma = numpyro.sample("sigma",
dist.LogNormal(np.log(1/E_duration_est),sigma_shape))
gamma = numpyro.sample("gamma",
dist.LogNormal(np.log(1/I_duration_est),gamma_shape))
# gamma = numpyro.sample("gamma",
# dist.TruncatedNormal(loc = 1./I_duration_est, scale = 0.25)
beta0 = numpyro.sample("beta0",
dist.Gamma(beta_shape, beta_shape * I_duration_est/R0_est))
det_prob = numpyro.sample("det_prob",
dist.Beta(det_prob_est * det_prob_conc,
(1-det_prob_est) * det_prob_conc))
hosp_prob = numpyro.sample("hosp_prob",
dist.Beta(.1 * 100,
(1-.1) * 100))
if drift_scale is not None:
drift = numpyro.sample("drift", dist.Normal(loc=0, scale=drift_scale))
else:
drift = 0
x0 = SEIHRModel.seed(N=N, I=I0, E=E0, R=R0, H=H0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
#confirmed0, confirmed = (None, None) if confirmed is None else (confirmed[0], confirmed[1:])
hosp0, hosp = (None, None) if hosp is None else (hosp[0], hosp[1:])
# First observation
with numpyro.handlers.scale(scale_factor=2.0):
y0 = observe_nonrandom("y0", x0[3], det_noise_scale, obs=hosp0)
params = (beta0, sigma, gamma,
rw_scale, drift,
det_prob, det_noise_scale,
hosp_prob)
beta, x, y = self.dynamics(T, params, x0, hosp=hosp)
x = np.vstack((x0, x))
y = np.append(y0, y)
#z = np.append(z0, z)
if T_future > 0:
params = (beta[-1], sigma, gamma,
forecast_rw_scale, drift,
det_prob, det_noise_scale,
hosp_prob)
beta_f, x_f, y_f = self.dynamics(T_future+1, params, x[-1,:],
suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
#z = np.append(z, z_f)
return beta, x, y, det_prob, hosp_prob
def model(data):
alpha = np.array([1.1, 1.1])
beta = np.array([1.1, 1.1])
p_latent = sample('p_latent', dist.Beta(alpha, beta))
sample('obs', dist.Bernoulli(p_latent), obs=data)
return p_latent
def model(data):
f = numpyro.sample("beta", dist.Beta(1., 1.))
numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def guide(data):
alpha_q = numpyro.param("alpha_q",
1.0,
constraint=constraints.positive)
beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive)
numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
def model(data=None):
with numpyro.plate("dim", 2):
beta = numpyro.sample("beta", dist.Beta(1., 1.))
with numpyro.plate("plate", N, dim=-2):
numpyro.deterministic("beta_sq", beta**2)
numpyro.sample("obs", dist.Bernoulli(beta), obs=data)
def model(data):
f = numpyro.sample("beta", dist.Beta(1., 1.))
with numpyro.plate("plate", 10):
numpyro.deterministic("beta_sq", f**2)
numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def model():
c = numpyro.sample("c", dist.Gamma(1, 1))
with handlers.collapse():
probs = numpyro.sample("probs", dist.Beta(c, 2))
with numpyro.plate("plate", len(data)):
numpyro.sample("obs", dist.Binomial(10, probs), obs=data)
def model():
c = numpyro.sample("c", dist.Gamma(1, 1))
with handlers.collapse():
probs = numpyro.sample("probs", dist.Beta(c, 2))
numpyro.sample("obs", dist.Bernoulli(probs), obs=data)
def model1():
c1 = numpyro.param("c1", 0.5, constraint=dist.constraints.positive)
c0 = numpyro.param("c0", 1.5, constraint=dist.constraints.positive)
with handlers.collapse():
probs = numpyro.sample("probs", dist.Beta(c1, c0))
numpyro.sample("obs", dist.Binomial(total_count, probs), obs=data)
def model(data):
f = numpyro.sample('beta', dist.Beta(np.ones(2), np.ones(2)))
numpyro.sample('obs', dist.Bernoulli(f), obs=data)
def model(data):
y_prob = numpyro.sample("y_prob", dist.Beta(1., 1.))
with numpyro.plate("data", data.shape[0]):
y = numpyro.sample("y", dist.Bernoulli(y_prob))
z = numpyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
numpyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)