def plot_waterfall(explanation: Any, i: int = 0, *, max_features: int = 10, ax: Any = None) -> Any:
"""Plot one instance's I/D/M decomposition as a stacked waterfall."""
try:
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
except ImportError as exc:
raise ImportError(
"waterfall plots require matplotlib; install with `pip install triadxai[viz]`"
) from exc
rows = explanation.channels[explanation.channels["instance"] == i]
if rows.empty:
raise ValueError(f"no channels for instance {i}")
rows = rows.assign(total=rows["I"] + rows["D"] + rows["M"])
rows = rows.reindex(rows["total"].abs().sort_values(ascending=False).index)
top = rows.head(max_features)
rest = rows.iloc[max_features:]
labels: list[str] = []
segments: list[tuple[float, float, float]] = []
for row in top.itertuples(index=False):
marker = " *" if row.oov else ""
labels.append(f"{row.feature}{marker}")
segments.append((row.I, row.D, row.M))
if not rest.empty:
labels.append(f"other ({len(rest)})")
segments.append((rest["I"].sum(), rest["D"].sum(), rest["M"].sum()))
if ax is None:
_, ax = plt.subplots(figsize=(8, 0.5 * (len(labels) + 2) + 1.5))
y_positions = np.arange(len(labels), 0, -1)
for y, (seg_i, seg_d, seg_m) in zip(y_positions, segments, strict=True):
left = 0.0
for value, color, hatch, alpha in (
(seg_i, _COLOR_I, None, 1.0),
(seg_d, _COLOR_D, "///", 0.6),
(seg_m, _COLOR_M, None, 0.9),
):
if value != 0.0:
ax.barh(y, value, left=left, color=color, hatch=hatch, alpha=alpha, height=0.7)
left += value
total = float(explanation.score[i] - explanation.intercept)
ax.barh(0, total, color="none", edgecolor="black", height=0.7)
ax.axvline(0.0, color="black", linewidth=0.8)
ax.set_yticks([*y_positions, 0])
ax.set_yticklabels([*labels, "total (eta - beta0)"])
ax.set_xlabel("logit contribution")
title = f"TRIAD waterfall - instance {i}"
if explanation.approximate:
title += " (approximate: SHAP-based channels)"
ax.set_title(title)
ax.legend(
handles=[
mpatches.Patch(color=_COLOR_I, label="I - information"),
mpatches.Patch(
facecolor=_COLOR_D, hatch="///", alpha=0.6, label="D - epistemic (guess)"
),
mpatches.Patch(color=_COLOR_M, label="M - missingness"),
],
loc="best",
fontsize=8,
)
if top["oov"].any():
ax.annotate(
"* outside training range / unseen category",
xy=(0, -0.12),
xycoords="axes fraction",
fontsize=7,
)
return ax