Update tests for new missing values approach

Author: brandonwillard

pymc3/tests/test_idata_conversion.py

@@ -5,6 +5,7 @@
 import pandas as pd
 import pytest
 
+from aesara.tensor.subtensor import AdvancedIncSubtensor
 from arviz import InferenceData
 from arviz.tests.helpers import check_multiple_attrs
 from numpy import ma
@@ -295,27 +296,23 @@ def test_ovewrite_model_coords_dims(self):
         assert np.all(idata2.constant_data.x.dim1.values == np.array(new_dim1))
         assert np.all(idata2.constant_data.x.dim2.values == np.array(["c1", "c2"]))
 
-    @pytest.mark.xfail(reason="Missing data not refactored for v4")
     def test_missing_data_model(self):
         # source pymc3/pymc3/tests/test_missing.py
         data = ma.masked_values([1, 2, -1, 4, -1], value=-1)
         model = pm.Model()
         with model:
             x = pm.Normal("x", 1, 1)
-            pm.Normal("y", x, 1, observed=data)
+            y = pm.Normal("y", x, 1, observed=data)
             inference_data = pm.sample(100, chains=2, return_inferencedata=True)
 
         # make sure that data is really missing
-        (y_missing,) = model.missing_values
-        # TODO: Test values aren't enabled anymore
-        assert y_missing.tag.test_value.shape == (2,)
+        assert isinstance(y.owner.op, AdvancedIncSubtensor)
 
         test_dict = {"posterior": ["x"], "observed_data": ["y"], "log_likelihood": ["y"]}
         fails = check_multiple_attrs(test_dict, inference_data)
         assert not fails
 
     @pytest.mark.xfail(reason="LKJCholeskyCov not refactored for v4")
-    @pytest.mark.xfail(reason="Missing data not refactored for v4")
     def test_mv_missing_data_model(self):
         data = ma.masked_values([[1, 2], [2, 2], [-1, 4], [2, -1], [-1, -1]], value=-1)
 
@@ -324,13 +321,11 @@ def test_mv_missing_data_model(self):
             mu = pm.Normal("mu", 0, 1, size=2)
             sd_dist = pm.HalfNormal.dist(1.0)
             chol, *_ = pm.LKJCholeskyCov("chol_cov", n=2, eta=1, sd_dist=sd_dist, compute_corr=True)
-            pm.MvNormal("y", mu=mu, chol=chol, observed=data)
+            y = pm.MvNormal("y", mu=mu, chol=chol, observed=data)
             inference_data = pm.sample(100, chains=2, return_inferencedata=True)
 
         # make sure that data is really missing
-        (y_missing,) = model.missing_values
-        # TODO: Test values aren't enabled anymore
-        assert y_missing.tag.test_value.shape == (4,)
+        assert isinstance(y.owner.op, AdvancedIncSubtensor)
 
         test_dict = {
             "posterior": ["mu", "chol_cov"],

pymc3/tests/test_missing.py

@@ -16,63 +16,45 @@
 import pandas as pd
 import pytest
 
+from aesara.tensor.subtensor import AdvancedIncSubtensor
 from numpy import array, ma
 
 from pymc3 import ImputationWarning, Model, Normal, sample, sample_prior_predictive
 
 
-# @pytest.mark.xfail(reason="Missing values not fully refactored")
-def test_missing():
-    data = ma.masked_values([1, 2, -1, 4, -1], value=-1)
-    with Model() as model:
-        x = Normal("x", 1, 1)
-        with pytest.warns(ImputationWarning):
-            Normal("y", x, 1, observed=data)
-
-    (y_missing,) = model.missing_values
-    assert y_missing.eval().shape == (2,)
-
-    # In v3, the log-likelihoods for these missing points are zero, and the
-    # missing data point values are the `Distribution`'s "default" values.
-    test_point = model.initial_point
-    model.logp(test_point)
+@pytest.mark.parametrize(
+    "data",
+    [ma.masked_values([1, 2, -1, 4, -1], value=-1), pd.DataFrame([1, 2, numpy.nan, 4, numpy.nan])],
+)
+def test_missing(data):
 
-    with model:
-        prior_trace = sample_prior_predictive()
-    assert {"x", "y"} <= set(prior_trace.keys())
-
-
-@pytest.mark.xfail(reason="Missing values not fully refactored")
-def test_missing_pandas():
-    data = pd.DataFrame([1, 2, numpy.nan, 4, numpy.nan])
     with Model() as model:
         x = Normal("x", 1, 1)
         with pytest.warns(ImputationWarning):
-            Normal("y", x, 1, observed=data)
+            y = Normal("y", x, 1, observed=data)
 
-    (y_missing,) = model.missing_values
-    assert y_missing.tag.test_value.shape == (2,)
+    assert isinstance(y.owner.op, AdvancedIncSubtensor)
 
-    model.logp(model.initial_point)
+    test_point = model.initial_point
+    assert not numpy.isnan(model.logp(test_point))
 
     with model:
         prior_trace = sample_prior_predictive()
     assert {"x", "y"} <= set(prior_trace.keys())
 
 
-@pytest.mark.xfail(reason="Missing values not fully refactored")
 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):
