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 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)