black

Author: AlexAndorra

pymc/tests/test_sampling.py

@@ -204,7 +204,7 @@ def test_return_inferencedata(self):
             assert len(result._groups_warmup) > 0
 
             # inferencedata without tuning, with idata_kwargs
-            prior = pm.sample_prior_predictive()
+            prior = pm.sample_prior_predictive(return_inferencedata=False)
             result = pm.sample(
                 **kwargs,
                 return_inferencedata=True,
@@ -472,7 +472,6 @@ def test_normal_scalar(self):
             trace = pm.sample(
                 draws=ndraws,
                 chains=nchains,
-                return_inferencedata=False,
             )
 
         with model:
@@ -486,43 +485,23 @@ def test_normal_scalar(self):
 
             # test keep_size parameter
             ppc = pm.sample_posterior_predictive(trace, keep_size=True)
-            assert ppc["a"].shape == (nchains, ndraws)
+            assert ppc.posterior_predictive["a"].shape == (1, nchains, ndraws)
 
             # test default case
             ppc = pm.sample_posterior_predictive(trace, var_names=["a"])
-            assert "a" in ppc
-            assert ppc["a"].shape == (nchains * ndraws,)
+            assert "a" in ppc.posterior_predictive.data_vars
+            assert ppc.posterior_predictive["a"].shape == (1, nchains * ndraws)
             # mu's standard deviation may have changed thanks to a's observed
-            _, pval = stats.kstest(ppc["a"] - trace["mu"], stats.norm(loc=0, scale=1).cdf)
+            _, pval = stats.kstest(
+                ppc.posterior_predictive["a"] - trace.posterior["mu"],
+                stats.norm(loc=0, scale=1).cdf,
+            )
             assert pval > 0.001
 
         # size argument not introduced to fast version [2019/08/20:rpg]
         with model:
             ppc = pm.sample_posterior_predictive(trace, size=5, var_names=["a"])
-            assert ppc["a"].shape == (nchains * ndraws, 5)
-
-    def test_normal_scalar_idata(self):
-        nchains = 2
-        ndraws = 500
-        with pm.Model() as model:
-            mu = pm.Normal("mu", 0.0, 1.0)
-            a = pm.Normal("a", mu=mu, sigma=1, observed=0.0)
-            trace = pm.sample(
-                draws=ndraws,
-                chains=nchains,
-                return_inferencedata=False,
-                discard_tuned_samples=False,
-            )
-
-        assert not isinstance(trace, InferenceData)
-
-        with model:
-            # test keep_size parameter and idata input
-            idata = pm.to_inference_data(trace)
-            assert isinstance(idata, InferenceData)
-
-            ppc = pm.sample_posterior_predictive(idata, keep_size=True)
-            assert ppc["a"].shape == (nchains, ndraws)
+            assert ppc.posterior_predictive["a"].shape == (1, nchains * ndraws, 5)
 
     def test_normal_vector(self, caplog):
         with pm.Model() as model:
@@ -532,25 +511,35 @@ def test_normal_vector(self, caplog):
 
         with model:
             # test list input
-            ppc0 = pm.sample_posterior_predictive([model.initial_point], samples=10)
-            ppc = pm.sample_posterior_predictive(trace, samples=12, var_names=[])
+            ppc0 = pm.sample_posterior_predictive(
+                [model.initial_point], return_inferencedata=False, samples=10
+            )
+            ppc = pm.sample_posterior_predictive(
+                trace, return_inferencedata=False, samples=12, var_names=[]
+            )
             assert len(ppc) == 0
 
             # test keep_size parameter
-            ppc = pm.sample_posterior_predictive(trace, keep_size=True)
+            ppc = pm.sample_posterior_predictive(trace, return_inferencedata=False, keep_size=True)
             assert ppc["a"].shape == (trace.nchains, len(trace), 2)
             with pytest.warns(UserWarning):
-                ppc = pm.sample_posterior_predictive(trace, samples=12, var_names=["a"])
+                ppc = pm.sample_posterior_predictive(
+                    trace, return_inferencedata=False, samples=12, var_names=["a"]
+                )
             assert "a" in ppc
             assert ppc["a"].shape == (12, 2)
 
             with pytest.warns(UserWarning):
-                ppc = pm.sample_posterior_predictive(trace, samples=12, var_names=["a"])
+                ppc = pm.sample_posterior_predictive(
+                    trace, return_inferencedata=False, samples=12, var_names=["a"]
+                )
             assert "a" in ppc
             assert ppc["a"].shape == (12, 2)
 
             # size unsupported by fast_ version  argument. [2019/08/19:rpg]
-            ppc = pm.sample_posterior_predictive(trace, samples=10, var_names=["a"], size=4)
+            ppc = pm.sample_posterior_predictive(
+                trace, return_inferencedata=False, samples=10, var_names=["a"], size=4
+            )
             assert "a" in ppc
             assert ppc["a"].shape == (10, 4, 2)
 
@@ -567,7 +556,7 @@ def test_normal_vector_idata(self, caplog):
             idata = pm.to_inference_data(trace)
             assert isinstance(idata, InferenceData)
 
-            ppc = pm.sample_posterior_predictive(idata, keep_size=True)
+            ppc = pm.sample_posterior_predictive(idata, return_inferencedata=False, keep_size=True)
             assert ppc["a"].shape == (trace.nchains, len(trace), 2)
 
     def test_exceptions(self, caplog):
@@ -600,11 +589,15 @@ def test_vector_observed(self):
             # TODO: Assert something about the output
             # ppc = pm.sample_posterior_predictive(idata, samples=12, var_names=[])
             # assert len(ppc) == 0
-            ppc = pm.sample_posterior_predictive(idata, samples=12, var_names=["a"])
+            ppc = pm.sample_posterior_predictive(
+                idata, return_inferencedata=False, samples=12, var_names=["a"]
+            )
             assert "a" in ppc
             assert ppc["a"].shape == (12, 2)
 
-            ppc = pm.sample_posterior_predictive(idata, samples=10, var_names=["a"], size=4)
+            ppc = pm.sample_posterior_predictive(
+                idata, return_inferencedata=False, samples=10, var_names=["a"], size=4
+            )
             assert "a" in ppc
             assert ppc["a"].shape == (10, 4, 2)
 
@@ -616,9 +609,13 @@ def test_sum_normal(self):
 
         with model:
             # test list input
-            ppc0 = pm.sample_posterior_predictive([model.initial_point], samples=10)
+            ppc0 = pm.sample_posterior_predictive(
+                [model.initial_point], return_inferencedata=False, samples=10
+            )
             assert ppc0 == {}
-            ppc = pm.sample_posterior_predictive(idata, samples=1000, var_names=["b"])
+            ppc = pm.sample_posterior_predictive(
+                idata, return_inferencedata=False, samples=1000, var_names=["b"]
+            )
             assert len(ppc) == 1
             assert ppc["b"].shape == (1000,)
             scale = np.sqrt(1 + 0.2 ** 2)
@@ -637,7 +634,7 @@ def test_model_not_drawable_prior(self):
             with pytest.raises(NotImplementedError) as excinfo:
                 pm.sample_prior_predictive(50)
             assert "Cannot sample" in str(excinfo.value)
-            samples = pm.sample_posterior_predictive(idata, 40)
+            samples = pm.sample_posterior_predictive(idata, 40, return_inferencedata=False)
             assert samples["foo"].shape == (40, 200)
 
     def test_model_shared_variable(self):
@@ -660,7 +657,7 @@ def test_model_shared_variable(self):
         samples = 100
         with model:
             post_pred = pm.sample_posterior_predictive(
-                trace, samples=samples, var_names=["p", "obs"]
+                trace, return_inferencedata=False, samples=samples, var_names=["p", "obs"]
             )
 
         expected_p = np.array([logistic.eval({coeff: val}) for val in trace["x"][:samples]])
@@ -694,6 +691,7 @@ def test_deterministic_of_observed(self):
             rtol = 1e-5 if aesara.config.floatX == "float64" else 1e-4
 
             ppc = pm.sample_posterior_predictive(
+                return_inferencedata=False,
                 model=model,
                 trace=trace,
                 samples=len(trace) * nchains,
@@ -728,6 +726,7 @@ def test_deterministic_of_observed_modified_interface(self):
                 trace, varnames=[n for n in trace.varnames if n != "out"]
             ).to_dict("records")
             ppc = pm.sample_posterior_predictive(
+                return_inferencedata=False,
                 model=model,
                 trace=ppc_trace,
                 samples=len(ppc_trace),
@@ -745,7 +744,7 @@ def test_variable_type(self):
             trace = pm.sample(compute_convergence_checks=False, return_inferencedata=False)
 
         with model:
-            ppc = pm.sample_posterior_predictive(trace, samples=1)
+            ppc = pm.sample_posterior_predictive(trace, return_inferencedata=False, samples=1)
             assert ppc["a"].dtype.kind == "f"
             assert ppc["b"].dtype.kind == "i"
 
@@ -918,7 +917,7 @@ def test_ignores_observed(self):
             positive_mu = pm.Deterministic("positive_mu", np.abs(mu))
             z = -1 - positive_mu
             pm.Normal("x_obs", mu=z, sigma=1, observed=observed_data)
-            prior = pm.sample_prior_predictive()
+            prior = pm.sample_prior_predictive(return_inferencedata=False)
 
         assert "observed_data" not in prior
         assert (prior["mu"] < -90).all()
@@ -932,8 +931,12 @@ def test_respects_shape(self):
             with pm.Model():
                 mu = pm.Gamma("mu", 3, 1, size=1)
                 goals = pm.Poisson("goals", mu, size=shape)
-                trace1 = pm.sample_prior_predictive(10, var_names=["mu", "mu", "goals"])
-                trace2 = pm.sample_prior_predictive(10, var_names=["mu", "goals"])
+                trace1 = pm.sample_prior_predictive(
+                    10, return_inferencedata=False, var_names=["mu", "mu", "goals"]
+                )
+                trace2 = pm.sample_prior_predictive(
+                    10, return_inferencedata=False, var_names=["mu", "goals"]
+                )
             if shape == 2:  # want to test shape as an int
                 shape = (2,)
             assert trace1["goals"].shape == (10,) + shape
@@ -944,7 +947,7 @@ def test_multivariate(self):
             m = pm.Multinomial("m", n=5, p=np.array([0.25, 0.25, 0.25, 0.25]))
             trace = pm.sample_prior_predictive(10)
 
-        assert trace["m"].shape == (10, 4)
+        assert trace.prior["m"].shape == (1, 10, 4)
 
     def test_multivariate2(self):
         # Added test for issue #3271
@@ -955,8 +958,12 @@ def test_multivariate2(self):
             burned_trace = pm.sample(
                 20, tune=10, cores=1, return_inferencedata=False, compute_convergence_checks=False
             )
-        sim_priors = pm.sample_prior_predictive(samples=20, model=dm_model)
-        sim_ppc = pm.sample_posterior_predictive(burned_trace, samples=20, model=dm_model)
+        sim_priors = pm.sample_prior_predictive(
+            return_inferencedata=False, samples=20, model=dm_model
+        )
+        sim_ppc = pm.sample_posterior_predictive(
+            burned_trace, return_inferencedata=False, samples=20, model=dm_model
+        )
         assert sim_priors["probs"].shape == (20, 6)
         assert sim_priors["obs"].shape == (20,) + mn_data.shape
         assert sim_ppc["obs"].shape == (20,) + mn_data.shape
@@ -987,9 +994,9 @@ def test_transformed(self):
             y = pm.Binomial("y", n=at_bats, p=thetas, observed=hits)
             gen = pm.sample_prior_predictive(draws)
 
-        assert gen["phi"].shape == (draws,)
-        assert gen["y"].shape == (draws, n)
-        assert "thetas" in gen
+        assert gen.prior["phi"].shape == (1, draws)
+        assert gen.prior_predictive["y"].shape == (1, draws, n)
+        assert "thetas" in gen.prior.data_vars
 
     def test_shared(self):
         n1 = 10
@@ -1002,16 +1009,16 @@ def test_shared(self):
             o = pm.Deterministic("o", obs)
             gen1 = pm.sample_prior_predictive(draws)
 
-        assert gen1["y"].shape == (draws, n1)
-        assert gen1["o"].shape == (draws, n1)
+        assert gen1.prior["y"].shape == (1, draws, n1)
+        assert gen1.prior["o"].shape == (1, draws, n1)
 
         n2 = 20
         obs.set_value(np.random.rand(n2) < 0.5)
         with m:
             gen2 = pm.sample_prior_predictive(draws)
 
-        assert gen2["y"].shape == (draws, n2)
-        assert gen2["o"].shape == (draws, n2)
+        assert gen2.prior["y"].shape == (1, draws, n2)
+        assert gen2.prior["o"].shape == (1, draws, n2)
 
     def test_density_dist(self):
         obs = np.random.normal(-1, 0.1, size=10)
@@ -1025,7 +1032,7 @@ def test_density_dist(self):
                 random=lambda mu, sd, rng=None, size=None: rng.normal(loc=mu, scale=sd, size=size),
                 observed=obs,
             )
-            prior = pm.sample_prior_predictive()
+            prior = pm.sample_prior_predictive(return_inferencedata=False)
 
         npt.assert_almost_equal(prior["a"].mean(), 0, decimal=1)
 
@@ -1035,17 +1042,17 @@ def test_shape_edgecase(self):
             sd = pm.Uniform("sd", lower=2, upper=3)
             x = pm.Normal("x", mu=mu, sigma=sd, size=5)
             prior = pm.sample_prior_predictive(10)
-        assert prior["mu"].shape == (10, 5)
+        assert prior.prior["mu"].shape == (1, 10, 5)
 
     def test_zeroinflatedpoisson(self):
         with pm.Model():
             theta = pm.Beta("theta", alpha=1, beta=1)
             psi = pm.HalfNormal("psi", sd=1)
             pm.ZeroInflatedPoisson("suppliers", psi=psi, theta=theta, size=20)
             gen_data = pm.sample_prior_predictive(samples=5000)
-            assert gen_data["theta"].shape == (5000,)
-            assert gen_data["psi"].shape == (5000,)
-            assert gen_data["suppliers"].shape == (5000, 20)
+            assert gen_data.prior["theta"].shape == (1, 5000)
+            assert gen_data.prior["psi"].shape == (1, 5000)
+            assert gen_data.prior["suppliers"].shape == (1, 5000, 20)
 
     def test_potentials_warning(self):
         warning_msg = "The effect of Potentials on other parameters is ignored during"
@@ -1075,10 +1082,10 @@ def ub_interval_forward(x, ub):
             )
 
         # Check values are correct
