Fix bug in plot_ice, and clean docstring of plot_ice and plot_pdp (#135)

* fix plot_pdp/ice

* fix test

* fix type hints

Author: aloctavodia

pymc_bart/utils.py

@@ -155,14 +155,12 @@ def plot_ice(
     bartrv: Variable,
     X: npt.NDArray[np.float_],
     Y: Optional[npt.NDArray[np.float_]] = None,
-    xs_interval: str = "quantiles",
-    xs_values: Optional[Union[int, List[float]]] = None,
     var_idx: Optional[List[int]] = None,
     var_discrete: Optional[List[int]] = None,
     func: Optional[Callable] = None,
     centered: Optional[bool] = True,
-    samples: int = 50,
-    instances: int = 10,
+    samples: int = 100,
+    instances: int = 30,
     random_seed: Optional[int] = None,
     sharey: bool = True,
     smooth: bool = True,
@@ -185,16 +183,6 @@ def plot_ice(
         The covariate matrix.
     Y : Optional[npt.NDArray[np.float_]], by default None.
         The response vector.
-    xs_interval : str
-        Method used to compute the values X used to evaluate the predicted function. "linear",
-        evenly spaced values in the range of X. "quantiles", the evaluation is done at the specified
-        quantiles of X. "insample", the evaluation is done at the values of X.
-        For discrete variables these options are ommited.
-    xs_values : Optional[Union[int, List[float]]], by default None.
-        Values of X used to evaluate the predicted function. If ``xs_interval="linear"`` number of
-        points in the evenly spaced grid. If ``xs_interval="quantiles"``quantile or sequence of
-        quantiles to compute, which must be between 0 and 1 inclusive.
-        Ignored when ``xs_interval="insample"``.
     var_idx : Optional[List[int]], by default None.
         List of the indices of the covariate for which to compute the pdp or ice.
     var_discrete : Optional[List[int]], by default None.
@@ -205,22 +193,20 @@ def plot_ice(
         If True the result is centered around the partial response evaluated at the lowest value in
         ``xs_interval``. Defaults to True.
     samples : int
-        Number of posterior samples used in the predictions. Defaults to 50
+        Number of posterior samples used in the predictions. Defaults to 100
     instances : int
-        Number of instances of X to plot. Defaults to 10.
+        Number of instances of X to plot. Defaults to 30.
     random_seed : Optional[int], by default None.
         Seed used to sample from the posterior. Defaults to None.
     sharey : bool
         Controls sharing of properties among y-axes. Defaults to True.
-    rug : bool
-        Whether to include a rugplot. Defaults to True.
     smooth : bool
         If True the result will be smoothed by first computing a linear interpolation of the data
         over a regular grid and then applying the Savitzky-Golay filter to the interpolated data.
         Defaults to True.
     grid : str or tuple
         How to arrange the subplots. Defaults to "long", one subplot below the other.
-        Other options are "wide", one subplot next to eachother or a tuple indicating the number of
+        Other options are "wide", one subplot next to each other or a tuple indicating the number of
         rows and columns.
     color : matplotlib valid color
         Color used to plot the pdp or ice. Defaults to "C0"
@@ -257,17 +243,17 @@ def identity(x):
         indices,
         var_idx,
         var_discrete,
-        xs_interval,
-        xs_values,
-    ) = _prepare_plot_data(X, Y, xs_interval, xs_values, var_idx, var_discrete)
+        _,
+        _,
+    ) = _prepare_plot_data(X, Y, "linear", None, var_idx, var_discrete)
 
     fig, axes, shape = _get_axes(bartrv, var_idx, grid, sharey, figsize, ax)
 
     instances_ary = rng.choice(range(X.shape[0]), replace=False, size=instances)
     idx_s = list(range(X.shape[0]))
 
     count = 0
-    for var in range(len(var_idx)):
+    for i_var, var in enumerate(var_idx):
         indices_mi = indices[:]
         indices_mi.remove(var)
         y_pred = []
@@ -283,6 +269,7 @@ def identity(x):
 
         new_x = fake_X[:, var]
         p_d = np.array(y_pred)
+        print(p_d.shape)
 
         for s_i in range(shape):
             if centered:
@@ -301,7 +288,7 @@ def identity(x):
                     idx = np.argsort(new_x)
                     axes[count].plot(new_x[idx], p_di.mean(0)[idx], color=color_mean)
                     axes[count].plot(new_x[idx], p_di.T[idx], color=color, alpha=alpha)
-            axes[count].set_xlabel(x_labels[var])
+            axes[count].set_xlabel(x_labels[i_var])
 
             count += 1
 
@@ -349,7 +336,7 @@ def plot_pdp(
         For discrete variables these options are ommited.
     xs_values : Optional[Union[int, List[float]]], by default None.
         Values of X used to evaluate the predicted function. If ``xs_interval="linear"`` number of
-        points in the evenly spaced grid. If ``xs_interval="quantiles"``quantile or sequence of
+        points in the evenly spaced grid. If ``xs_interval="quantiles"`` quantile or sequence of
         quantiles to compute, which must be between 0 and 1 inclusive.
         Ignored when ``xs_interval="insample"``.
     var_idx : Optional[List[int]], by default None.
@@ -717,7 +704,8 @@ def plot_variable_importance(
     xlabel_angle: float = 0,
     samples: int = 100,
     random_seed: Optional[int] = None,
-) -> Tuple[List[int], List[plt.Axes]]:
+    ax: Optional[plt.Axes] = None,
+) -> Tuple[List[int], Union[List[plt.Axes], Any]]:
     """
     Estimates variable importance from the BART-posterior.
 
@@ -747,6 +735,8 @@ def plot_variable_importance(
         Number of predictions used to compute correlation for subsets of variables. Defaults to 100
     random_seed : Optional[int]
         random_seed used to sample from the posterior. Defaults to None.
+    ax : axes
+        Matplotlib axes.
 
     Returns
     -------
@@ -771,7 +761,8 @@ def plot_variable_importance(
     if figsize is None:
         figsize = (8, 3)
 
-    _, ax = plt.subplots(1, 1, figsize=figsize)
+    if ax is None:
+        _, ax = plt.subplots(1, 1, figsize=figsize)
 
     if labels is None:
         labels_ary = np.arange(n_vars).astype(str)

tests/test_bart.py

@@ -153,8 +153,6 @@ def test_sample_posterior(self):
             {},
             {
                 "samples": 2,
-                "xs_interval": "quantiles",
-                "xs_values": [0.25, 0.5, 0.75],
                 "var_discrete": [3],
             },
             {"instances": 2},