"""SHAP-based model explainability: global importance, summary plots, and single-row waterfalls.""" import logging import os from dataclasses import dataclass, field from pathlib import Path import numpy as np import shap logger = logging.getLogger(__name__) # Silence SHAP's own verbose internal logger logging.getLogger("shap").setLevel(logging.WARNING) # Tree-based model class names that support TreeExplainer (fast) _TREE_MODEL_CLASSES = ( "XGBClassifier", "XGBRegressor", "LGBMClassifier", "LGBMRegressor", "CatBoostClassifier", "CatBoostRegressor", "RandomForestClassifier", "RandomForestRegressor", "ExtraTreesClassifier", "ExtraTreesRegressor", "GradientBoostingClassifier", "GradientBoostingRegressor", ) # Max background samples for KernelExplainer (slow fallback) _KERNEL_BACKGROUND_SAMPLES = 100 # Max rows to compute SHAP values on (keeps it fast for large datasets) _MAX_EXPLAIN_ROWS = 100 @dataclass class ExplainerResult: """Output of a SHAP explanation run.""" feature_importance: dict[str, float] # feature name → mean |SHAP|, sorted desc plot_paths: dict[str, str] # plot type → file path shap_values: np.ndarray = field(repr=False, default=None) feature_names: list[str] = field(default_factory=list) def _is_tree_model(model) -> bool: """Check if a model supports the fast TreeExplainer.""" return type(model).__name__ in _TREE_MODEL_CLASSES def _get_explainer(model, X_background: np.ndarray) -> shap.Explainer: """Return the appropriate SHAP explainer for the given model. Uses TreeExplainer for tree-based models (fast, exact). Falls back to KernelExplainer for all others (slow, model-agnostic). Args: model: Fitted sklearn-compatible estimator. X_background: Background dataset used by KernelExplainer. Returns: A SHAP explainer instance. """ if _is_tree_model(model): logger.info("Using TreeExplainer for %s.", type(model).__name__) return shap.TreeExplainer(model) logger.info( "Using KernelExplainer for %s (background=%d rows).", type(model).__name__, len(X_background), ) # KernelExplainer needs a callable — use predict_proba for classifiers, predict otherwise if hasattr(model, "predict_proba"): fn = model.predict_proba else: fn = model.predict background = shap.sample(X_background, min(_KERNEL_BACKGROUND_SAMPLES, len(X_background))) return shap.KernelExplainer(fn, background) def _compute_shap_values(explainer, X: np.ndarray) -> np.ndarray: """Compute SHAP values, handling both old and new SHAP API shapes. For binary classification TreeExplainer returns shape (n, features) for class 1. For multiclass it returns (n_classes, n, features). We normalise to always return (n, features). Args: explainer: Fitted SHAP explainer. X: Feature matrix to explain. Returns: SHAP values array, shape (n_samples, n_features). """ values = explainer.shap_values(X) # List of arrays (older SHAP API): one array per class if isinstance(values, list): if len(values) == 2: return values[1] # binary: positive class return np.mean(np.abs(np.stack(values, axis=0)), axis=0) # 3D array (newer KernelExplainer with predict_proba): (n_samples, n_features, n_classes) if isinstance(values, np.ndarray) and values.ndim == 3: if values.shape[2] == 2: return values[:, :, 1] # binary: positive class return np.abs(values).mean(axis=2) return values def explain( model, X: np.ndarray, feature_names: list[str], output_dir: str = "shap_plots", row_index: int = 0, ) -> ExplainerResult: """Run a full SHAP explanation pass and save plots to disk. Generates three artefacts: - Bar chart of global mean |SHAP| feature importance - Beeswarm summary plot (distribution of SHAP values per feature) - Waterfall plot for a single prediction (row_index) Args: model: Fitted estimator (tree-based or any sklearn model). X: Preprocessed feature matrix (numpy array). feature_names: Feature names aligned with X columns. output_dir: Directory to save plot images. row_index: Row in X to use for the single-prediction waterfall plot. Returns: ExplainerResult with feature_importance dict and plot_paths dict. """ import matplotlib matplotlib.use("Agg") # headless — no display required import matplotlib.pyplot as plt Path(output_dir).mkdir(parents=True, exist_ok=True) # Cap rows for speed n = min(_MAX_EXPLAIN_ROWS, X.shape[0]) X_sample = X[:n] try: explainer = _get_explainer(model, X_sample) shap_values = _compute_shap_values(explainer, X_sample) except Exception as exc: logger.warning("SHAP explanation failed: %s", exc) return ExplainerResult(feature_importance={}, plot_paths={}) # Global importance: mean absolute SHAP per feature mean_abs = np.abs(shap_values).mean(axis=0) importance = dict( sorted( zip(feature_names, mean_abs.tolist()), key=lambda kv: kv[1], reverse=True, ) ) plot_paths: dict[str, str] = {} # --- Bar chart --- try: fig, ax = plt.subplots(figsize=(8, 6)) top_n = 20 names = list(importance.keys())[:top_n] vals = list(importance.values())[:top_n] ax.barh(names[::-1], vals[::-1], color="#1f77b4") ax.set_xlabel("Mean |SHAP value|") ax.set_title("Feature Importance (SHAP)") plt.tight_layout() path = os.path.join(output_dir, "shap_bar.png") fig.savefig(path, dpi=120, bbox_inches="tight") plt.close(fig) plot_paths["bar"] = path logger.info("Saved SHAP bar chart → %s", path) except Exception as exc: logger.warning("Failed to save SHAP bar chart: %s", exc) # --- Beeswarm summary plot --- try: plt.figure(figsize=(8, 6)) shap.summary_plot(shap_values, X_sample, feature_names=feature_names, show=False) path = os.path.join(output_dir, "shap_beeswarm.png") plt.savefig(path, dpi=120, bbox_inches="tight") plt.close() plot_paths["beeswarm"] = path logger.info("Saved SHAP beeswarm → %s", path) except Exception as exc: logger.warning("Failed to save SHAP beeswarm: %s", exc) # --- Waterfall plot for a single row --- try: row_idx = min(row_index, n - 1) expected_value = ( explainer.expected_value[1] if isinstance(explainer.expected_value, (list, np.ndarray)) else explainer.expected_value ) explanation = shap.Explanation( values=shap_values[row_idx], base_values=expected_value, data=X_sample[row_idx], feature_names=feature_names, ) plt.figure() shap.plots.waterfall(explanation, show=False) path = os.path.join(output_dir, "shap_waterfall.png") plt.savefig(path, dpi=120, bbox_inches="tight") plt.close() plot_paths["waterfall"] = path logger.info("Saved SHAP waterfall → %s", path) except Exception as exc: logger.warning("Failed to save SHAP waterfall: %s", exc) return ExplainerResult( feature_importance=importance, plot_paths=plot_paths, shap_values=shap_values, feature_names=feature_names, )