def model(data):
concentration = jnp.array([1.0, 1.0, 1.0])
p_latent = numpyro.sample('p_latent', dist.Dirichlet(concentration))
numpyro.sample('obs', dist.Categorical(p_latent), obs=data)
return p_latent
def SIR_stochastic(T=50,
N=1e5,
T_future=0,
duration_mean=10,
R0_mean=2.2,
gamma_shape=5,
beta_shape=5,
det_rate_mean=0.3,
det_rate_conc=50,
det_conc=100,
drift_scale=5e-2,
obs=None):
'''
Stochastic SIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(0, N * 0.02))
# Sample parameters
gamma = numpyro.sample(
"gamma", dist.Gamma(gamma_shape, gamma_shape * duration_mean))
beta0 = numpyro.sample(
"beta0", dist.Gamma(beta_shape, beta_shape * duration_mean / R0_mean))
det_rate = numpyro.sample(
"det_rate",
dist.Beta(det_rate_mean * det_rate_conc,
(1 - det_rate_mean) * det_rate_conc))
x0 = SIRModel.seed(N, I0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
obs0, obs = (None, None) if obs is None else (obs[0], obs[1:])
# First observation
y0 = observe("y0", x0[3], det_rate, det_conc, obs=obs0)
# Run dynamics
params = (beta0, gamma, det_rate, det_conc, drift_scale)
beta, x, y = SIR_dynamics(T, params, x0, obs=obs)
x = np.vstack((x0, x))
y = np.append(y0, y)
if T_future > 0:
params = (beta[-1], gamma, det_rate, det_conc, drift_scale)
beta_f, x_f, y_f = SIR_dynamics(T_future + 1,
params,
x[-1, :],
suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
return beta, x, y, det_rate
def model():
a = numpyro.param('a',
a_init,
constraint=constraints.greater_than(a_minval))
b = numpyro.param('b', b_init, constraint=constraints.positive)
numpyro.sample('x', dist.Normal(a, b))
def model(returns):
step_size = numpyro.sample('sigma', dist.Exponential(50.))
s = numpyro.sample('s', dist.GaussianRandomWalk(scale=step_size, num_steps=jnp.shape(returns)[0]))
nu = numpyro.sample('nu', dist.Exponential(.1))
return numpyro.sample('r', dist.StudentT(df=nu, loc=0., scale=jnp.exp(s)),
obs=returns)
def SEIRD_stochastic(T=50,
N=1e5,
T_future=0,
E_duration_est=4.0,
I_duration_est=3.0,
R0_est=3.0,
beta_shape=1,
sigma_shape=10,
gamma_shape=10,
det_prob_est=0.3,
det_prob_conc=50,
det_noise_scale=0.15,
rw_scale=1e-1,
drift_scale=None,
obs=None,
death=None,
hosp=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))
# 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))
det_prob_d = numpyro.sample("det_prob_d",
dist.Beta(.9 * 100, (1 - .9) * 100))
death_prob = numpyro.sample("death_prob",
dist.Beta(.1 * 100, (1 - .1) * 100))
death_rate = numpyro.sample("death_rate", dist.Gamma(10, 10 * 10))
if True: #drift_scale is not None:
drift = numpyro.sample("drift", dist.Normal(loc=np.log(.5), scale=.01))
drift = -np.exp(drift)
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
obs0, obs = (None, None) if obs is None else (obs[0], obs[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_prob, det_noise_scale, obs=obs0)
z0 = observe("z0", x0[5], det_prob_d, det_noise_scale, obs=death0)
params = (beta0, sigma, gamma, rw_scale, drift, det_prob, det_noise_scale,
death_prob, death_rate, det_prob_d)
beta, x, y, z = SEIRD_dynamics(T, params, x0, obs=obs, 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, rw_scale, drift, det_prob,
det_noise_scale, death_prob, death_rate, det_prob_d)
beta_f, x_f, y_f, z_f = SEIRD_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 expected_model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
loc = numpyro.sample('loc', dist.Uniform(0, 1)) * alpha
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def model(x, y):
nn = numpyro.module("nn", Dense(1), (10, ))
mu = nn(x).squeeze(-1)
sigma = numpyro.sample("sigma", dist.HalfNormal(1))
numpyro.sample('y', dist.Normal(mu, sigma), obs=y)
def model_pred() -> np.ndarray:
mu = numpyro.sample("mu", dist.Normal(0, 5))
tau = numpyro.sample("tau", dist.HalfCauchy(5))
return numpyro.sample("obs", dist.Normal(mu, tau))
def model(data):
loc = numpyro.sample("loc", dist.Normal(0., 1.))
numpyro.sample("obs", dist.Normal(loc, 1.), obs=data)
def model(data, labels):
dim = data.shape[1]
coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
logits = jnp.dot(data, coefs)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
def model_c(nu1, y1, e1):
Rp = sample('Rp', dist.Uniform(0.5, 1.5))
Mp = sample('Mp', dist.Normal(33.5, 0.3))
RV = sample('RV', dist.Uniform(26.0, 30.0))
MMR_CO = sample('MMR_CO', dist.Uniform(0.0, maxMMR_CO))
MMR_H2O = sample('MMR_H2O', dist.Uniform(0.0, maxMMR_H2O))
T0 = sample('T0', dist.Uniform(1000.0, 1700.0))
alpha = sample('alpha', dist.Uniform(0.05, 0.15))
vsini = sample('vsini', dist.Uniform(10.0, 20.0))
# Kipping Limb Darkening Prior
q1 = sample('q1', dist.Uniform(0.0, 1.0))
q2 = sample('q2', dist.Uniform(0.0, 1.0))
u1, u2 = ld_kipping(q1, q2)
# GP
logtau = sample('logtau', dist.Uniform(-1.5, 0.5)) #tau=1 <=> 5A
tau = 10**(logtau)
loga = sample('loga', dist.Uniform(-4.0, -2.0))
a = 10**(loga)
#gravity
g = getjov_gravity(Rp, Mp)
# T-P model
Tarr = T0 * (Parr / Pref)**alpha
# Line computation
qt_CO = vmap(mdbCO.qr_interp)(Tarr)
qt_H2O = vmap(mdbH2O.qr_interp)(Tarr)
def obyo(y, tag, nusdx, nus, mdbCO, mdbH2O, cdbH2H2, cdbH2He):
#CO
SijM_CO, ngammaLM_CO, nsigmaDl_CO = exomol(mdbCO, Tarr, Parr, R_CO,
molmassCO)
xsm_CO = xsmatrix(cnu_CO, indexnu_CO, R_CO, pmarray_CO, nsigmaDl_CO,
ngammaLM_CO, SijM_CO, nus, dgm_ngammaL_CO)
dtaumCO = dtauM(dParr, jnp.abs(xsm_CO), MMR_CO * ONEARR, molmassCO, g)
#H2O
SijM_H2O, ngammaLM_H2O, nsigmaDl_H2O = exomol(mdbH2O, Tarr, Parr,
R_H2O, molmassH2O)
xsm_H2O = xsmatrix(cnu_H2O, indexnu_H2O, R_H2O, pmarray_H2O,
nsigmaDl_H2O, ngammaLM_H2O, SijM_H2O, nus,
dgm_ngammaL_H2O)
dtaumH2O = dtauM(dParr, jnp.abs(xsm_H2O), MMR_H2O * ONEARR, molmassH2O,
g)
#CIA
dtaucH2H2=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrH2,\
mmw,g,cdbH2H2.nucia,cdbH2H2.tcia,cdbH2H2.logac)
dtaucH2He=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrHe,\
mmw,g,cdbH2He.nucia,cdbH2He.tcia,cdbH2He.logac)
dtau = dtaumCO + dtaumH2O + dtaucH2H2 + dtaucH2He
sourcef = planck.piBarr(Tarr, nus)
Ftoa = Fref / Rp**2
F0 = rtrun(dtau, sourcef) / baseline / Ftoa
Frot = response.rigidrot(nus, F0, vsini, u1, u2)
mu = response.ipgauss_sampling(nusdx, nus, Frot, beta, RV)
cov = gpkernel_RBF(nu1, tau, a, e1)
sample(tag,
dist.MultivariateNormal(loc=mu, covariance_matrix=cov),
obs=y)
obyo(y1, "y1", nu1, nus, mdbCO, mdbH2O, cdbH2H2, cdbH2He)
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):
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(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 body_fn(z, val):
z = numpyro.sample("z", dist.Normal(z, 1))
return z, z
def guide(data):
guide_loc = numpyro.param("guide_loc", 0.)
guide_scale = np.exp(numpyro.param("guide_scale_log", 0.))
numpyro.sample("loc", dist.Normal(guide_loc, guide_scale))
def transition(state, i):
x, mu = state
x = numpyro.sample("x", dist.LogNormal(phi * x, q))
mu = beta * mu + x
numpyro.sample("y", dist.Normal(mu, r))
return (x, mu), None
def model(deterministic=True):
GLOBAL["count"] += 1
x = numpyro.sample("x", dist.Normal())
if deterministic:
numpyro.deterministic("x_copy", x)
def model(x, y):
a = numpyro.sample("a", dist.Normal(0, 10))
b = numpyro.sample("b", dist.Normal(0, 10), sample_shape=(3, ))
mu = a + b[0] * x + b[1] * x**2 + b[2] * x**3
numpyro.sample("y", dist.Normal(mu, 0.001), obs=y)
def model(X,ndims,ndata,y_obs=None):
w = numpyro.sample('w', dist.Normal(np.zeros(ndims), np.ones(ndims)))
sigma = numpyro.sample('sigma', dist.Exponential(1.))
y = numpyro.sample('y', dist.Normal(np.matmul(X, w), sigma * np.ones(ndata)), obs=y_obs)
def model(data):
f = numpyro.sample('beta', dist.Beta(jnp.ones(2), jnp.ones(2)))
numpyro.sample('obs', dist.Bernoulli(f), obs=data)
def true_fun(_):
x = numpyro.sample("x", dist.Normal(4.0))
numpyro.deterministic("z", x - 4.0)
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
def false_fun(_):
x = numpyro.sample("x", dist.Normal(0.0))
numpyro.deterministic("z", x)
def dual_moon_model():
numpyro.sample('x', DualMoonDistribution())
def true_fun(_):
numpyro.sample("x", dist.Normal(m1, s1))
def model(data, labels):
coefs = numpyro.sample('coefs',
dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
offset = numpyro.sample('offset', dist.Uniform(-1, 1))
logits = offset + jnp.sum(coefs * data, axis=-1)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
def false_fun(_):
numpyro.sample("x", dist.Normal(m2, s2))
def actual_model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with numpyro.handlers.reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample('loc', dist.Uniform(0, alpha))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def model():
x = numpyro.sample("x", dist.Normal())
with numpyro.handlers.mask(mask_array=False):
numpyro.sample("y", dist.Normal(x), obs=1)
