add tests

Author: aloctavodia

pymc/bart/utils.py

@@ -51,7 +51,7 @@ def predict(idata, rng, X_new=None, size=None):
 
 
 def plot_dependence(
-    bart_trees,
+    idata,
     X=None,
     Y=None,
     kind="pdp",
@@ -78,8 +78,8 @@ def plot_dependence(
 
     Parameters
     ----------
-    bart_trees: DataArray of trees
-        BART trees
+    idata: InferenceData
+        InferenceData containing a collection of BART_trees in sample_stats group
     X : array-like
         The covariate matrix.
     Y : array-like
@@ -144,6 +144,7 @@ def plot_dependence(
         )
 
     rng = RandomState(seed=random_seed)
+    bart_trees = idata.sample_stats.bart_trees
 
     if isinstance(X, pd.DataFrame):
         X_names = list(X.columns)

pymc/tests/test_bart.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 from numpy.random import RandomState
 from numpy.testing import assert_almost_equal
@@ -47,33 +48,47 @@ def test_bart_vi():
         assert_almost_equal(var_imp.sum(), 1)
 
 
-def test_bart_random():
+def test_missing_data():
     X = np.random.normal(0, 1, size=(2, 50)).T
     Y = np.random.normal(0, 1, size=50)
+    X[10:20, 0] = np.nan
 
     with pm.Model() as model:
         mu = pm.BART("mu", X, Y, m=10)
         sigma = pm.HalfNormal("sigma", 1)
         y = pm.Normal("y", mu, sigma, observed=Y)
-        idata = pm.sample(random_seed=3415, chains=1)
-
-    rng = RandomState(12345)
-    pred_all = pm.bart.utils.predict(idata, rng, size=2)
-    rng = RandomState(12345)
-    pred_first = pm.bart.utils.predict(idata, rng, X_new=X[:10])
-
-    assert_almost_equal(pred_first, pred_all[0, :10], decimal=4)
-    assert pred_all.shape == (2, 50)
-    assert pred_first.shape == (10,)
+        idata = pm.sample(random_seed=3415)
 
 
-def test_missing_data():
+def test_utils():
     X = np.random.normal(0, 1, size=(2, 50)).T
     Y = np.random.normal(0, 1, size=50)
-    X[10:20, 0] = np.nan
 
     with pm.Model() as model:
         mu = pm.BART("mu", X, Y, m=10)
         sigma = pm.HalfNormal("sigma", 1)
         y = pm.Normal("y", mu, sigma, observed=Y)
         idata = pm.sample(random_seed=3415)
+
+    def test_predict():
+        rng = RandomState(12345)
+        pred_all = pm.bart.utils.predict(idata, rng, size=2)
+        rng = RandomState(12345)
+        pred_first = pm.bart.utils.predict(idata, rng, X_new=X[:10])
+
+        assert_almost_equal(pred_first, pred_all[0, :10], decimal=4)
+        assert pred_all.shape == (2, 50)
+        assert pred_first.shape == (10,)
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            {},
+            {"kind": "pdp", "xs_interval": "quantiles", "xs_values": [0.25, 0.5, 0.75]},
+            {"kind": "ice", "instances": 5},
+            {"var_idx": 0, "rug": False, "smooth": False, "color": "k"},
+            {"grid": (1, 2), "sharey": "False", "alpha": 1},
+        ],
+    )
+    def test_pdp():
+        pm.bart.utils.plot_dependence(idata, X=None, Y=None, **kwargs)