import warnings
from typing import TYPE_CHECKING, Any, List, Optional, Tuple
import numpy as np
from .settings import logger
# Optional SHAP dependency
try:
import shap
HAS_SHAP = True
except ImportError:
shap = None # type: ignore[assignment]
HAS_SHAP = False
# Optional matplotlib dependency
try:
import matplotlib.pyplot as plt
HAS_MATPLOTLIB = True
except ImportError:
plt = None # type: ignore[assignment]
Figure = None # type: ignore[assignment]
HAS_MATPLOTLIB = False
# Import BaseModelWrapper with TYPE_CHECKING to avoid circular imports
if TYPE_CHECKING:
from .core import BaseModelWrapper
def _check_shap():
"""Check SHAP availability."""
if not HAS_SHAP:
raise ImportError(
"SHAP is required for explainability features. Install with: pip install shap"
)
def _warn_if_deep_explainer_deprecated() -> None:
"""Emit DeprecationWarning if shap >= 0.45 where DeepExplainer is deprecated."""
try:
_parts = shap.__version__.split(".")
if (int(_parts[0]), int(_parts[1])) >= (0, 45):
warnings.warn(
f"shap.DeepExplainer is deprecated in SHAP >= 0.45 "
f"(detected {shap.__version__}). It may be removed in a future "
"release. Consider shap.GradientExplainer as an alternative.",
DeprecationWarning,
stacklevel=3,
)
except (AttributeError, ValueError, IndexError):
pass
def _check_matplotlib():
"""Check matplotlib availability."""
if not HAS_MATPLOTLIB:
raise ImportError(
"matplotlib is required for SHAP plotting. Install with: pip install matplotlib"
)
# Tree model detection for SHAP explainer selection.
# Imported lazily at module load so sklearn is not a hard dependency.
try:
from sklearn.ensemble import (
AdaBoostClassifier,
ExtraTreesClassifier,
ExtraTreesRegressor,
GradientBoostingClassifier,
GradientBoostingRegressor,
HistGradientBoostingClassifier,
HistGradientBoostingRegressor,
RandomForestClassifier,
RandomForestRegressor,
)
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
_SKLEARN_TREE_TYPES: Tuple = ()
_SKLEARN_TREE_TYPES = (
AdaBoostClassifier,
DecisionTreeClassifier,
DecisionTreeRegressor,
ExtraTreesClassifier,
ExtraTreesRegressor,
GradientBoostingClassifier,
GradientBoostingRegressor,
HistGradientBoostingClassifier,
HistGradientBoostingRegressor,
RandomForestClassifier,
RandomForestRegressor,
)
_HAS_SKLEARN_TREES = True
except ImportError:
_SKLEARN_TREE_TYPES = ()
_HAS_SKLEARN_TREES = False
def _is_tree_model(model: Any) -> bool:
"""Return ``True`` if *model* should use ``shap.TreeExplainer``.
Uses ``isinstance()`` for known sklearn tree types. Falls back to
class-name fragment matching for XGBoost, LightGBM, and CatBoost,
which are not sklearn classes but are supported by TreeExplainer.
"""
if _HAS_SKLEARN_TREES and isinstance(model, _SKLEARN_TREE_TYPES):
return True
model_type = type(model).__name__
return any(hint in model_type for hint in ("XGB", "LGBM", "LightGBM", "CatBoost", "Booster"))
[docs]
def plot_shap_summary(
model_wrapper: "BaseModelWrapper",
X_data: np.ndarray,
feature_names: Optional[List[str]] = None,
plot_type: str = "bar",
):
"""
Generates a SHAP summary plot for a model.
Args:
model_wrapper (BaseModelWrapper): The trained model wrapper.
X_data (np.ndarray): The data to explain.
feature_names (Optional[List[str]]): List of feature names.
plot_type (str): The type of plot to generate ("bar", "dot", "violin").
"""
_check_shap()
_check_matplotlib()
if _is_tree_model(model_wrapper.model):
explainer = shap.TreeExplainer(model_wrapper.model)
shap_values = explainer.shap_values(X_data)
elif hasattr(model_wrapper.model, "layers"):
# For neural networks (Keras/TensorFlow models)
logger.info("Using DeepExplainer for neural network. This may take some time...")
try:
# Use a subset of the data as background for efficiency
background_size = min(100, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
background = X_data[_bg_idx]
_warn_if_deep_explainer_deprecated()
explainer = shap.DeepExplainer(model_wrapper.model, background)
shap_values = explainer.shap_values(X_data)
except Exception as e:
logger.warning(f"DeepExplainer failed: {e}")
logger.warning("Falling back to KernelExplainer (this will be slower)...")
# Fallback to KernelExplainer
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
if hasattr(model_wrapper.model, "predict_proba"):
predict_fn = model_wrapper.predict_proba
else:
predict_fn = model_wrapper.predict
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
shap_values = explainer.shap_values(X_data)
else:
# For other models, use KernelExplainer
logger.info(
"Using KernelExplainer, which can be very slow for large datasets. Consider using a smaller sample of your data for faster results."
)
# Use model's predict_proba if available, otherwise predict
if hasattr(model_wrapper, "predict_proba"):
predict_fn = model_wrapper.predict_proba
else:
predict_fn = model_wrapper.predict
# Use a smaller background dataset for efficiency
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
if hasattr(model_wrapper.model, "predict_proba"):
predict_fn = model_wrapper.predict_proba
else:
predict_fn = model_wrapper.predict
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
shap_values = explainer.shap_values(X_data)
# Handle the different formats SHAP might return
if isinstance(shap_values, list):
# Multi-class models return a list of arrays (one per class)
# Plot the first class by default
if len(shap_values) > 1:
logger.info("Multi-class model detected. Plotting SHAP values for class 0.")
logger.info(
f"Model has {len(shap_values)} classes. You may want to plot other classes separately."
)
shap.summary_plot(shap_values[0], X_data, feature_names=feature_names, plot_type=plot_type)
else:
# Single output (binary classification or regression)
shap.summary_plot(shap_values, X_data, feature_names=feature_names, plot_type=plot_type)
[docs]
def plot_shap_waterfall(
model_wrapper: "BaseModelWrapper",
X_data: np.ndarray,
sample_index: int = 0,
feature_names: Optional[List[str]] = None,
class_index: int = 0,
):
"""
Generates a SHAP waterfall plot for a single prediction.
Args:
model_wrapper (BaseModelWrapper): The trained model wrapper.
X_data (np.ndarray): The data to explain.
sample_index (int): Index of the sample to explain.
feature_names (Optional[List[str]]): List of feature names.
class_index (int): For multi-class models, which class to explain.
"""
_check_shap()
_check_matplotlib()
if sample_index >= len(X_data):
raise ValueError(
f"sample_index {sample_index} is out of range for data with {len(X_data)} samples"
)
if _is_tree_model(model_wrapper.model):
explainer = shap.TreeExplainer(model_wrapper.model)
elif hasattr(model_wrapper.model, "layers"):
background_size = min(100, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
background = X_data[_bg_idx]
_warn_if_deep_explainer_deprecated()
explainer = shap.DeepExplainer(model_wrapper.model, background)
else:
predict_fn = getattr(model_wrapper, "predict_proba", model_wrapper.predict)
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
# Get SHAP values for the specific sample
sample_data = X_data[sample_index : sample_index + 1]
shap_values = explainer.shap_values(sample_data)
# Handle multi-class vs single output across SHAP API versions.
# - SHAP < 0.45 multiclass: list of per-class 2-D arrays, each (1, n_features)
# - SHAP >= 0.45 multiclass: 3-D ndarray (1, n_features, n_classes)
# - Binary classifier or regression: 2-D ndarray (1, n_features)
if isinstance(shap_values, list):
# SHAP < 0.45 multiclass path
if class_index >= len(shap_values):
raise ValueError(
f"class_index {class_index} is out of range for model with "
f"{len(shap_values)} classes"
)
shap_values_to_plot = shap_values[class_index][0]
elif isinstance(shap_values, np.ndarray) and shap_values.ndim == 3:
# SHAP >= 0.45 multiclass path — 3-D: (n_samples, n_features, n_classes)
n_classes = shap_values.shape[2]
if class_index >= n_classes:
raise ValueError(
f"class_index {class_index} is out of range for model with " f"{n_classes} classes"
)
shap_values_to_plot = shap_values[0, :, class_index]
else:
# Binary classifier or regression — 2-D: (n_samples, n_features)
shap_values_to_plot = shap_values[0]
# Create explanation object for waterfall plot
if hasattr(shap, "Explanation"):
# expected_value shape across SHAP API versions:
# - list of scalars (one per class): SHAP < 0.45 multiclass
# - 1-D ndarray of scalars: SHAP >= 0.45 multiclass
# - scalar: binary or regression
expected_value = explainer.expected_value
if isinstance(expected_value, list):
base_value = expected_value[class_index]
elif isinstance(expected_value, np.ndarray) and expected_value.ndim >= 1:
base_value = float(expected_value[class_index])
else:
base_value = expected_value
explanation = shap.Explanation(
values=shap_values_to_plot,
base_values=base_value,
data=sample_data[0],
feature_names=feature_names,
)
shap.waterfall_plot(explanation)
else:
# Older SHAP versions - use force plot as fallback
logger.info("Waterfall plot not available in this SHAP version. Using force plot instead.")
# Same expected_value handling as the Explanation branch above.
expected_value = explainer.expected_value
if isinstance(expected_value, list):
base_value = expected_value[class_index]
elif isinstance(expected_value, np.ndarray) and expected_value.ndim >= 1:
base_value = float(expected_value[class_index])
else:
base_value = expected_value
shap.force_plot(
base_value, shap_values_to_plot, sample_data[0], feature_names=feature_names
)
[docs]
def plot_shap_dependence(
model_wrapper: "BaseModelWrapper",
X_data: np.ndarray,
feature_name: str,
feature_names: Optional[List[str]] = None,
interaction_index: Optional[int] = None,
class_index: int = 0,
):
"""
Generates a SHAP dependence plot showing how a feature affects predictions.
Args:
model_wrapper (BaseModelWrapper): The trained model wrapper.
X_data (np.ndarray): The data to explain.
feature_name (str): Name of the feature to plot (or index if feature_names not provided).
feature_names (Optional[List[str]]): List of feature names.
interaction_index (Optional[int]): Feature index to use for coloring interaction effects.
class_index (int): For multi-class models, which class to explain.
"""
_check_shap()
_check_matplotlib()
# Convert feature name to index if necessary
if feature_names and isinstance(feature_name, str):
if feature_name not in feature_names:
raise ValueError(f"Feature '{feature_name}' not found in feature_names")
feature_index = feature_names.index(feature_name)
else:
# Assume feature_name is already an index
try:
feature_index = int(feature_name)
except (ValueError, TypeError):
raise ValueError(
f"feature_name must be either a string (if feature_names provided) or an integer index"
)
if feature_index >= X_data.shape[1]:
raise ValueError(
f"Feature index {feature_index} is out of range for data with {X_data.shape[1]} features"
)
try:
if _is_tree_model(model_wrapper.model):
explainer = shap.TreeExplainer(model_wrapper.model)
elif hasattr(model_wrapper.model, "layers"):
try:
background_size = min(100, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
background = X_data[_bg_idx]
_warn_if_deep_explainer_deprecated()
explainer = shap.DeepExplainer(model_wrapper.model, background)
except Exception:
predict_fn = getattr(model_wrapper, "predict_proba", model_wrapper.predict)
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
else:
predict_fn = getattr(model_wrapper, "predict_proba", model_wrapper.predict)
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
# Get SHAP values
shap_values = explainer.shap_values(X_data)
# Handle multi-class vs single output
if isinstance(shap_values, list):
if class_index >= len(shap_values):
raise ValueError(
f"class_index {class_index} is out of range for model with {len(shap_values)} classes"
)
shap_values_to_plot = shap_values[class_index]
else:
shap_values_to_plot = shap_values
# Create dependence plot
shap.dependence_plot(
feature_index,
shap_values_to_plot,
X_data,
feature_names=feature_names,
interaction_index=interaction_index,
show=False,
)
plt.show()
except Exception as e:
raise RuntimeError(f"SHAP dependence plot failed: {e}")
[docs]
def get_shap_feature_importance(
model_wrapper: "BaseModelWrapper",
X_data: np.ndarray,
feature_names: Optional[List[str]] = None,
class_index: int = 0,
) -> np.ndarray:
"""
Get feature importance scores based on mean absolute SHAP values.
Args:
model_wrapper (BaseModelWrapper): The trained model wrapper.
X_data (np.ndarray): The data to explain.
feature_names (Optional[List[str]]): List of feature names.
class_index (int): For multi-class models, which class to analyze.
Returns:
np.ndarray: Mean absolute SHAP values for each feature.
"""
_check_shap()
if _is_tree_model(model_wrapper.model):
explainer = shap.TreeExplainer(model_wrapper.model)
elif hasattr(model_wrapper.model, "layers"):
background_size = min(100, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
background = X_data[_bg_idx]
_warn_if_deep_explainer_deprecated()
explainer = shap.DeepExplainer(model_wrapper.model, background)
else:
if hasattr(model_wrapper.model, "predict_proba"):
predict_fn = model_wrapper.predict_proba
else:
predict_fn = model_wrapper.predict
background_size = min(50, len(X_data))
_bg_rng = np.random.default_rng(42)
_bg_idx = _bg_rng.choice(len(X_data), size=background_size, replace=False)
explainer = shap.KernelExplainer(predict_fn, X_data[_bg_idx])
# Get SHAP values
shap_values = explainer.shap_values(X_data)
# Handle multi-class vs single output
if isinstance(shap_values, list):
if class_index >= len(shap_values):
raise ValueError(
f"class_index {class_index} is out of range for model with {len(shap_values)} classes"
)
shap_values_to_use = shap_values[class_index]
else:
shap_values_to_use = shap_values
# Calculate mean absolute SHAP values
feature_importance = np.mean(np.abs(shap_values_to_use), axis=0)
return feature_importance