def test_list_normals_sampling(self):
norm_w = np.array([0.75, 0.25])
norm_mu = np.array([0.0, 5.0])
norm_sigma = np.ones_like(norm_mu)
norm_x = generate_normal_mixture_data(norm_w, norm_mu, norm_sigma, size=1000)
with Model() as model:
w = Dirichlet("w", floatX(np.ones_like(norm_w)), shape=norm_w.size)
mu = Normal("mu", 0.0, 10.0, shape=norm_w.size)
tau = Gamma("tau", 1.0, 1.0, shape=norm_w.size)
Mixture(
"x_obs",
w,
[Normal.dist(mu[0], tau=tau[0]), Normal.dist(mu[1], tau=tau[1])],
observed=norm_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(norm_w), rtol=0.1, atol=0.1)
assert_allclose(np.sort(trace["mu"].mean(axis=0)), np.sort(norm_mu), rtol=0.1, atol=0.1)
def __init__(self, mean=0, sigma=1, name="", model=None):
super().__init__(name, model)
self.register_rv(Normal.dist(mu=mean, sigma=sigma), "v1")
Normal("v2", mu=mean, sigma=sigma)
Normal("v3", mu=mean, sigma=Normal("sd", mu=10, sigma=1, initval=1.0))
Deterministic("v3_sq", self.v3**2)
Potential("p1", at.constant(1))
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_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_float64(self):
with Model() as model:
x = Normal("x", initval=np.array(1.0, dtype="float64"))
obs = Normal("obs", mu=x, sigma=1.0, observed=np.random.randn(5))
assert x.dtype == "float64"
assert obs.dtype == "float64"
for sampler in self.samplers:
with model:
sample(draws=10, tune=10, chains=1, step=sampler())
def test_missing_logp():
with Model() as m:
theta1 = Normal("theta1", 0, 5, observed=[0, 1, 2, 3, 4])
theta2 = Normal("theta2", mu=theta1, observed=[0, 1, 2, 3, 4])
m_logp = m.logp()
with Model() as m_missing:
theta1 = Normal("theta1", 0, 5, observed=np.array([0, 1, np.nan, 3, np.nan]))
theta2 = Normal("theta2", mu=theta1, observed=np.array([np.nan, np.nan, 2, np.nan, 4]))
m_missing_logp = m_missing.logp({"theta1_missing": [2, 4], "theta2_missing": [0, 1, 3]})
assert m_logp == m_missing_logp
def test_rv_size_is_none():
rv = Normal.dist(0, 1, size=None)
assert rv_size_is_none(rv.owner.inputs[1])
rv = Normal.dist(0, 1, size=1)
assert not rv_size_is_none(rv.owner.inputs[1])
size = Bernoulli.dist(0.5)
rv = Normal.dist(0, 1, size=size)
assert not rv_size_is_none(rv.owner.inputs[1])
size = Normal.dist(0, 1).size
rv = Normal.dist(0, 1, size=size)
assert not rv_size_is_none(rv.owner.inputs[1])
def test_missing(data):
with Model() as model:
x = Normal("x", 1, 1)
with pytest.warns(ImputationWarning):
_ = Normal("y", x, 1, observed=data)
assert "y_missing" in model.named_vars
test_point = model.compute_initial_point()
assert not np.isnan(model.compile_logp()(test_point))
with model:
prior_trace = sample_prior_predictive(return_inferencedata=False)
assert {"x", "y"} <= set(prior_trace.keys())
def test_normal_mixture_nd(self, nd, ncomp):
nd = to_tuple(nd)
ncomp = int(ncomp)
comp_shape = nd + (ncomp,)
test_mus = np.random.randn(*comp_shape)
test_taus = np.random.gamma(1, 1, size=comp_shape)
observed = generate_normal_mixture_data(
w=np.ones(ncomp) / ncomp, mu=test_mus, sigma=1 / np.sqrt(test_taus), size=10
)
with Model() as model0:
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,))
mixture0 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd, comp_shape=comp_shape)
obs0 = NormalMixture(
"obs", w=ws, mu=mus, tau=taus, comp_shape=comp_shape, observed=observed
)
with Model() as model1:
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,))
comp_dist = [
Normal.dist(mu=mus[..., i], tau=taus[..., i], shape=nd) for i in range(ncomp)
]
mixture1 = Mixture("m", w=ws, comp_dists=comp_dist, shape=nd)
obs1 = Mixture("obs", w=ws, comp_dists=comp_dist, observed=observed)
with Model() as model2:
# Test that results are correct without comp_shape being passed to the Mixture.
# This used to fail in V3
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,))
mixture2 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd)
obs2 = NormalMixture("obs", w=ws, mu=mus, tau=taus, observed=observed)
testpoint = model0.compute_initial_point()
testpoint["mus"] = test_mus
testpoint["taus_log__"] = np.log(test_taus)
for logp0, logp1, logp2 in zip(
model0.compile_logp(vars=[mixture0, obs0], sum=False)(testpoint),
model1.compile_logp(vars=[mixture1, obs1], sum=False)(testpoint),
model2.compile_logp(vars=[mixture2, obs2], sum=False)(testpoint),
):
assert_allclose(logp0, logp1)
assert_allclose(logp0, logp2)
def test_scalar_components(self):
nd = 3
npop = 4
# [[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]]
mus = at.constant(np.full((nd, npop), np.arange(npop)))
with Model(rng_seeder=self.get_random_state()) as model:
m = NormalMixture(
"m",
w=np.ones(npop) / npop,
mu=mus,
sigma=1e-5,
comp_shape=(nd, npop),
shape=nd,
)
z = Categorical("z", p=np.ones(npop) / npop, shape=nd)
mu = at.as_tensor_variable([mus[i, z[i]] for i in range(nd)])
latent_m = Normal("latent_m", mu=mu, sigma=1e-5, shape=nd)
size = 100
m_val = draw(m, draws=size)
latent_m_val = draw(latent_m, draws=size)
assert m_val.shape == latent_m_val.shape
# Test that each element in axis = -1 can come from independent
# components
assert not all(np.all(np.diff(m_val) < 1e-3, axis=-1))
assert not all(np.all(np.diff(latent_m_val) < 1e-3, axis=-1))
self.samples_from_same_distribution(m_val, latent_m_val)
# Check that logp is the same whether elements of the last axis are mixed or not
logp_fn = model.compile_logp(vars=[m])
assert np.isclose(logp_fn({"m": [0, 0, 0]}), logp_fn({"m": [0, 1, 2]}))
self.logp_matches(m, latent_m, z, npop, model=model)
def test_vector_components(self):
nd = 3
npop = 4
# [[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]]
mus = at.constant(np.full((nd, npop), np.arange(npop)))
with Model(rng_seeder=self.get_random_state()) as model:
m = Mixture(
"m",
w=np.ones(npop) / npop,
# MvNormal distribution with squared sigma diagonal covariance should
# be equal to vector of Normals from latent_m
comp_dists=[MvNormal.dist(mus[:, i], np.eye(nd) * 1e-5**2) for i in range(npop)],
)
z = Categorical("z", p=np.ones(npop) / npop)
latent_m = Normal("latent_m", mu=mus[..., z], sigma=1e-5, shape=nd)
size = 100
m_val = draw(m, draws=size)
latent_m_val = draw(latent_m, draws=size)
assert m_val.shape == latent_m_val.shape
# Test that each element in axis = -1 comes from the same mixture
# component
assert np.all(np.diff(m_val) < 1e-3)
assert np.all(np.diff(latent_m_val) < 1e-3)
# TODO: The following statistical test appears to be more flaky than expected
# even though the distributions should be the same. Seeding should make it
# stable but might be worth investigating further
self.samples_from_same_distribution(m_val, latent_m_val)
# Check that mixing of values in the last axis leads to smaller logp
logp_fn = model.compile_logp(vars=[m])
assert logp_fn({"m": [0, 0, 0]}) > logp_fn({"m": [0, 1, 0]}) > logp_fn({"m": [0, 1, 2]})
self.logp_matches(m, latent_m, z, npop, model=model)
def test_missing_dual_observations():
with Model() as model:
obs1 = ma.masked_values([1, 2, -1, 4, -1], value=-1)
obs2 = ma.masked_values([-1, -1, 6, -1, 8], value=-1)
beta1 = Normal("beta1", 1, 1)
beta2 = Normal("beta2", 2, 1)
latent = Normal("theta", size=5)
with pytest.warns(ImputationWarning):
ovar1 = Normal("o1", mu=beta1 * latent, observed=obs1)
with pytest.warns(ImputationWarning):
ovar2 = Normal("o2", mu=beta2 * latent, observed=obs2)
prior_trace = sample_prior_predictive(return_inferencedata=False)
assert {"beta1", "beta2", "theta", "o1", "o2"} <= set(prior_trace.keys())
# TODO: Assert something
trace = sample(chains=1, draws=50)
def test_missing_vector_parameter():
with Model() as m:
x = Normal(
"x",
np.array([-10, 10]),
0.1,
observed=np.array([[np.nan, 10], [-10, np.nan], [np.nan, np.nan]]),
)
x_draws = x.eval()
assert x_draws.shape == (3, 2)
assert np.all(x_draws[:, 0] < 0)
assert np.all(x_draws[:, 1] > 0)
assert np.isclose(
m.logp({"x_missing": np.array([-10, 10, -10, 10])}),
scipy.stats.norm(scale=0.1).logpdf(0) * 6,
)
def test_missing_with_predictors():
predictors = array([0.5, 1, 0.5, 2, 0.3])
data = ma.masked_values([1, 2, -1, 4, -1], value=-1)
with Model() as model:
x = Normal("x", 1, 1)
with pytest.warns(ImputationWarning):
y = Normal("y", x * predictors, 1, observed=data)
assert "y_missing" in model.named_vars
test_point = model.compute_initial_point()
assert not np.isnan(model.compile_logp()(test_point))
with model:
prior_trace = sample_prior_predictive(return_inferencedata=False)
assert {"x", "y"} <= set(prior_trace.keys())
def test_float32_MLDA(self):
data = np.random.randn(5).astype("float32")
with Model() as coarse_model:
x = Normal("x", initval=np.array(1.0, dtype="float32"))
obs = Normal("obs", mu=x, sigma=1.0, observed=data + 0.5)
with Model() as model:
x = Normal("x", initval=np.array(1.0, dtype="float32"))
obs = Normal("obs", mu=x, sigma=1.0, observed=data)
assert x.dtype == "float32"
assert obs.dtype == "float32"
with model:
sample(draws=10,
tune=10,
chains=1,
step=MLDA(coarse_models=[coarse_model]))
def test_interval_missing_observations():
with Model() as model:
obs1 = ma.masked_values([1, 2, -1, 4, -1], value=-1)
obs2 = ma.masked_values([-1, -1, 6, -1, 8], value=-1)
rng = aesara.shared(np.random.RandomState(2323), borrow=True)
with pytest.warns(ImputationWarning):
theta1 = Uniform("theta1", 0, 5, observed=obs1, rng=rng)
with pytest.warns(ImputationWarning):
theta2 = Normal("theta2", mu=theta1, observed=obs2, rng=rng)
assert "theta1_observed" in model.named_vars
assert "theta1_missing_interval__" in model.named_vars
assert not hasattr(
model.rvs_to_values[model.named_vars["theta1_observed"]].tag, "transform"
)
prior_trace = sample_prior_predictive(return_inferencedata=False)
# Make sure the observed + missing combined deterministics have the
# same shape as the original observations vectors
assert prior_trace["theta1"].shape[-1] == obs1.shape[0]
assert prior_trace["theta2"].shape[-1] == obs2.shape[0]
# Make sure that the observed values are newly generated samples
assert np.all(np.var(prior_trace["theta1_observed"], 0) > 0.0)
assert np.all(np.var(prior_trace["theta2_observed"], 0) > 0.0)
# Make sure the missing parts of the combined deterministic matches the
# sampled missing and observed variable values
assert np.mean(prior_trace["theta1"][:, obs1.mask] - prior_trace["theta1_missing"]) == 0.0
assert np.mean(prior_trace["theta1"][:, ~obs1.mask] - prior_trace["theta1_observed"]) == 0.0
assert np.mean(prior_trace["theta2"][:, obs2.mask] - prior_trace["theta2_missing"]) == 0.0
assert np.mean(prior_trace["theta2"][:, ~obs2.mask] - prior_trace["theta2_observed"]) == 0.0
assert {"theta1", "theta2"} <= set(prior_trace.keys())
trace = sample(
chains=1, draws=50, compute_convergence_checks=False, return_inferencedata=False
)
assert np.all(0 < trace["theta1_missing"].mean(0))
assert np.all(0 < trace["theta2_missing"].mean(0))
assert "theta1" not in trace.varnames
assert "theta2" not in trace.varnames
# Make sure that the observed values are newly generated samples and that
# the observed and deterministic matche
pp_trace = sample_posterior_predictive(trace, return_inferencedata=False, keep_size=False)
assert np.all(np.var(pp_trace["theta1"], 0) > 0.0)
assert np.all(np.var(pp_trace["theta2"], 0) > 0.0)
assert np.mean(pp_trace["theta1"][:, ~obs1.mask] - pp_trace["theta1_observed"]) == 0.0
assert np.mean(pp_trace["theta2"][:, ~obs2.mask] - pp_trace["theta2_observed"]) == 0.0
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_observed_rv_fail(self):
with pytest.raises(TypeError):
with pm.Model():
x = Normal("x")
Normal("n", observed=x)
def test_normal_moment(mu, sigma, size, expected):
with Model() as model:
Normal("x", mu=mu, sigma=sigma, size=size)
assert_moment_is_expected(model, expected)
def setup_class(self):
# True parameter values
alpha, sigma = 1, 1
beta = [1, 2.5]
# Size of dataset
size = 100
# Predictor variable
X = np.random.normal(size=(size, 2)).dot(np.array([[1, 0], [0, 0.2]]))
# Simulate outcome variable
Y = alpha + X.dot(beta) + np.random.randn(size) * sigma
with Model() as self.model:
# TODO: some variables commented out here as they're not working properly
# in v4 yet (9-jul-2021), so doesn't make sense to test str/latex for them
# Priors for unknown model parameters
alpha = Normal("alpha", mu=0, sigma=10)
b = Normal("beta", mu=0, sigma=10, size=(2, ), observed=beta)
sigma = HalfNormal("sigma", sigma=1)
# Test Cholesky parameterization
Z = MvNormal("Z", mu=np.zeros(2), chol=np.eye(2), size=(2, ))
# NegativeBinomial representations to test issue 4186
# nb1 = pm.NegativeBinomial(
# "nb_with_mu_alpha", mu=pm.Normal("nbmu"), alpha=pm.Gamma("nbalpha", mu=6, sigma=1)
# )
nb2 = NegativeBinomial("nb_with_p_n", p=Uniform("nbp"), n=10)
# Symbolic distribution
zip = ZeroInflatedPoisson("zip", 0.5, 5)
# Expected value of outcome
mu = Deterministic("mu", floatX(alpha + dot(X, b)))
# add a bounded variable as well
# bound_var = Bound(Normal, lower=1.0)("bound_var", mu=0, sigma=10)
# KroneckerNormal
n, m = 3, 4
covs = [np.eye(n), np.eye(m)]
kron_normal = KroneckerNormal("kron_normal",
mu=np.zeros(n * m),
covs=covs,
size=n * m)
# MatrixNormal
# matrix_normal = MatrixNormal(
# "mat_normal",
# mu=np.random.normal(size=n),
# rowcov=np.eye(n),
# colchol=np.linalg.cholesky(np.eye(n)),
# size=(n, n),
# )
# DirichletMultinomial
dm = DirichletMultinomial("dm", n=5, a=[1, 1, 1], size=(2, 3))
# Likelihood (sampling distribution) of observations
Y_obs = Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)
# add a potential as well
pot = Potential("pot", mu**2)
self.distributions = [alpha, sigma, mu, b, Z, nb2, zip, Y_obs, pot]
self.deterministics_or_potentials = [mu, pot]
# tuples of (formatting, include_params
self.formats = [("plain", True), ("plain", False), ("latex", True),
("latex", False)]
self.expected = {
("plain", True): [
r"alpha ~ N(0, 10)",
r"sigma ~ N**+(0, 1)",
r"mu ~ Deterministic(f(beta, alpha))",
r"beta ~ N(0, 10)",
r"Z ~ N(f(), f())",
r"nb_with_p_n ~ NB(10, nbp)",
r"zip ~ MarginalMixtureRV{inline=False}",
r"Y_obs ~ N(mu, sigma)",
r"pot ~ Potential(f(beta, alpha))",
],
("plain", False): [
r"alpha ~ N",
r"sigma ~ N**+",
r"mu ~ Deterministic",
r"beta ~ N",
r"Z ~ N",
r"nb_with_p_n ~ NB",
r"zip ~ MarginalMixtureRV{inline=False}",
r"Y_obs ~ N",
r"pot ~ Potential",
],
("latex", True): [
r"$\text{alpha} \sim \operatorname{N}(0,~10)$",
r"$\text{sigma} \sim \operatorname{N^{+}}(0,~1)$",
r"$\text{mu} \sim \operatorname{Deterministic}(f(\text{beta},~\text{alpha}))$",
r"$\text{beta} \sim \operatorname{N}(0,~10)$",
r"$\text{Z} \sim \operatorname{N}(f(),~f())$",
r"$\text{nb_with_p_n} \sim \operatorname{NB}(10,~\text{nbp})$",
r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
r"$\text{Y_obs} \sim \operatorname{N}(\text{mu},~\text{sigma})$",
r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
],
("latex", False): [
r"$\text{alpha} \sim \operatorname{N}$",
r"$\text{sigma} \sim \operatorname{N^{+}}$",
r"$\text{mu} \sim \operatorname{Deterministic}$",
r"$\text{beta} \sim \operatorname{N}$",
r"$\text{Z} \sim \operatorname{N}$",
r"$\text{nb_with_p_n} \sim \operatorname{NB}$",
r"$\text{zip} \sim \text{MarginalMixtureRV{inline=False}}$",
r"$\text{Y_obs} \sim \operatorname{N}$",
r"$\text{pot} \sim \operatorname{Potential}$",
],
}
assert np.allclose(mix_logp_eval, expected_logp)
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
@pytest.mark.parametrize(
"comp_dists",
(
[Normal.dist(size=(2,))],
[Normal.dist(), Normal.dist()],
[MvNormal.dist(np.ones(3), np.eye(3), size=(2,))],
[
MvNormal.dist(np.ones(3), np.eye(3)),
MvNormal.dist(np.ones(3), np.eye(3)),
],
),
)
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,
class TestMixtureMoments:
@pytest.mark.parametrize(
"weights, comp_dists, size, expected",
[
(
np.array([0.4, 0.6]),
Normal.dist(mu=np.array([-2, 6]), sigma=np.array([5, 3])),
None,
2.8,
),
(
np.tile(1 / 13, 13),
Normal.dist(-2, 1, size=(13,)),
(3,),
np.full((3,), -2),
),
(
np.array([0.4, 0.6]),
Normal.dist([-2, 6], 3),
(5, 3),
np.full((5, 3), 2.8),
),
(
np.broadcast_to(np.array([0.4, 0.6]), (5, 3, 2)),
Normal.dist(np.array([-2, 6]), np.array([5, 3])),
None,
np.full(shape=(5, 3), fill_value=2.8),
),
(
np.array([0.4, 0.6]),
Normal.dist(np.array([-2, 6]), np.array([5, 3]), size=(5, 3, 2)),
None,
np.full(shape=(5, 3), fill_value=2.8),
),
(
np.array([[0.8, 0.2], [0.2, 0.8]]),
Normal.dist(np.array([-2, 6])),
None,
np.array([-0.4, 4.4]),
),
# implied size = (11, 7) will be overwritten by (5, 3)
(
np.array([0.4, 0.6]),
Normal.dist(np.array([-2, 6]), np.array([5, 3]), size=(11, 7, 2)),
(5, 3),
np.full(shape=(5, 3), fill_value=2.8),
),
],
)
def test_single_univariate_component(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)
@pytest.mark.parametrize(
"weights, comp_dists, size, expected",
[
(
np.array([1, 0]),
[Normal.dist(-2, 5), Normal.dist(6, 3)],
None,
-2,
),
(
np.array([0.4, 0.6]),
[Normal.dist(-2, 5, size=(2,)), Normal.dist(6, 3, size=(2,))],
None,
np.full((2,), 2.8),
),
(
np.array([0.5, 0.5]),
[Normal.dist(-2, 5), Exponential.dist(lam=1 / 3)],
(3, 5),
np.full((3, 5), 0.5),
),
(
np.broadcast_to(np.array([0.4, 0.6]), (5, 3, 2)),
[Normal.dist(-2, 5), Normal.dist(6, 3)],
None,
np.full(shape=(5, 3), fill_value=2.8),
),
(
np.array([[0.8, 0.2], [0.2, 0.8]]),
[Normal.dist(-2, 5), Normal.dist(6, 3)],
None,
np.array([-0.4, 4.4]),
),
(
np.array([[0.8, 0.2], [0.2, 0.8]]),
[Normal.dist(-2, 5), Normal.dist(6, 3)],
(3, 2),
np.full((3, 2), np.array([-0.4, 4.4])),
),
(
# implied size = (11, 7) will be overwritten by (5, 3)
np.array([0.4, 0.6]),
[Normal.dist(-2, 5, size=(11, 7)), Normal.dist(6, 3, size=(11, 7))],
(5, 3),
np.full(shape=(5, 3), fill_value=2.8),
),
],
)
def test_list_univariate_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)
@pytest.mark.parametrize(
"weights, comp_dists, size, expected",
[
(
np.array([0.4, 0.6]),
MvNormal.dist(mu=np.array([[-1, -2], [3, 5]]), cov=np.eye(2) * 0.3),
None,
np.array([1.4, 2.2]),
),
(
np.array([0.5, 0.5]),
Dirichlet.dist(a=np.array([[0.0001, 0.0001, 1000], [2, 4, 6]])),
(4,),
np.array(np.full((4, 3), [1 / 12, 1 / 6, 3 / 4])),
),
(
np.array([0.4, 0.6]),
MvNormal.dist(mu=np.array([-10, 0, 10]), cov=np.eye(3) * 3, size=(4, 2)),
None,
np.full((4, 3), [-10, 0, 10]),
),
(
np.array([[1.0, 0], [0.0, 1.0]]),
MvNormal.dist(
mu=np.array([[-5, -10, -15], [5, 10, 15]]), cov=np.eye(3) * 3, size=(2,)
),
(3, 2),
np.full((3, 2, 3), [[-5, -10, -15], [5, 10, 15]]),
),
],
)
def test_single_multivariate_component(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)
@pytest.mark.parametrize(
"weights, comp_dists, size, expected",
[
(
np.array([0.4, 0.6]),
[
MvNormal.dist(mu=np.array([-1, -2]), cov=np.eye(2) * 0.3),
MvNormal.dist(mu=np.array([3, 5]), cov=np.eye(2) * 0.8),
],
None,
np.array([1.4, 2.2]),
),
(
np.array([0.4, 0.6]),
[
Dirichlet.dist(a=np.array([2, 3, 5])),
MvNormal.dist(mu=np.array([-10, 0, 10]), cov=np.eye(3) * 3),
],
(4,),
np.array(np.full((4, 3), [-5.92, 0.12, 6.2])),
),
(
np.array([0.4, 0.6]),
[
Dirichlet.dist(a=np.array([2, 3, 5]), size=(2,)),
MvNormal.dist(mu=np.array([-10, 0, 10]), cov=np.eye(3) * 3, size=(2,)),
],
None,
np.full((2, 3), [-5.92, 0.12, 6.2]),
),
(
np.array([[1.0, 0], [0.0, 1.0]]),
[
MvNormal.dist(mu=np.array([-5, -10, -15]), cov=np.eye(3) * 3, size=(2,)),
MvNormal.dist(mu=np.array([5, 10, 15]), cov=np.eye(3) * 3, size=(2,)),
],
(3, 2),
np.full((3, 2, 3), [[-5, -10, -15], [5, 10, 15]]),
),
],
)
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)
class TestMixture(SeededTest):
def get_inital_point(self, model):
"""Get initial point with untransformed variables for posterior predictive sampling"""
return {
var.name: initial_point
for var, initial_point in zip(
model.unobserved_value_vars,
model.compile_fn(model.unobserved_value_vars)(model.compute_initial_point()),
)
}
@pytest.mark.parametrize(
"weights",
[
np.array([1, 0]),
np.array([[1, 0], [0, 1], [1, 0]]),
],
)
@pytest.mark.parametrize(
"component",
[
Normal.dist([-10, 10]),
Normal.dist([-10, 10], size=(3, 2)),
Normal.dist([[-15, 15], [-10, 10], [-5, 5]], 1e-3),
Normal.dist([-10, 10], size=(4, 3, 2)),
],
)
@pytest.mark.parametrize("size", [None, (3,), (4, 3)])
def test_single_univariate_component_deterministic_weights(self, weights, component, size):
# Size can't be smaller than what is implied by replication dimensions
if size is not None and len(size) < max(component.ndim - 1, weights.ndim - 1):
return
mix = Mixture.dist(weights, component, size=size)
mix_eval = mix.eval()
# Test shape
# component shape is either (4, 3, 2), (3, 2) or (2,)
# weights shape is either (3, 2) or (2,)
if size is not None:
expected_shape = size
elif component.ndim == 3:
expected_shape = (4, 3)
elif component.ndim == 2 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 = logp(component, mix_eval[..., None]).eval()[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_shape)
assert np.allclose(mix_logp_eval, expected_logp)
@pytest.mark.parametrize(
"weights",
[
np.array([1, 0]),
np.array([[1, 0], [0, 1], [1, 0]]),
],
)
@pytest.mark.parametrize(
"components",
[
(Normal.dist(-10, 1e-3), Normal.dist(10, 1e-3)),
(Normal.dist(-10, 1e-3, size=(3,)), Normal.dist(10, 1e-3, size=(3,))),
(Normal.dist([-15, -10, -5], 1e-3), Normal.dist([15, 10, 5], 1e-3)),
(Normal.dist(-10, 1e-3, size=(4, 3)), Normal.dist(10, 1e-3, size=(4, 3))),
],
)
@pytest.mark.parametrize("size", [None, (3,), (4, 3)])
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)
@pytest.mark.parametrize(
"weights",
[
np.array([1, 0]),
np.array([[1, 0], [0, 1], [1, 0], [0, 1]]),
],
)
@pytest.mark.parametrize(
"component",
[
DirichletMultinomial.dist(n=[5_000, 10_000], a=np.ones((3,))),
DirichletMultinomial.dist(n=[5_000, 10_000], a=np.ones((3,)), size=(4, 2)),
],
)
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)