-        assert np.allclose(prior["ub_log__"], np.log(prior["ub"]))
+        assert np.allclose(prior.prior["ub_log__"].data, np.log(prior.prior["ub"].data))
         assert np.allclose(
-            prior["x_interval__"],
-            ub_interval_forward(prior["x"], prior["ub"]),
+            prior.prior["x_interval__"].data,
+            ub_interval_forward(prior.prior["x"].data, prior.prior["ub"].data),
         )
 
         # Check that it works when the original RVs are not mentioned in var_names
@@ -1090,9 +1097,16 @@ def ub_interval_forward(x, ub):
                 var_names=["ub_log__", "x_interval__"],
                 samples=10,
             )
-        assert "ub" not in prior_transformed_only and "x" not in prior_transformed_only
-        assert np.allclose(prior["ub_log__"], prior_transformed_only["ub_log__"])
-        assert np.allclose(prior["x_interval__"], prior_transformed_only["x_interval__"])
+        assert (
+            "ub" not in prior_transformed_only.prior.data_vars
+            and "x" not in prior_transformed_only.prior.data_vars
+        )
+        assert np.allclose(
+            prior.prior["ub_log__"].data, prior_transformed_only.prior["ub_log__"].data
+        )
+        assert np.allclose(
+            prior.prior["x_interval__"], prior_transformed_only.prior["x_interval__"].data
+        )
 
     def test_issue_4490(self):
         # Test that samples do not depend on var_name order or, more fundamentally,
