def test_single_multivariate_component_deterministic_weights(self, weights, component, size):
# This test needs seeding to avoid repetitions
rngs = [
aesara.shared(np.random.default_rng(seed))
for seed in self.get_random_state().randint(2**30, size=2)
]
mix = Mixture.dist(weights, component, size=size, rngs=rngs)
mix_eval = mix.eval()
# Test shape
# component shape is either (4, 2, 3), (2, 3)
# weights shape is either (4, 2) or (2,)
if size is not None:
expected_shape = size + (3,)
elif component.ndim == 3 or weights.ndim == 2:
expected_shape = (4, 3)
else:
expected_shape = (3,)
assert mix_eval.shape == expected_shape
# Test draws
totals = mix_eval.sum(-1)
expected_large_count = (weights == 1)[..., 1]
assert np.all((totals == 10_000) == expected_large_count)
repetitions = np.unique(mix_eval[..., 0]).size < totals.size
assert not repetitions
# Test logp
mix_logp_eval = logp(mix, mix_eval).eval()
expected_logp_shape = expected_shape[:-1]
assert mix_logp_eval.shape == expected_logp_shape
bcast_weights = np.broadcast_to(weights, (*expected_logp_shape, 2))
expected_logp = logp(component, mix_eval[..., None, :]).eval()[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_logp_shape)
assert np.allclose(mix_logp_eval, expected_logp)
def test_broadcasting_in_shape(self):
with Model() as model:
mu = Gamma("mu", 1.0, 1.0, shape=2)
comp_dists = Poisson.dist(mu, shape=2)
mix = Mixture("mix", w=np.ones(2) / 2, comp_dists=comp_dists, shape=(1000,))
prior = sample_prior_predictive(samples=self.n_samples, return_inferencedata=False)
assert prior["mix"].shape == (self.n_samples, 1000)
def test_mixture_dtype(self):
mix_dtype = Mixture.dist(
w=[0.5, 0.5],
comp_dists=[
Multinomial.dist(n=5, p=[0.5, 0.5]),
Multinomial.dist(n=5, p=[0.5, 0.5]),
],
).dtype
assert mix_dtype == "int64"
mix_dtype = Mixture.dist(
w=[0.5, 0.5],
comp_dists=[
Dirichlet.dist(a=[0.5, 0.5]),
Dirichlet.dist(a=[0.5, 0.5]),
],
).dtype
assert mix_dtype == aesara.config.floatX
def test_change_size(self, comp_dists, expand):
univariate = comp_dists[0].owner.op.ndim_supp == 0
mix = Mixture.dist(w=Dirichlet.dist([1, 1]), comp_dists=comp_dists)
mix = Mixture.change_size(mix, new_size=(4,), expand=expand)
draws = mix.eval()
expected_shape = (4,) if univariate else (4, 3)
assert draws.shape == expected_shape
assert np.unique(draws).size == draws.size
mix = Mixture.dist(w=Dirichlet.dist([1, 1]), comp_dists=comp_dists, size=(3,))
mix = Mixture.change_size(mix, new_size=(5, 4), expand=expand)
draws = mix.eval()
expected_shape = (5, 4) if univariate else (5, 4, 3)
if expand:
expected_shape = expected_shape + (3,)
assert draws.shape == expected_shape
assert np.unique(draws).size == draws.size
def test_component_choice_random(self):
"""Test that mixture choices change over evaluations"""
with Model() as m:
weights = [0.5, 0.5]
components = [Normal.dist(-10, 0.01), Normal.dist(10, 0.01)]
mix = Mixture.dist(weights, components)
draws = draw(mix, draws=10)
# Probability of coming from same component 10 times is 0.5**10
assert np.unique(draws > 0).size == 2
def test_broadcast_components(self, comp_dists, expected_shape):
n_dists = len(comp_dists)
mix = Mixture.dist(w=np.ones(n_dists) / n_dists, comp_dists=comp_dists)
mix_eval = mix.eval()
assert tuple(mix_eval.shape) == expected_shape
assert np.unique(mix_eval).size == mix.eval().size
for comp_dist in mix.owner.inputs[2:]:
# We check that the input is a "pure" RandomVariable and not a broadcast
# operation. This confirms that all draws will be unique
assert isinstance(comp_dist.owner.op, RandomVariable)
assert tuple(comp_dist.shape.eval()) == expected_shape
def test_components_expanded_by_weights(self, comp_dists):
"""Test that components are expanded when size or weights are larger than components"""
univariate = comp_dists[0].owner.op.ndim_supp == 0
mix = Mixture.dist(
w=Dirichlet.dist([1, 1], shape=(3, 2)),
comp_dists=comp_dists,
size=(3,),
)
draws = mix.eval()
assert draws.shape == (3,) if univariate else (3, 3)
assert np.unique(draws).size == draws.size
mix = Mixture.dist(
w=Dirichlet.dist([1, 1], shape=(4, 3, 2)),
comp_dists=comp_dists,
size=(3,),
)
draws = mix.eval()
assert draws.shape == (4, 3) if univariate else (4, 3, 3)
assert np.unique(draws).size == draws.size
def test_iterable_single_component_warning(self):
with pytest.warns(None) as record:
Mixture.dist(w=[0.5, 0.5], comp_dists=Normal.dist(size=2))
Mixture.dist(w=[0.5, 0.5], comp_dists=[Normal.dist(size=2), Normal.dist(size=2)])
assert not record
with pytest.warns(UserWarning, match="Single component will be treated as a mixture"):
Mixture.dist(w=[0.5, 0.5], comp_dists=[Normal.dist(size=2)])
def test_single_poisson_predictive_sampling_shape(self):
# test the shape broadcasting in mixture random
rng = self.get_random_state()
y = np.concatenate([rng.poisson(5, size=10), rng.poisson(9, size=10)])
with Model() as model:
comp0 = Poisson.dist(mu=np.ones(2))
w0 = Dirichlet("w0", a=np.ones(2), shape=(2,))
like0 = Mixture("like0", w=w0, comp_dists=comp0, observed=y)
comp1 = Poisson.dist(mu=np.ones((20, 2)), shape=(20, 2))
w1 = Dirichlet("w1", a=np.ones(2), shape=(2,))
like1 = Mixture("like1", w=w1, comp_dists=comp1, observed=y)
comp2 = Poisson.dist(mu=np.ones(2))
w2 = Dirichlet("w2", a=np.ones(2), shape=(20, 2))
like2 = Mixture("like2", w=w2, comp_dists=comp2, observed=y)
comp3 = Poisson.dist(mu=np.ones(2), shape=(20, 2))
w3 = Dirichlet("w3", a=np.ones(2), shape=(20, 2))
like3 = Mixture("like3", w=w3, comp_dists=comp3, observed=y)
n_samples = 30
with model:
prior = sample_prior_predictive(samples=n_samples, return_inferencedata=False)
ppc = sample_posterior_predictive(
[self.get_inital_point(model)], samples=n_samples, return_inferencedata=False
)
assert prior["like0"].shape == (n_samples, 20)
assert prior["like1"].shape == (n_samples, 20)
assert prior["like2"].shape == (n_samples, 20)
assert prior["like3"].shape == (n_samples, 20)
assert ppc["like0"].shape == (n_samples, 20)
assert ppc["like1"].shape == (n_samples, 20)
assert ppc["like2"].shape == (n_samples, 20)
assert ppc["like3"].shape == (n_samples, 20)
def test_list_univariate_components_deterministic_weights(self, weights, components, size):
# Size can't be smaller than what is implied by replication dimensions
if size is not None and len(size) < max(components[0].ndim, weights.ndim - 1):
return
mix = Mixture.dist(weights, components, size=size)
mix_eval = mix.eval()
# Test shape
# components[0] shape is either (4, 3), (3,) or ()
# weights shape is either (3, 2) or (2,)
if size is not None:
expected_shape = size
elif components[0].ndim == 2:
expected_shape = (4, 3)
elif components[0].ndim == 1 or weights.ndim == 2:
expected_shape = (3,)
else:
expected_shape = ()
assert mix_eval.shape == expected_shape
# Test draws
expected_positive = np.zeros_like(mix_eval)
if expected_positive.ndim > 0:
expected_positive[..., :] = (weights == 1)[..., 1]
assert np.all((mix_eval > 0) == expected_positive)
repetitions = np.unique(mix_eval).size < mix_eval.size
assert not repetitions
# Test logp
mix_logp_eval = logp(mix, mix_eval).eval()
assert mix_logp_eval.shape == expected_shape
bcast_weights = np.broadcast_to(weights, (*expected_shape, 2))
expected_logp = np.stack(
(
logp(components[0], mix_eval).eval(),
logp(components[1], mix_eval).eval(),
),
axis=-1,
)[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_shape)
assert np.allclose(mix_logp_eval, expected_logp)
def test_list_poissons_sampling(self):
pois_w = np.array([0.4, 0.6])
pois_mu = np.array([5.0, 20.0])
pois_x = generate_poisson_mixture_data(pois_w, pois_mu, size=1000)
with Model() as model:
w = Dirichlet("w", floatX(np.ones_like(pois_w)), shape=pois_w.shape)
mu = Gamma("mu", 1.0, 1.0, shape=pois_w.size)
Mixture("x_obs", w, [Poisson.dist(mu[0]), Poisson.dist(mu[1])], observed=pois_x)
trace = sample(
5000,
chains=1,
step=Metropolis(),
random_seed=self.random_seed,
progressbar=False,
return_inferencedata=False,
)
assert_allclose(np.sort(trace["w"].mean(axis=0)), np.sort(pois_w), rtol=0.1, atol=0.1)
assert_allclose(np.sort(trace["mu"].mean(axis=0)), np.sort(pois_mu), rtol=0.1, atol=0.1)
def test_list_mvnormals_predictive_sampling_shape(self):
N = 100 # number of data points
K = 3 # number of mixture components
D = 3 # dimensionality of the data
X = MvNormal.dist(np.zeros(D), np.eye(D), size=N).eval()
with Model() as model:
pi = Dirichlet("pi", np.ones(K), shape=(K,))
comp_dist = []
mu = []
packed_chol = []
chol = []
for i in range(K):
mu.append(Normal(f"mu{i}", 0, 10, shape=D))
packed_chol.append(
LKJCholeskyCov(
f"chol_cov_{i}",
eta=2,
n=D,
sd_dist=HalfNormal.dist(2.5, size=D),
compute_corr=False,
)
)
chol.append(expand_packed_triangular(D, packed_chol[i], lower=True))
comp_dist.append(MvNormal.dist(mu=mu[i], chol=chol[i], shape=D))
Mixture("x_obs", pi, comp_dist, observed=X)
n_samples = 20
with model:
prior = sample_prior_predictive(samples=n_samples, return_inferencedata=False)
ppc = sample_posterior_predictive(
[self.get_inital_point(model)], samples=n_samples, return_inferencedata=False
)
assert ppc["x_obs"].shape == (n_samples,) + X.shape
assert prior["x_obs"].shape == (n_samples,) + X.shape
assert prior["mu0"].shape == (n_samples, D)
assert prior["chol_cov_0"].shape == (n_samples, D * (D + 1) // 2)
def test_list_multivariate_components_deterministic_weights(self, weights, components, size):
mix = Mixture.dist(weights, components, size=size)
mix_eval = mix.eval()
# Test shape
# components[0] shape is either (4, 3) or (3,)
# weights shape is either (4, 2) or (2,)
if size is not None:
expected_shape = size + (3,)
elif components[0].ndim == 2 or weights.ndim == 2:
expected_shape = (4, 3)
else:
expected_shape = (3,)
assert mix_eval.shape == expected_shape
# Test draws
expected_positive = np.zeros_like(mix_eval)
expected_positive[..., :] = (weights == 1)[..., 1, None]
assert np.all((mix_eval > 0) == expected_positive)
repetitions = np.unique(mix_eval).size < mix_eval.size
assert not repetitions
# Test logp
# MvNormal logp is currently limited to 2d values
expectation = pytest.raises(ValueError) if mix_eval.ndim > 2 else does_not_raise()
with expectation:
mix_logp_eval = logp(mix, mix_eval).eval()
assert mix_logp_eval.shape == expected_shape[:-1]
bcast_weights = np.broadcast_to(weights, (*expected_shape[:-1], 2))
expected_logp = np.stack(
(
logp(components[0], mix_eval).eval(),
logp(components[1], mix_eval).eval(),
),
axis=-1,
)[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_shape[:-1])
assert np.allclose(mix_logp_eval, expected_logp)
def test_nested_mixture(self):
if aesara.config.floatX == "float32":
rtol = 1e-4
else:
rtol = 1e-7
nbr = 4
with Model() as model:
# mixtures components
g_comp = Normal.dist(
mu=Exponential("mu_g", lam=1.0, shape=nbr, transform=None), sigma=1, shape=nbr
)
l_comp = LogNormal.dist(
mu=Exponential("mu_l", lam=1.0, shape=nbr, transform=None), sigma=1, shape=nbr
)
# weight vector for the mixtures
g_w = Dirichlet("g_w", a=floatX(np.ones(nbr) * 0.0000001), transform=None, shape=(nbr,))
l_w = Dirichlet("l_w", a=floatX(np.ones(nbr) * 0.0000001), transform=None, shape=(nbr,))
# mixture components
g_mix = Mixture.dist(w=g_w, comp_dists=g_comp)
l_mix = Mixture.dist(w=l_w, comp_dists=l_comp)
# mixture of mixtures
mix_w = Dirichlet("mix_w", a=floatX(np.ones(2)), transform=None, shape=(2,))
mix = Mixture("mix", w=mix_w, comp_dists=[g_mix, l_mix], observed=np.exp(self.norm_x))
test_point = model.compute_initial_point()
def mixmixlogp(value, point):
floatX = aesara.config.floatX
priorlogp = (
st.dirichlet.logpdf(
x=point["g_w"],
alpha=np.ones(nbr) * 0.0000001,
).astype(floatX)
+ st.expon.logpdf(x=point["mu_g"]).sum(dtype=floatX)
+ st.dirichlet.logpdf(
x=point["l_w"],
alpha=np.ones(nbr) * 0.0000001,
).astype(floatX)
+ st.expon.logpdf(x=point["mu_l"]).sum(dtype=floatX)
+ st.dirichlet.logpdf(
x=point["mix_w"],
alpha=np.ones(2),
).astype(floatX)
)
complogp1 = st.norm.logpdf(x=value, loc=point["mu_g"]).astype(floatX)
mixlogp1 = logsumexp(
np.log(point["g_w"]).astype(floatX) + complogp1, axis=-1, keepdims=True
)
complogp2 = st.lognorm.logpdf(value, 1.0, 0.0, np.exp(point["mu_l"])).astype(floatX)
mixlogp2 = logsumexp(
np.log(point["l_w"]).astype(floatX) + complogp2, axis=-1, keepdims=True
)
complogp_mix = np.concatenate((mixlogp1, mixlogp2), axis=1)
mixmixlogpg = logsumexp(
np.log(point["mix_w"]).astype(floatX) + complogp_mix, axis=-1, keepdims=False
)
return priorlogp, mixmixlogpg
value = np.exp(self.norm_x)[:, None]
priorlogp, mixmixlogpg = mixmixlogp(value, test_point)
# check logp of mixture
assert_allclose(mixmixlogpg, mix.logp_elemwise(test_point), rtol=rtol)
# check model logp
assert_allclose(priorlogp + mixmixlogpg.sum(), model.logp(test_point), rtol=rtol)
# check input and check logp again
test_point["g_w"] = np.asarray([0.1, 0.1, 0.2, 0.6])
test_point["mu_g"] = np.exp(np.random.randn(nbr))
priorlogp, mixmixlogpg = mixmixlogp(value, test_point)
assert_allclose(mixmixlogpg, mix.logp_elemwise(test_point), rtol=rtol)
assert_allclose(priorlogp + mixmixlogpg.sum(), model.logp(test_point), rtol=rtol)
def test_list_multivariate_components(self, weights, comp_dists, size, expected):
with Model() as model:
Mixture("x", weights, comp_dists, size=size)
assert_moment_is_expected(model, expected, check_finite_logp=False)