Fix ICE plot when there is a discrete variable (#107)

* fix tata plot shape

* add test case for discrete variables

Author: Juan Orduz

pymc_bart/utils.py

@@ -289,7 +289,7 @@ def plot_ice(
                 p_di = func(p_d[:, :, s_i])
             if var in var_discrete:
                 axes[count].plot(new_x, p_di.mean(0), "o", color=color_mean)
-                axes[count].plot(new_x, p_di, ".", color=color, alpha=alpha)
+                axes[count].plot(new_x, p_di.T, ".", color=color, alpha=alpha)
             else:
                 if smooth:
                     x_data, y_data = _smooth_mean(new_x, p_di, "ice", smooth_kwargs)

tests/test_bart.py

@@ -160,6 +160,7 @@ def test_sample_posterior(self):
             {"instances": 2},
             {"var_idx": [0], "smooth": False, "color": "k"},
             {"grid": (1, 2), "sharey": "none", "alpha": 1},
+            {"var_discrete": [0]}
         ],
     )
     def test_ice(self, kwargs):
@@ -177,6 +178,7 @@ def test_ice(self, kwargs):
             },
             {"var_idx": [0], "smooth": False, "color": "k"},
             {"grid": (1, 2), "sharey": "none", "alpha": 1},
+            {"var_discrete": [0]}
         ],
     )
     def test_pdp(self, kwargs):