@@ -1112,27 +1126,34 @@ def test_issue_4490(self):
             d = pm.Normal("d")
             prior2 = pm.sample_prior_predictive(samples=1, var_names=["b", "a", "d", "c"])
 
-        assert prior1["a"] == prior2["a"]
-        assert prior1["b"] == prior2["b"]
-        assert prior1["c"] == prior2["c"]
-        assert prior1["d"] == prior2["d"]
+        assert prior1.prior["a"] == prior2.prior["a"]
+        assert prior1.prior["b"] == prior2.prior["b"]
+        assert prior1.prior["c"] == prior2.prior["c"]
+        assert prior1.prior["d"] == prior2.prior["d"]
 
 
 class TestSamplePosteriorPredictive:
     def test_point_list_arg_bug_spp(self, point_list_arg_bug_fixture):
         pmodel, trace = point_list_arg_bug_fixture
         with pmodel:
-            pp = pm.sample_posterior_predictive([trace[15]], var_names=["d"])
+            pp = pm.sample_posterior_predictive(
+                [trace[15]], return_inferencedata=False, var_names=["d"]
+            )
 
     def test_sample_from_xarray_prior(self, point_list_arg_bug_fixture):
         pmodel, trace = point_list_arg_bug_fixture
 
         with pmodel:
-            prior = pm.sample_prior_predictive(samples=20)
+            prior = pm.sample_prior_predictive(
+                samples=20,
+                return_inferencedata=False,
+            )
             idat = pm.to_inference_data(trace, prior=prior)
 
         with pmodel:
-            pp = pm.sample_posterior_predictive(idat.prior, var_names=["d"])
+            pp = pm.sample_posterior_predictive(
+                idat.prior, return_inferencedata=False, var_names=["d"]
+            )
 
     def test_sample_from_xarray_posterior(self, point_list_arg_bug_fixture):
         pmodel, trace = point_list_arg_bug_fixture