CLN Remove pandas from bart.utils.plot_dependence (#5558)

Author: thomasjpfan

pymc/bart/utils.py

@@ -1,7 +1,6 @@
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 
 from numpy.random import RandomState
 from scipy.interpolate import griddata
@@ -147,13 +146,13 @@ def plot_dependence(
 
     rng = RandomState(seed=random_seed)
 
-    if isinstance(X, pd.DataFrame):
+    if hasattr(X, "columns") and hasattr(X, "values"):
         X_names = list(X.columns)
         X = X.values
     else:
         X_names = []
 
-    if isinstance(Y, pd.DataFrame):
+    if hasattr(Y, "name"):
         Y_label = f"Predicted {Y.name}"
     else:
         Y_label = "Predicted Y"

pymc/tests/test_bart.py

@@ -2,7 +2,7 @@
 import pytest
 
 from numpy.random import RandomState
-from numpy.testing import assert_almost_equal
+from numpy.testing import assert_almost_equal, assert_array_equal
 
 import pymc as pm
 
@@ -103,6 +103,18 @@ def test_predict(self):
     def test_pdp(self, kwargs):
         pm.bart.utils.plot_dependence(self.idata, X=self.X, Y=self.Y, **kwargs)
 
+    def test_pdp_pandas_labels(self):
+        pd = pytest.importorskip("pandas")
+
+        X_names = ["norm1", "norm2", "binom"]
+        X_pd = pd.DataFrame(self.X, columns=X_names)
+        Y_pd = pd.Series(self.Y, name="response")
+        axes = pm.bart.utils.plot_dependence(self.idata, X=X_pd, Y=Y_pd)
+
+        figure = axes[0].figure
+        assert figure.texts[0].get_text() == "Predicted response"
+        assert_array_equal([ax.get_xlabel() for ax in axes], X_names)
+
 
 @pytest.mark.parametrize(
     "size, expected",