Add softmax option to plot_pdp

Author: AlexAndorra

pymc_bart/utils.py

@@ -290,6 +290,7 @@ def plot_pdp(
     var_idx: Optional[list[int]] = None,
     var_discrete: Optional[list[int]] = None,
     func: Optional[Callable] = None,
+    softmax_link: Optional[bool] = False,
     samples: int = 200,
     ref_line: bool = True,
     random_seed: Optional[int] = None,
@@ -330,6 +331,9 @@ def plot_pdp(
         List of the indices of the covariate treated as discrete.
     func : Optional[Callable], by default None.
         Arbitrary function to apply to the predictions. Defaults to the identity function.
+    softmax_link: Optional[bool] = False,
+        If True the predictions are transformed using the softmax function. Only works when
+        likelihood is categorical. Defaults to False.
     samples : int
         Number of posterior samples used in the predictions. Defaults to 200
     ref_line : bool
@@ -396,6 +400,12 @@ def identity(x):
         p_d = _sample_posterior(
             all_trees, X=fake_X, rng=rng, size=samples, excluded=excluded, shape=shape
         )
+        if softmax_link is True:
+            from scipy.special import softmax
+
+            # categories are the last dimension
+            p_d = softmax(p_d, axis=-1)
+
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="hdi currently interprets 2d data")
             new_x = fake_X[:, var]