def transition_fn(x, y):
probs = init_probs if x is None else transition_probs[x]
with numpyro.plate("D", D, dim=-1):
x = numpyro.sample("x", dist.Categorical(probs))
numpyro.sample("y", dist.Normal(locs[x], 1), obs=y)
return x, None
def model(data):
mean = numpyro.param('mean', 0.)
std = numpyro.param('std', 1., constraint=constraints.positive)
return numpyro.sample('obs', dist.Normal(mean, std), obs=data)
def schools_model():
mu = numpyro.sample('mu', dist.Normal(0, 5))
tau = numpyro.sample('tau', dist.HalfCauchy(5))
theta = numpyro.sample('theta', dist.Normal(mu, tau), sample_shape=(data['J'],))
numpyro.sample('obs', dist.Normal(theta, data['sigma']), obs=data['y'])
def model_c(nu1, y1, e1):
Rp = numpyro.sample('Rp', dist.Uniform(0.5, 1.5))
Mp = numpyro.sample('Mp', dist.Normal(33.5, 0.3))
sigma = numpyro.sample('sigma', dist.Exponential(1.0))
RV = numpyro.sample('RV', dist.Uniform(25.0, 35.0))
MMR_CO = numpyro.sample('MMR_CO', dist.Uniform(0.0, maxMMR_CO))
MMR_H2O = numpyro.sample('MMR_H2O', dist.Uniform(0.0, maxMMR_H2O))
T0 = numpyro.sample('T0', dist.Uniform(1000.0, 1700.0))
alpha = numpyro.sample('alpha', dist.Uniform(0.05, 0.2))
vsini = numpyro.sample('vsini', dist.Uniform(10.0, 20.0))
g = 2478.57730044555 * Mp / Rp**2 #gravity
#Kipping Limb Darkening Prior arxiv:1308.0009
q1 = numpyro.sample('q1', dist.Uniform(0.0, 1.0))
q2 = numpyro.sample('q2', dist.Uniform(0.0, 1.0))
sqrtq1 = jnp.sqrt(q1)
u1 = 2.0 * sqrtq1 * q2
u2 = sqrtq1 * (1.0 - 2.0 * q2)
#T-P model//
Tarr = T0 * (Parr / Pref)**alpha
#line computation CO
qt_CO = vmap(mdbCO1.qr_interp)(Tarr)
qt_H2O = vmap(mdbH2O1.qr_interp)(Tarr)
def obyo(y, tag, nusd, nus, numatrix_CO, numatrix_H2O, mdbCO, mdbH2O,
cdbH2H2, cdbH2He):
#CO
SijM_CO=jit(vmap(SijT,(0,None,None,None,0)))\
(Tarr,mdbCO.logsij0,mdbCO.dev_nu_lines,mdbCO.elower,qt_CO)
gammaLMP_CO = jit(vmap(gamma_exomol,(0,0,None,None)))\
(Parr,Tarr,mdbCO.n_Texp,mdbCO.alpha_ref)
gammaLMN_CO = gamma_natural(mdbCO.A)
gammaLM_CO = gammaLMP_CO + gammaLMN_CO[None, :]
sigmaDM_CO=jit(vmap(doppler_sigma,(None,0,None)))\
(mdbCO.dev_nu_lines,Tarr,molmassCO)
xsm_CO = xsmatrix(numatrix_CO, sigmaDM_CO, gammaLM_CO, SijM_CO)
dtaumCO = dtauM(dParr, xsm_CO, MMR_CO * ONEARR, molmassCO, g)
#H2O
SijM_H2O=jit(vmap(SijT,(0,None,None,None,0)))\
(Tarr,mdbH2O.logsij0,mdbH2O.dev_nu_lines,mdbH2O.elower,qt_H2O)
gammaLMP_H2O = jit(vmap(gamma_exomol,(0,0,None,None)))\
(Parr,Tarr,mdbH2O.n_Texp,mdbH2O.alpha_ref)
gammaLMN_H2O = gamma_natural(mdbH2O.A)
gammaLM_H2O = gammaLMP_H2O + gammaLMN_H2O[None, :]
sigmaDM_H2O=jit(vmap(doppler_sigma,(None,0,None)))\
(mdbH2O.dev_nu_lines,Tarr,molmassH2O)
xsm_H2O = xsmatrix(numatrix_H2O, sigmaDM_H2O, gammaLM_H2O, SijM_H2O)
dtaumH2O = dtauM(dParr, 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(nusd, nus, Frot, beta, RV)
errall = jnp.sqrt(e1**2 + sigma**2)
numpyro.sample(tag, dist.Normal(mu, errall), obs=y)
obyo(y1, "y1", nusd1, nus1, numatrix_CO1, numatrix_H2O1, mdbCO1, mdbH2O1,
cdbH2H21, cdbH2He1)
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=8,
gamma_shape=8,
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,
death=None):
'''
Stochastic SEIR 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))
H0 = numpyro.sample("H0", dist.Uniform(0, 0.01 * N))
D0 = numpyro.sample("D0", dist.Uniform(0, 0.01 * 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 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_prob,
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_prob,
det_noise_scale, death_prob, death_rate, det_prob_d)
beta, x, y, z = self.dynamics(T,
params,
x0,
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, det_noise_scale, death_prob, death_rate,
det_prob_d)
beta_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 test_block():
with handlers.trace() as trace:
with handlers.block(hide=["x"]):
with handlers.seed(rng_seed=0):
numpyro.sample("x", dist.Normal())
assert "x" not in trace
def model(y):
mu = numpyro.sample("mu", dist.Normal())
sigma = numpyro.sample("sigma", dist.HalfNormal())
with numpyro.plate("plate1", y.shape[0]):
numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
def model(data, labels):
coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim),
np.ones(dim)))
offset = numpyro.sample('offset', dist.Uniform(-1, 1))
logits = offset + np.sum(coefs * data, axis=-1)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
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 transition_fn(name, probs, locs, x, y):
x = numpyro.sample(name, dist.Categorical(probs[x]))
numpyro.sample("y_" + name, dist.Normal(locs[x], 1), obs=y)
return x, None
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)
def model(data):
x = w = 0
for i, y in markov(enumerate(data)):
x = numpyro.sample(f"x_{i}", dist.Categorical(probs_x[x]))
w = numpyro.sample(f"w_{i}", dist.Categorical(probs_w[w]))
numpyro.sample(f"y_{i}", dist.Normal(locs[w, x], 1), obs=y)
def transition_fn(x, y):
x = numpyro.sample("x", dist.Categorical(probs[w, x]))
numpyro.sample("y", dist.Normal(locs[x], 1), obs=y)
return x, None
def model(data):
w = numpyro.sample("w", dist.Bernoulli(0.6))
x = 0
for i, y in markov(enumerate(data)):
x = numpyro.sample(f"x_{i}", dist.Categorical(probs[w, x]))
numpyro.sample(f"y_{i}", dist.Normal(locs[x], 1), obs=y)
def guide(subsample):
scale = numpyro.param("scale", 1.0)
with handlers.substitute(data={"data": subsample}):
with numpyro.plate("data", len(data), subsample_size):
loc = numpyro.param("loc", jnp.zeros(len(data)), event_dim=0)
numpyro.sample("z", dist.Normal(loc, scale))
def actual_model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
loc = numpyro.sample('loc', dist.Uniform(0, alpha))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def guide():
with numpyro.plate_stack("plates", shape):
return numpyro.sample("x", dist.Normal(0, 1))
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 fn():
return numpyro.sample("x", dist.Normal())
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():
alpha = numpyro.sample("alpha", dist.Normal())
numpyro.deterministic("alpha", alpha * 2)
def model(data):
x = numpyro.sample("x", dist.Normal(0, 1))
with optional(use_context_manager, handlers.scale(scale=10)):
numpyro.sample("obs", dist.Normal(x, 1), obs=data)
def model():
x = numpyro.sample("x", dist.Normal())
with numpyro.handlers.mask(mask=False):
numpyro.sample("y", dist.Normal(x), obs=1)
def _sample():
x = numpyro.sample("x", dist.Normal(0.0, 1.0))
y = numpyro.sample("y", dist.Normal(1.0, 2.0))
return jnp.stack([x, y])
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
loc = numpyro.param('loc', 0., constraint=constraints.interval(0., alpha))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def model():
x = numpyro.sample("x", dist.Normal(0.0, 1.0))
y = numpyro.sample("y", dist.Normal(0.0, 1.0))
return x + y
def model(data):
loc = numpyro.sample('loc', dist.Normal(0., 1.))
numpyro.sample('obs', dist.Normal(loc, 1), obs=data)
def model(subsample):
with handlers.substitute(data={"data": subsample}):
with numpyro.plate("data", len(data), subsample_size) as ind:
x = data[ind]
z = numpyro.sample("z", dist.Normal(0, 1))
numpyro.sample("x", dist.Normal(z, 1), obs=x)
def model(labels):
coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim), np.ones(dim)))
logits = np.sum(coefs * data, axis=-1)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
def likelihood_func(self, yhat):
"""Return a normal likelihood with fitted sigma."""
_sigma = numpyro.sample("_sigma", dist.Exponential(self.sigma_prior))
return dist.Normal(yhat, _sigma)
