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)