def test_wrong_distribution_argument_in_list_fails():
with pytest.raises(TypeError, match=r"every element in 'distribution' "):
pm.Mixture(
"mix",
p=[0.5, 0.5],
distributions=[pm.Normal("comp1", loc=0.0, scale=1.0), "not a distribution"],
)
def _mixture_same_family(k, p, loc, scale, dat):
m = yield pm.Normal("means", loc=loc, scale=scale)
distribution = pm.Normal("d", loc=m, scale=scale)
obs = yield pm.Mixture(
"mixture", p=p, distributions=distribution, validate_args=True, observed=dat
)
return obs
def _mixture(k, p, loc, scale, dat):
m = yield pm.Normal("means", loc=loc, scale=scale)
distributions = [pm.Normal("d" + str(i), loc=m[..., i], scale=scale) for i in range(k)]
obs = yield pm.Mixture(
"mixture", p=p, distributions=distributions, validate_args=True, observed=dat
)
return obs
def _base_dist(self, psi: TensorLike, theta: TensorLike, *args, **kwargs):
return pm.Mixture(
p=[psi, 1.0 - psi],
distributions=[
pm.Constant(value=0, name="Zero"),
pm.Poisson(mu=theta, name="Poisson")
],
name="ZeroInflatedPoisson",
)._distribution
def _base_dist(self, psi: TensorLike, mu: TensorLike, alpha: TensorLike,
*args, **kwargs):
return pm.Mixture(
p=[psi, 1.0 - psi],
distributions=[
pm.Constant(value=0, name="Zero"),
pm.NegativeBinomial(mu=mu, alpha=alpha, name="NegBin"),
],
name="ZeroInflatedNegativeBinomial",
)._distribution
def _base_dist(self, psi: TensorLike, n: TensorLike, p: TensorLike, *args,
**kwargs):
"""
Zero-inflated binomial base distribution.
A ZeroInflatedBinomial is a mixture between a deterministic
distribution and a Binomial distribution.
"""
return pm.Mixture(
p=[psi, 1.0 - psi],
distributions=[
pm.Constant(0, name="Zero"),
pm.Binomial(n, p, name="Bin")
],
name="ZeroInflatedBinomial",
)._distribution
def test_wrong_distribution_argument_batched_fails():
with pytest.raises(TypeError, match=r"sequence of distributions"):
pm.Mixture("mix", p=[0.5, 0.5], distributions=tfd.Normal(0, 1))