-            Normal("y", x * predictors, 1, observed=data)
+            y = Normal("y", x * predictors, 1, observed=data)
 
-    (y_missing,) = model.missing_values
-    assert y_missing.tag.test_value.shape == (2,)
+    assert isinstance(y.owner.op, AdvancedIncSubtensor)
 
-    model.logp(model.initial_point)
+    test_point = model.initial_point
+    assert not numpy.isnan(model.logp(test_point))
 
     with model:
         prior_trace = sample_prior_predictive()
@@ -93,9 +75,15 @@ def test_missing_dual_observations():
 
         prior_trace = sample_prior_predictive()
         assert {"beta1", "beta2", "theta", "o1", "o2"} <= set(prior_trace.keys())
+        # TODO: Assert something
         sample()
 
 
+@pytest.mark.skip(
+    reason="This doesn't make sense in v4, because there are no "
+    "explicit variables to sample.  The missing values are "
+    "implicit random variables."
+)
 def test_internal_missing_observations():
     with Model() as model:
         obs1 = ma.masked_values([1, 2, -1, 4, -1], value=-1)
@@ -107,4 +95,5 @@ def test_internal_missing_observations():
 
         prior_trace = sample_prior_predictive()
         assert {"theta1", "theta2"} <= set(prior_trace.keys())
+        # TODO: Assert something
         sample()

pymc3/tests/test_model.py

@@ -24,16 +24,17 @@
 import pytest
 import scipy.sparse as sps
 
+from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor
 from aesara.tensor.var import TensorConstant
 from numpy.testing import assert_almost_equal
 
 import pymc3 as pm
 
 from pymc3 import Deterministic, Potential
-from pymc3.blocking import RaveledVars
+from pymc3.blocking import DictToArrayBijection, RaveledVars
 from pymc3.distributions import Normal, logpt_sum, transforms
-from pymc3.model import ValueGradFunction
+from pymc3.model import Point, ValueGradFunction
 from pymc3.tests.helpers import SeededTest
 
 
@@ -201,20 +202,9 @@ def test_duplicate_vars():
 def test_empty_observed():
     data = pd.DataFrame(np.ones((2, 3)) / 3)
     data.values[:] = np.nan
-    with pm.Model(aesara_config={"compute_test_value": "raise"}):
+    with pm.Model():
         a = pm.Normal("a", observed=data)
-
-        assert isinstance(a.tag.observations.owner.op, AdvancedIncSubtensor)
-        # The masked observations are replaced by elements of the RV `a`,
-        # which means that they should all have the same sample test values
-        a_data = a.tag.observations.owner.inputs[1]
-        npt.assert_allclose(a.tag.test_value.flatten(), a_data.tag.test_value)
-
-        # Let's try this again with another distribution
-        b = pm.Gamma("b", alpha=1, beta=1, observed=data)
-        assert isinstance(b.tag.observations.owner.op, AdvancedIncSubtensor)
-        b_data = b.tag.observations.owner.inputs[1]
-        npt.assert_allclose(b.tag.test_value.flatten(), b_data.tag.test_value)
+        assert not hasattr(a.tag, "observations")
 
 
 class TestValueGradFunction(unittest.TestCase):
@@ -302,8 +292,8 @@ def test_edge_case(self):
         assert dlogp.size == 4
         npt.assert_allclose(dlogp, 0.0, atol=1e-5)
 
-    def test_tensor_type_conversion(self):
-        # case described in #3122
+    def test_missing_data(self):
+        # Originally from a case described in #3122
         X = np.random.binomial(1, 0.5, 10)
         X[0] = -1  # masked a single value
         X = np.ma.masked_values(X, value=-1)
@@ -312,9 +302,16 @@ def test_tensor_type_conversion(self):
             x2 = pm.Bernoulli("x2", x1, observed=X)
 
         gf = m.logp_dlogp_function()
+        gf._extra_are_set = True
 
-        # TODO: Assert something.
-        # assert m["x2_missing"].type == gf._extra_vars_shared["x2_missing"].type
+        m.default_rng.get_value(borrow=True).seed(102)
+
+        # The gradient should have random values as inputs, so its value should
+        # change every time we evaluate it at the same point
+        #
+        # TODO: We could probably use a better test than this.
+        res = [gf(DictToArrayBijection.map(Point(m.test_point, model=m))) for i in range(20)]
+        assert np.var(res) > 0.0
 
     def test_aesara_switch_broadcast_edge_cases_1(self):
         # Tests against two subtle issues related to a previous bug in Theano
@@ -474,7 +471,7 @@ def test_make_obs_var():
 
     # Ensure that the missing values are appropriately set to None
     for func_output in [dense_output, sparse_output]:
-        assert func_output.tag.missing_values is None
+        assert isinstance(func_output.owner.op, RandomVariable)
 
     # Ensure that the Aesara variable names are correctly set.
     # Note that the output for masked inputs do not have their names set
@@ -488,9 +485,7 @@ def test_make_obs_var():
 
     # Masked output is something weird. Just ensure it has missing values
     # self.assertIsInstance(masked_output, TensorConstant)
-    assert masked_output.tag.missing_values is not None
-
-    return None
+    assert isinstance(masked_output.owner.op, AdvancedIncSubtensor)
 
 
 def test_initial_point():