def model(data):
p = numpyro.sample('p', dist.Beta(1., 1.))
if with_logits:
logits = logit(p)
numpyro.sample('obs', dist.Binomial(data['n'], logits=logits), obs=data['x'])
else:
numpyro.sample('obs', dist.Binomial(data['n'], probs=p), obs=data['x'])
def init_agents(self,
num_steps=1000,
pop_size=(1e2, 1e2),
initial_infections=2):
self.age = ny.sample('age', dist.Uniform(0, 100))
self.sex = ny.sample('sex', dist.Binomial(1, .5))
self.risk_tolerance = ny.sample('risk', dist.Beta(2, 5))
# self.risk_factors = ny.sample('health', dist.Binomial(5, .3))
self.hygiene = ny.sample('hygiene', dist.Beta(2, 5))
# self.worker_type = ny.sample('worker_type', dist.Categorical((.6, .1, .2, .1)))
self.epidemic_state = ny.sample(
'state',
dist.Binomial(1, initial_infections / pop_size[0] * pop_size[1]))
self.social_radius = ny.sample('radius', dist.Binomial(10, .2))
self.base_isolation = ny.sample('base_isolation', dist.Beta(2, 2))
# TODO: make these depend on risk factors as well
self.will_be_hospitalized = ny.sample(
'hosp', dist.Binomial(1, self.params.HOSP_AGE_MAP[self.age]))
self.will_die = ny.sample(
'die', dist.Binomial(1, self.params.DEATH_MAP[self.age]))
# The lengths of the infection are handled on a per agent basis via scenarios, these are just placeholders.
self.date_infected = np.where(
self.epidemic_state > 0, np.zeros(shape=pop_size),
np.full(shape=pop_size, fill_value=np.inf))
self.date_contagious = np.where(
self.epidemic_state > 0, np.ceil(self.params.EXPOSED_PERIOD),
np.full(shape=pop_size, fill_value=np.inf))
self.date_symptomatic = np.full(shape=pop_size, fill_value=np.inf)
self.date_recovered = np.full(shape=pop_size, fill_value=np.inf)
self.date_hospitalized = np.full(shape=pop_size, fill_value=np.inf)
self.date_died = np.full(shape=pop_size, fill_value=np.inf)
def model(data):
p = numpyro.sample("p", dist.Beta(1.0, 1.0))
if with_logits:
logits = logit(p)
numpyro.sample(
"obs", dist.Binomial(data["n"], logits=logits), obs=data["x"]
)
else:
numpyro.sample("obs", dist.Binomial(data["n"], probs=p), obs=data["x"])
def partially_pooled_with_logit(at_bats: jnp.ndarray, hits: Optional[jnp.ndarray] = None) -> None:
loc = numpyro.sample("loc", dist.Normal(-1, 1))
scale = numpyro.sample("scale", dist.HalfCauchy(1))
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
alpha = numpyro.sample("alpha", dist.Normal(loc, scale))
numpyro.sample("obs", dist.Binomial(at_bats, logits=alpha), obs=hits)
def create_model(yT, yC, num_components):
# Cosntants
nC = yC.shape[0]
nT = yT.shape[0]
zC = jnp.isinf(yC).sum().item()
zT = jnp.isinf(yT).sum().item()
yT_finite = yT[jnp.isinf(yT) == False]
yC_finite = yC_finite = yC[jnp.isinf(yC) == False]
K = num_components
p = numpyro.sample('p', dist.Beta(.5, .5))
gammaC = numpyro.sample('gammaC', dist.Beta(1, 1))
gammaT = numpyro.sample('gammaT', dist.Beta(1, 1))
etaC = numpyro.sample('etaC', dist.Dirichlet(jnp.ones(K) / K))
etaT = numpyro.sample('etaT', dist.Dirichlet(jnp.ones(K) / K))
with numpyro.plate('mixutre_components', K):
nu = numpyro.sample('nu', dist.LogNormal(3.5, 0.5))
mu = numpyro.sample('mu', dist.Normal(0, 3))
sigma = numpyro.sample('sigma', dist.LogNormal(0, .5))
phi = numpyro.sample('phi', dist.Normal(0, 3))
gammaT_star = simulate_data.compute_gamma_T_star(gammaC, gammaT, p)
etaT_star = simulate_data.compute_eta_T_star(etaC, etaT, p, gammaC, gammaT,
gammaT_star)
with numpyro.plate('y_C', nC - zC):
numpyro.sample('finite_obs_C',
Mix(nu[None, :],
mu[None, :],
sigma[None, :],
phi[None, :],
etaC[None, :]), obs=yC_finite[:, None])
with numpyro.plate('y_T', nT - zT):
numpyro.sample('finite_obs_T',
Mix(nu[None, :],
mu[None, :],
sigma[None, :],
phi[None, :],
etaT_star[None, :]), obs=yT_finite[:, None])
numpyro.sample('N_C', dist.Binomial(nC, gammaC), obs=zC)
numpyro.sample('N_T', dist.Binomial(nT, gammaT_star), obs=zT)
def partially_pooled(at_bats: jnp.ndarray, hits: Optional[jnp.ndarray] = None) -> None:
m = numpyro.sample("m", dist.Uniform(0, 1))
kappa = numpyro.sample("kappa", dist.Pareto(1, 1.5))
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
phi = numpyro.sample("phi", dist.Beta(m * kappa, (1 - m) * kappa))
numpyro.sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def fully_pooled(at_bats, hits=None):
r"""
Number of hits in $K$ at bats for each player has a Binomial
distribution with a common probability of success, $\phi$.
:param (np.DeviceArray) at_bats: Number of at bats for each player.
:param (np.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
phi_prior = dist.Uniform(np.array([0.]), np.array([1.]))
phi = sample("phi", phi_prior)
return sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def test_beta_binomial_log_prob(total_count, shape):
concentration0 = onp.exp(onp.random.normal(size=shape))
concentration1 = onp.exp(onp.random.normal(size=shape))
value = np.arange(1 + total_count)
num_samples = 100000
probs = onp.random.beta(concentration1, concentration0, size=(num_samples,) + shape)
log_probs = dist.Binomial(total_count, probs).log_prob(value)
expected = logsumexp(log_probs, 0) - np.log(num_samples)
actual = dist.BetaBinomial(concentration1, concentration0, total_count).log_prob(value)
assert_allclose(actual, expected, rtol=0.02)
def glmm(dept, male, applications, admit):
v_mu = numpyro.sample('v_mu', dist.Normal(0, np.array([4., 1.])))
sigma = numpyro.sample('sigma', dist.HalfNormal(np.ones(2)))
L_Rho = numpyro.sample('L_Rho', dist.LKJCholesky(2))
scale_tril = sigma[..., np.newaxis] * L_Rho
# non-centered parameterization
num_dept = len(onp.unique(dept))
z = numpyro.sample('z', dist.Normal(np.zeros((num_dept, 2)), 1))
v = np.dot(scale_tril, z.T).T
logits = v_mu[0] + v[dept, 0] + (v_mu[1] + v[dept, 1]) * male
numpyro.sample('admit', dist.Binomial(applications, logits=logits), obs=admit)
def not_pooled(at_bats, hits=None):
r"""
Number of hits in $K$ at bats for each player has a Binomial
distribution with independent probability of success, $\phi_i$.
:param (np.DeviceArray) at_bats: Number of at bats for each player.
:param (np.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
num_players = at_bats.shape[0]
phi_prior = dist.Uniform(np.zeros((num_players,)),
np.ones((num_players,)))
phi = sample("phi", phi_prior)
return sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def glmm(dept, male, applications, admit):
v_mu = sample('v_mu', dist.Normal(0, np.array([4., 1.])))
sigma = sample('sigma', dist.HalfNormal(np.ones(2)))
L_Rho = sample('L_Rho', dist.LKJCholesky(2))
scale_tril = np.expand_dims(sigma, axis=-1) * L_Rho
# non-centered parameterization
num_dept = len(onp.unique(dept))
z = sample('z', dist.Normal(np.zeros((num_dept, 2)), 1))
v = np.squeeze(np.matmul(np.expand_dims(scale_tril, axis=-3), np.expand_dims(z, axis=-1)),
axis=-1)
logits = v_mu[..., :1] + v[..., dept, 0] + (v_mu[..., 1:] + v[..., dept, 1]) * male
sample('admit', dist.Binomial(applications, logits=logits), obs=admit)
def fully_pooled(at_bats, hits=None):
r"""
Number of hits in $K$ at bats for each player has a Binomial
distribution with a common probability of success, $\phi$.
:param (jnp.DeviceArray) at_bats: Number of at bats for each player.
:param (jnp.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
phi_prior = dist.Uniform(0, 1)
phi = numpyro.sample("phi", phi_prior)
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
return numpyro.sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def partially_pooled_with_logit(at_bats, hits=None):
r"""
Number of hits has a Binomial distribution with a logit link function.
The logits $\alpha$ for each player is normally distributed with the
mean and scale parameters sharing a common prior.
:param (jnp.DeviceArray) at_bats: Number of at bats for each player.
:param (jnp.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
loc = numpyro.sample("loc", dist.Normal(-1, 1))
scale = numpyro.sample("scale", dist.HalfCauchy(1))
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
alpha = numpyro.sample("alpha", dist.Normal(loc, scale))
return numpyro.sample("obs", dist.Binomial(at_bats, logits=alpha), obs=hits)
def partially_pooled_with_logit(at_bats, hits=None):
r"""
Number of hits has a Binomial distribution with a logit link function.
The logits $\alpha$ for each player is normally distributed with the
mean and scale parameters sharing a common prior.
:param (np.DeviceArray) at_bats: Number of at bats for each player.
:param (np.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
num_players = at_bats.shape[0]
loc = sample("loc", dist.Normal(np.array([-1.]), np.array([1.])))
scale = sample("scale", dist.HalfCauchy(np.array([1.])))
shape = np.shape(loc)[:np.ndim(loc) - 1] + (num_players,)
alpha = sample("alpha", dist.Normal(np.broadcast_to(loc, shape),
np.broadcast_to(scale, shape)))
return sample("obs", dist.Binomial(at_bats, logits=alpha), obs=hits)
def glmm(dept, male, applications, admit=None):
v_mu = numpyro.sample('v_mu', dist.Normal(0, jnp.array([4., 1.])))
sigma = numpyro.sample('sigma', dist.HalfNormal(jnp.ones(2)))
L_Rho = numpyro.sample('L_Rho', dist.LKJCholesky(2, concentration=2))
scale_tril = sigma[..., jnp.newaxis] * L_Rho
# non-centered parameterization
num_dept = len(np.unique(dept))
z = numpyro.sample('z', dist.Normal(jnp.zeros((num_dept, 2)), 1))
v = jnp.dot(scale_tril, z.T).T
logits = v_mu[0] + v[dept, 0] + (v_mu[1] + v[dept, 1]) * male
if admit is None:
# we use a Delta site to record probs for predictive distribution
probs = expit(logits)
numpyro.sample('probs', dist.Delta(probs), obs=probs)
numpyro.sample('admit', dist.Binomial(applications, logits=logits), obs=admit)
def partially_pooled(at_bats, hits=None):
r"""
Number of hits has a Binomial distribution with independent
probability of success, $\phi_i$. Each $\phi_i$ follows a Beta
distribution with concentration parameters $c_1$ and $c_2$, where
$c_1 = m * kappa$, $c_2 = (1 - m) * kappa$, $m ~ Uniform(0, 1)$,
and $kappa ~ Pareto(1, 1.5)$.
:param (jnp.DeviceArray) at_bats: Number of at bats for each player.
:param (jnp.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
m = numpyro.sample("m", dist.Uniform(0, 1))
kappa = numpyro.sample("kappa", dist.Pareto(1, 1.5))
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
phi_prior = dist.Beta(m * kappa, (1 - m) * kappa)
phi = numpyro.sample("phi", phi_prior)
return numpyro.sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def partially_pooled(at_bats, hits=None):
r"""
Number of hits has a Binomial distribution with independent
probability of success, $\phi_i$. Each $\phi_i$ follows a Beta
distribution with concentration parameters $c_1$ and $c_2$, where
$c_1 = m * kappa$, $c_2 = (1 - m) * kappa$, $m ~ Uniform(0, 1)$,
and $kappa ~ Pareto(1, 1.5)$.
:param (np.DeviceArray) at_bats: Number of at bats for each player.
:param (np.DeviceArray) hits: Number of hits for the given at bats.
:return: Number of hits predicted by the model.
"""
num_players = at_bats.shape[0]
m = sample("m", dist.Uniform(np.array([0.]), np.array([1.])))
kappa = sample("kappa", dist.Pareto(np.array([1.5])))
shape = np.shape(kappa)[:np.ndim(kappa) - 1] + (num_players,)
phi_prior = dist.Beta(np.broadcast_to(m * kappa, shape),
np.broadcast_to((1 - m) * kappa, shape))
phi = sample("phi", phi_prior)
return sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
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():
numpyro.sample("x", dist.Bernoulli(0.7))
y = numpyro.sample("y", dist.Binomial(10, 0.3))
numpyro.deterministic("y2", y**2)
def model():
numpyro.sample("x", dist.Bernoulli(0.7).expand([3]))
numpyro.sample("y", dist.Binomial(10, 0.3))
def not_pooled(at_bats: jnp.ndarray, hits: Optional[jnp.ndarray] = None) -> None:
num_players = at_bats.shape[0]
with numpyro.plate("num_players", num_players):
phi = numpyro.sample("phi", dist.Uniform(0, 1))
numpyro.sample("obs", dist.Binomial(at_bats, probs=phi), obs=hits)
def model_meta(prior_std, obs=None):
a = numpyro.sample("a", dist.Normal(0, prior_std))
numpyro.sample("obs", dist.Binomial(logits=a), obs=obs)
def model():
numpyro.sample("x",
dist.Bernoulli(0.7),
infer={"enumerate": "parallel"})
y = numpyro.sample("y", dist.Binomial(10, 0.3))
numpyro.deterministic("y2", y**2)
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)
