feat: add waterfall plot for local explanations (#453)

python/interpret-core/interpret/core/templates.py

@@ -61,7 +61,7 @@ def data(self, key=None):
             return None
         return self._internal_obj["specific"][key]
 
-    def visualize(self, key=None, title=None):
+    def visualize(self, key=None, title=None, plot_type="bar", **kwargs):
         """Provides interactive visualizations.
 
         Args:
@@ -82,6 +82,7 @@ def visualize(self, key=None, title=None):
             plot_horizontal_bar,
             plot_line,
             plot_pairwise_heatmap,
+            plot_waterfall,
             sort_take,
         )
 
@@ -98,6 +99,9 @@ def visualize(self, key=None, title=None):
 
         # Handle local instance graphs
         if self.explanation_type == "local":
+            if plot_type == "waterfall":
+                return plot_waterfall(data_dict, title=title, **kwargs)
+
             if isinstance(key, tuple) and len(key) == 2:
                 provider, key = key
                 # TODO: MLI should handle multiclass at a future date.

python/interpret-core/interpret/glassbox/_ebm.py

@@ -226,7 +226,13 @@ class EBMExplanation(FeatureValueExplanation):
 
     explanation_type = None
 
-    def visualize(self, key: int | None = None, title: str | None = None) -> Any:
+    def visualize(
+        self,
+        key: int | None = None,
+        title: str | None = None,
+        plot_type: str = "bar",
+        **kwargs: Any,
+    ) -> Any:
         """Provide interactive visualizations.
 
         Args:
@@ -320,11 +326,23 @@ def visualize(self, key: int | None = None, title: str | None = None) -> Any:
 
         # Local explanation graph
         if self.explanation_type == "local":
-            figure = super().visualize(key)
-            figure.update_layout(
-                title="Local Explanation (" + figure.layout.title.text + ")",
-                xaxis_title="Contribution to Prediction",
-            )
+            figure = super().visualize(key, title=title, plot_type=plot_type, **kwargs)
+            if figure is None:
+                return None
+            orig_title = figure.layout.title.text if figure.layout.title else ""
+            if plot_type == "waterfall":
+                figure.update_layout(
+                    title="Local Explanation (" + orig_title + ")"
+                    if orig_title
+                    else "Local Explanation",
+                )
+            else:
+                figure.update_layout(
+                    title="Local Explanation (" + orig_title + ")"
+                    if orig_title
+                    else "Local Explanation",
+                    xaxis_title="Contribution to Prediction",
+                )
             figure._interpret_help_text = (
                 "A local explanation shows the breakdown of how much "
                 "each term contributed to the prediction for a single sample. The intercept "

python/interpret-core/interpret/glassbox/_linear.py

@@ -367,7 +367,7 @@ class LinearExplanation(FeatureValueExplanation):
 
     explanation_type = None
 
-    def visualize(self, key=None):
+    def visualize(self, key=None, title=None, plot_type="bar", **kwargs):
         """Provides interactive visualizations.
 
         Args:
@@ -428,7 +428,7 @@ def visualize(self, key=None):
                 data_dict, title="Overall Importance:<br>Coefficients"
             )
 
-        return super().visualize(key)
+        return super().visualize(key, title=title, plot_type=plot_type, **kwargs)
 
 
 class LinearRegression(_SKRegressorMixin, BaseLinear):

python/interpret-core/interpret/provider/_visualize.py

@@ -80,14 +80,15 @@ def render(self, explanation, key=-1, **kwargs):
 class PreserveProvider(VisualizeProvider):
     def render(self, explanation, key=-1, **kwargs):
         file_name = kwargs.pop("file_name", None)
+        plot_type = kwargs.pop("plot_type", "bar")
 
         # NOTE: Preserve didn't support returning everything. If key is -1 default to key is None.
         # This is for backward-compatibility. All of this will be deprecated shortly anyway.
         if key == -1:
             key = None
 
         # Get visual object
-        visual = explanation.visualize(key=key)
+        visual = explanation.visualize(key=key, plot_type=plot_type)
 
         # Output to front-end/file
 
@@ -234,4 +235,5 @@ def render(self, explanation, key=-1, **kwargs):
             default_key=key,
             detected_envs=self.detected_envs,
             js_url=self.js_url,
+            **kwargs,
         )

python/interpret-core/interpret/visual/_inline.py

@@ -106,15 +106,15 @@ def _build_viz_err_obj(err_msg):
     }
 
 
-def _build_viz_obj(explanation, detected_envs):
-    overall = _build_viz_figure(explanation.visualize(), detected_envs)
+def _build_viz_obj(explanation, detected_envs, **kwargs):
+    overall = _build_viz_figure(explanation.visualize(**kwargs), detected_envs)
     if explanation.selector is None:
         # NOTE: Unsure if this should be a list or None in the long term.
         specific = []
         selector_obj = {"columns": [], "data": []}
     else:
         specific = [
-            _build_viz_figure(explanation.visualize(i), detected_envs)
+            _build_viz_figure(explanation.visualize(i, **kwargs), detected_envs)
             for i in range(len(explanation.selector["data"]))
         ]
         selector_obj = explanation.selector
@@ -218,7 +218,14 @@ def _render_databricks(js):  # pragma: no cover
     _databricks_displayHTML(js)
 
 
-def render(explanation, id_str=None, default_key=-1, detected_envs=None, js_url=None):
+def render(
+    explanation,
+    id_str=None,
+    default_key=-1,
+    detected_envs=None,
+    js_url=None,
+    **kwargs,
+):
     from IPython.display import HTML, display
 
     global _jupyter_initialized
@@ -227,7 +234,7 @@ def render(explanation, id_str=None, default_key=-1, detected_envs=None, js_url=
         msg = "Dashboard not yet supported in cloud environments."
         viz_obj = _build_viz_err_obj(msg)
     else:
-        viz_obj = _build_viz_obj(explanation, detected_envs)
+        viz_obj = _build_viz_obj(explanation, detected_envs, **kwargs)
 
     init_js, body_js = _build_javascript(
         viz_obj, id_str, default_key=default_key, js_url=js_url

python/interpret-core/interpret/visual/plot.py

@@ -605,6 +605,167 @@ def plot_horizontal_bar(
     return go.Figure(data=traces, layout=layout)
 
 
+def plot_waterfall(
+    data_dict,
+    title="",
+    xtitle="",
+    ytitle="",
+    orientation="h",
+    top_n=15,
+    class_idx=None,
+    increasing_color=None,
+    decreasing_color=None,
+    total_color=None,
+):
+    """Produces a waterfall plot for local explanations.
+
+    Visualizes how the model prediction is obtained by accumulating the intercept
+    and individual feature contributions towards the final total prediction.
+
+    Args:
+        data_dict: Dictionary containing explanation data for a single instance.
+        title: Title of the graph. If empty, generates a summary title from perf metrics.
+        xtitle: Label for the x-axis.
+        ytitle: Label for the y-axis.
+        orientation: "h" for horizontal waterfall (default), "v" for vertical waterfall.
+        top_n: Number of top contributing terms to show individually. Remaining terms
+            are aggregated into an "Other features" step.
+        class_idx: For multiclass explanations, which class index to plot. Defaults to
+            the predicted class.
+        increasing_color: Color for positive contributions (defaults to COLORS[1]).
+        decreasing_color: Color for negative contributions (defaults to COLORS[0]).
+        total_color: Color for the total / prediction bar (defaults to COLORS[2]).
+
+    Returns:
+        A Plotly Figure object.
+    """
+    if data_dict.get("scores", None) is None:  # pragma: no cover
+        return None
+
+    is_multiclass = is_multiclass_local_data_dict(data_dict)
+    if is_multiclass:
+        if class_idx is None:
+            class_idx = (
+                data_dict["perf"]["predicted"]
+                if "perf" in data_dict and "predicted" in data_dict["perf"]
+                else 0
+            )
+        scores = [float(s[class_idx]) for s in data_dict["scores"]]
+    else:
+        scores = [float(s) for s in data_dict["scores"]]
+
+    names = list(data_dict["names"])
+    values = data_dict.get("values", None)
+    if values is not None:
+        values = list(values)
+        names = _names_with_values(names, values)
+
+    extra = data_dict.get("extra", None)
+    if extra is not None and "scores" in extra and len(extra["scores"]) > 0:
+        if is_multiclass:
+            intercept = float(extra["scores"][0][class_idx])
+        else:
+            intercept = float(extra["scores"][0])
+        intercept_name = extra.get("names", ["Intercept"])[0]
+    else:
+        intercept = 0.0
+        intercept_name = "Intercept"
+
+    if title == "":
+        title_items = []
+        if data_dict.get("perf", None) is not None:
+            predicted = data_dict["perf"]["predicted"]
+            actual = data_dict["perf"]["actual"]
+            predicted_score = data_dict["perf"]["predicted_score"]
+            actual_score = data_dict["perf"]["actual_score"]
+
+            if "meta" in data_dict and "label_names" in data_dict["meta"]:
+                label_names = data_dict["meta"]["label_names"]
+                pred_label = label_names[predicted]
+
+                title_str = ""
+                if not np.isnan(actual):
+                    actual_class = label_names[actual]
+                    title_str += f"Actual Class: {actual_class} | "
+
+                title_str += f"Predicted Class: {pred_label}"
+                title_str += f"<br />Pr(y = {pred_label}): {predicted_score:.3f}"
+
+                if not np.isnan(actual) and len({predicted, actual}) == 2:
+                    title_str += f" | Pr(y = {actual}): {actual_score:.3f}"
+                title_items.append(title_str)
+            else:
+                if not np.isnan(actual):
+                    actual_score_str = _pretty_number(actual_score)
+                    title_items.append(f"Actual: {actual_score_str}")
+
+                pred_score_str = _pretty_number(predicted_score)
+                title_items.append(f"Predicted: {pred_score_str}")
+
+        if is_multiclass and "meta" in data_dict and "label_names" in data_dict["meta"]:
+            cls_name = data_dict["meta"]["label_names"][class_idx]
+            title_items.append(f"Class: {cls_name}")
+
+        title = " | ".join(title_items)
+
+    pairs = sorted(zip(names, scores, strict=False), key=lambda x: -abs(x[1]))
+    if top_n is not None and len(pairs) > top_n:
+        top_pairs = pairs[:top_n]
+        remaining_sum = sum(s for _, s in pairs[top_n:])
+        if abs(remaining_sum) > 1e-9:
+            top_pairs.append(("Other features", remaining_sum))
+        pairs = top_pairs
+
+    step_names = [intercept_name] + [p[0] for p in pairs] + ["Prediction"]
+    step_scores = [intercept] + [p[1] for p in pairs] + [0.0]
+    measures = ["relative"] * (len(step_names) - 1) + ["total"]
+
+    inc_color = increasing_color or COLORS[1]
+    dec_color = decreasing_color or COLORS[0]
+    tot_color = total_color or COLORS[2]
+
+    is_horizontal = orientation.lower().startswith("h")
+    if is_horizontal:
+        waterfall_trace = go.Waterfall(
+            orientation="h",
+            measure=measures,
+            x=step_scores,
+            y=step_names,
+            connector={"line": {"color": "rgb(63, 63, 63)", "width": 1}},
+            increasing={"marker": {"color": inc_color}},
+            decreasing={"marker": {"color": dec_color}},
+            totals={"marker": {"color": tot_color}},
+        )
+        layout = {
+            "title": title,
+            "yaxis": {
+                "automargin": True,
+                "title": ytitle,
+                "dtick": 1,
+                "autorange": "reversed",
+            },
+            "xaxis": {"title": xtitle or "Score Contribution"},
+        }
+    else:
+        waterfall_trace = go.Waterfall(
+            orientation="v",
+            measure=measures,
+            x=step_names,
+            y=step_scores,
+            connector={"line": {"color": "rgb(63, 63, 63)", "width": 1}},
+            increasing={"marker": {"color": inc_color}},
+            decreasing={"marker": {"color": dec_color}},
+            totals={"marker": {"color": tot_color}},
+        )
+        layout = {
+            "title": title,
+            "xaxis": {"automargin": True, "title": xtitle, "dtick": 1},
+            "yaxis": {"title": ytitle or "Score Contribution"},
+        }
+
+    return go.Figure(data=[waterfall_trace], layout=layout)
+
+
 def mli_plot_horizontal_bar(
     scores,
     names,