def test_with_multinomial(self, batch_shape): p = np.random.uniform(size=(*batch_shape, self.mixture_comps, 3)) n = 100 * np.ones((*batch_shape, 1)) w = np.ones(self.mixture_comps) / self.mixture_comps mixture_axis = len(batch_shape) with Model() as model: comp_dists = Multinomial.dist(p=p, n=n, shape=(*batch_shape, self.mixture_comps, 3)) mixture = Mixture( "mixture", w=w, comp_dists=comp_dists, shape=(*batch_shape, 3), ) prior = sample_prior_predictive(samples=self.n_samples, return_inferencedata=False) assert prior["mixture"].shape == (self.n_samples, *batch_shape, 3) assert draw(mixture, draws=self.size).shape == (self.size, *batch_shape, 3) if aesara.config.floatX == "float32": rtol = 1e-4 else: rtol = 1e-7 initial_point = model.compute_initial_point() comp_logp = logp(comp_dists, initial_point["mixture"].reshape(*batch_shape, 1, 3)) log_sum_exp = logsumexp( comp_logp.eval() + np.log(w), axis=mixture_axis, keepdims=True ).sum() assert_allclose( model.compile_logp()(initial_point), log_sum_exp, rtol, )
def test_with_mvnormal(self): # 10 batch, 3-variate Gaussian mu = np.random.randn(self.mixture_comps, 3) mat = np.random.randn(3, 3) cov = mat @ mat.T chol = np.linalg.cholesky(cov) w = np.ones(self.mixture_comps) / self.mixture_comps with Model() as model: comp_dists = MvNormal.dist(mu=mu, chol=chol, shape=(self.mixture_comps, 3)) mixture = Mixture("mixture", w=w, comp_dists=comp_dists, shape=(3,)) prior = sample_prior_predictive(samples=self.n_samples, return_inferencedata=False) assert prior["mixture"].shape == (self.n_samples, 3) assert draw(mixture, draws=self.size).shape == (self.size, 3) if aesara.config.floatX == "float32": rtol = 1e-4 else: rtol = 1e-7 initial_point = model.compute_initial_point() comp_logp = logp(comp_dists, initial_point["mixture"].reshape(1, 3)) log_sum_exp = logsumexp(comp_logp.eval() + np.log(w), axis=0, keepdims=True).sum() assert_allclose( model.compile_logp()(initial_point), log_sum_exp, rtol, )
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 initialize_population(self) -> Dict[str, NDArray]: """Create an initial population from the prior distribution""" return sample_prior_predictive( self.draws, var_names=[v.name for v in self.model.unobserved_value_vars], model=self.model, return_inferencedata=False, )
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_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 initialize_population(self) -> Dict[str, np.ndarray]: """Create an initial population from the prior distribution""" sys.stdout.write(" ") # see issue #5828 with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="The effect of Potentials") result = sample_prior_predictive( self.draws, var_names=[v.name for v in self.model.unobserved_value_vars], model=self.model, return_inferencedata=False, ) return cast(Dict[str, np.ndarray], result)
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_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_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_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)