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 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_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 __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_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)
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)
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)
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,