openavmkit.shap_analysis

SHAP-value computation and reporting for tree-based models.

Computes SHAP (Shapley) values across all subsets (test, train, sales, universe) for XGBoost, LightGBM, and CatBoost models. Used by :mod:openavmkit.modeling to produce the per-feature params and contributions outputs that mirror what linear models produce as coefficients and coefficient×value contributions.

See Also

openavmkit.modeling : Calls into this module to build params/contribs files for tree-based models.

explanation_from_contributions

explanation_from_contributions(df_contrib, df_features, key_col='key_sale')

Rebuild a plottable shap.Explanation from a contributions_* table.

This is the inverse of :func:make_shap_table. It is used for models whose SHAP values exist only as on-disk contributions (e.g. the ensemble, whose contributions are assembled as a weighted combination of its members rather than computed from a live estimator), so there is no explainer to re-run.

Parameters:

Name	Type	Description	Default
df_contrib	DataFrame	A contributions table: a base_value (or intercept) column, one column per feature, and bookkeeping columns (key, key_sale, contribution_sum, prediction, check_delta) which are dropped.	required
df_features	DataFrame	Raw feature values for the same rows, used as the beeswarm data matrix (it colors points by feature magnitude). Aligned to df_contrib by key_col and reindexed to the feature columns; missing or non-numeric features become NaN (rendered gray by the beeswarm).	required
key_col	str	The column shared by both frames used to align rows. Defaults to "key_sale".	'key_sale'

Returns:

Type	Description
Explanation	With values (n, m), base_values (n,), data (n, m) and feature_names of length m.

Source code in openavmkit/shap_analysis.py

def explanation_from_contributions(
    df_contrib: pd.DataFrame,
    df_features: pd.DataFrame,
    key_col: str = "key_sale",
) -> shap.Explanation:
    """Rebuild a plottable ``shap.Explanation`` from a ``contributions_*`` table.

    This is the inverse of :func:`make_shap_table`. It is used for models whose
    SHAP values exist only as on-disk contributions (e.g. the ensemble, whose
    contributions are assembled as a weighted combination of its members rather
    than computed from a live estimator), so there is no explainer to re-run.

    Parameters
    ----------
    df_contrib : pd.DataFrame
        A contributions table: a ``base_value`` (or ``intercept``) column, one
        column per feature, and bookkeeping columns (``key``, ``key_sale``,
        ``contribution_sum``, ``prediction``, ``check_delta``) which are dropped.
    df_features : pd.DataFrame
        Raw feature values for the same rows, used as the beeswarm ``data`` matrix
        (it colors points by feature magnitude). Aligned to ``df_contrib`` by
        ``key_col`` and reindexed to the feature columns; missing or non-numeric
        features become NaN (rendered gray by the beeswarm).
    key_col : str
        The column shared by both frames used to align rows. Defaults to
        ``"key_sale"``.

    Returns
    -------
    shap.Explanation
        With ``values`` (n, m), ``base_values`` (n,), ``data`` (n, m) and
        ``feature_names`` of length m.
    """
    base_col = (
        "base_value" if "base_value" in df_contrib.columns
        else ("intercept" if "intercept" in df_contrib.columns else None)
    )
    feature_names = [
        c for c in df_contrib.columns if c not in _CONTRIB_NON_FEATURE_COLS
    ]

    values = df_contrib[feature_names].apply(
        pd.to_numeric, errors="coerce"
    ).to_numpy(dtype=float)

    if base_col is not None:
        base_values = pd.to_numeric(
            df_contrib[base_col], errors="coerce"
        ).to_numpy(dtype=float)
    else:
        base_values = np.zeros(len(df_contrib), dtype=float)

    # Align raw feature values to the contribution rows for the `data` matrix.
    if key_col in df_contrib.columns and key_col in df_features.columns:
        feats = df_features.copy()
        feats[key_col] = feats[key_col].astype(str)
        feats = feats[~feats[key_col].duplicated()].set_index(key_col)
        keys = df_contrib[key_col].astype(str)
        data = feats.reindex(keys).reindex(columns=feature_names)
    else:
        # No shared key: fall back to positional alignment.
        data = df_features.reindex(columns=feature_names)

    data = data.apply(pd.to_numeric, errors="coerce").to_numpy(dtype=float)

    return shap.Explanation(
        values=values,
        base_values=base_values,
        data=data,
        feature_names=feature_names,
    )

get_full_layeredcomp_shaps

get_full_layeredcomp_shaps(model, X_train, X_test, X_sales, X_univ, verbose=False, subsets=None)

Compute exact SHAP Explanations for all subsets of a LayeredComp model.

Mirrors :func:get_full_model_shaps / :func:get_full_ngboost_shaps for the LayeredComp engine's folded-tree decomposition.

Source code in openavmkit/shap_analysis.py

def get_full_layeredcomp_shaps(
    model: LayeredCompModel,
    X_train: pd.DataFrame,
    X_test: pd.DataFrame,
    X_sales: pd.DataFrame,
    X_univ: pd.DataFrame,
    verbose: bool = False,
    subsets: frozenset | set | None = None,
):
    """Compute exact SHAP Explanations for all subsets of a LayeredComp model.

    Mirrors :func:`get_full_model_shaps` / :func:`get_full_ngboost_shaps` for the
    LayeredComp engine's folded-tree decomposition.
    """
    explainer = _layeredcomp_shap(model)

    def explain(X, label):
        return _shap_explain("layeredcomp", explainer, X, verbose=verbose, label=label)

    if verbose:
        print("Generating LayeredComp SHAPs...")

    want = _wanted_subsets(subsets)
    return {
        "train": explain(X_train, "train") if "train" in want else None,
        "test": explain(X_test, "test") if "test" in want else None,
        "sales": explain(X_sales, "sales") if "sales" in want else None,
        "universe": explain(X_univ, "universe") if "universe" in want else None,
    }

get_full_model_shaps

get_full_model_shaps(model, X_train, X_test, X_sales, X_univ, verbose=False, subsets=None)

Calculates shaps for all subsets (test, train, sales, universe) of one model run

Parameters:

Name	Type	Description	Default
model	XGBoostModel | LightGBMModel | CatBoostModel	A trained prediction model	required
X_train	DataFrame	2D array of independent variables' values from the training set	required
X_test	DataFrame	2D array of independent variables' values from the testing set	required
X_sales	DataFrame	2D array of independent variables' values from the sales set	required
X_univ	DataFrame	2D array of independent variables' values from the universe set	required
verbose	bool	Whether to print verbose information. Defaults to False.	False
subsets	frozenset | set | None	Which subsets to actually explain. None (default) explains all four. Any subset left out gets None in its slot instead of an Explanation, and its SHAP pass is skipped entirely. X_train is still required either way: it backs the explainer.	None

Returns:

Type	Description
dict	A dict containing shap.Explanation objects (or None) keyed to "train", "test", "sales", and "universe"

Source code in openavmkit/shap_analysis.py

def get_full_model_shaps(
    model: XGBoostModel | LightGBMModel | CatBoostModel,
    X_train: pd.DataFrame,
    X_test: pd.DataFrame,
    X_sales: pd.DataFrame,
    X_univ: pd.DataFrame,
    verbose: bool = False,
    subsets: frozenset | set | None = None,
):
    """
    Calculates shaps for all subsets (test, train, sales, universe) of one model run

    Parameters
    ----------
    model: XGBoostModel | LightGBMModel | CatBoostModel
        A trained prediction model
    X_train: pd.DataFrame
        2D array of independent variables' values from the training set
    X_test: pd.DataFrame
        2D array of independent variables' values from the testing set
    X_sales: pd.DataFrame
        2D array of independent variables' values from the sales set
    X_univ: pd.DataFrame
        2D array of independent variables' values from the universe set
    verbose: bool
        Whether to print verbose information. Defaults to False.
    subsets: frozenset | set | None
        Which subsets to actually explain. ``None`` (default) explains all four. Any subset
        left out gets ``None`` in its slot instead of an Explanation, and its SHAP pass is
        skipped entirely. X_train is still required either way: it backs the explainer.

    Returns
    -------
    dict
        A dict containing shap.Explanation objects (or None) keyed to "train", "test",
        "sales", and "universe"

    """

    # NGBoost uses an exact additive tree decomposition rather than a single
    # TreeExplainer; delegate to its dedicated path (param_index=0 -> mean).
    if isinstance(model, NGBoostModel):
        return get_full_ngboost_shaps(
            model, X_train, X_test, X_sales, X_univ, param_index=0, verbose=verbose,
            subsets=subsets,
        )

    # LayeredComp folds its path-weighted prediction into per-leaf values and
    # gets exact path-dependent tree-SHAP via its own hand-rolled explainer.
    if isinstance(model, LayeredCompModel):
        return get_full_layeredcomp_shaps(
            model, X_train, X_test, X_sales, X_univ, verbose=verbose, subsets=subsets
        )

    tree_explainer: shap.TreeExplainer

    approximate = True
    cat_data = model.cat_data

    model_type = ""
    if isinstance(model, XGBoostModel):
        model_type = "xgboost"
        tree_explainer = _xgboost_shap(model, X_train)
    elif isinstance(model, LightGBMModel):
        model_type = "lightgbm"
        approximate = False # approx. not supported for LightGBM
        tree_explainer = _lightgbm_shap(model, X_train)
    elif isinstance(model, CatBoostModel):
        model_type = "catboost"
        tree_explainer = _catboost_shap(model, X_train)
    else:
        raise ValueError(f"Unsupported model type: {type(model)}")

    if verbose:
        print(f"Generating SHAPs...")

    want = _wanted_subsets(subsets)

    def _maybe(name, X):
        if name not in want:
            return None
        return _shap_explain(
            model_type, tree_explainer, X, cat_data=cat_data,
            approximate=approximate, verbose=verbose, label=name,
        )

    shap_sales = _maybe("sales", X_sales)
    shap_train = _maybe("train", X_train)
    shap_test  = _maybe("test", X_test)
    shap_univ  = _maybe("universe", X_univ)

    return {
        "train": shap_train,
        "test":  shap_test,
        "sales": shap_sales,
        "universe":  shap_univ,
    }

get_full_ngboost_shaps

get_full_ngboost_shaps(model, X_train, X_test, X_sales, X_univ, param_index=0, verbose=False, subsets=None)

Compute exact SHAP Explanations for all subsets for one NGBoost distribution parameter.

Mirrors :func:get_full_model_shaps but for NGBoost's additive tree decomposition. param_index selects loc (0, the mean) or logscale (1, log predictive std).

Source code in openavmkit/shap_analysis.py

def get_full_ngboost_shaps(
    model: NGBoostModel,
    X_train: pd.DataFrame,
    X_test: pd.DataFrame,
    X_sales: pd.DataFrame,
    X_univ: pd.DataFrame,
    param_index: int = 0,
    verbose: bool = False,
    subsets: frozenset | set | None = None,
):
    """Compute exact SHAP Explanations for all subsets for one NGBoost distribution parameter.

    Mirrors :func:`get_full_model_shaps` but for NGBoost's additive tree
    decomposition. ``param_index`` selects loc (0, the mean) or logscale
    (1, log predictive std).
    """
    cat_data = model.cat_data
    explainer = _ngboost_shap(model, param_index=param_index)

    def explain(X, label):
        return _shap_explain(
            "ngboost", explainer, X, cat_data=cat_data, verbose=verbose, label=label
        )

    if verbose:
        which = "loc (mean)" if param_index == 0 else "logscale (uncertainty)"
        print(f"Generating NGBoost SHAPs for {which}...")

    want = _wanted_subsets(subsets)
    return {
        "train": explain(X_train, "train") if "train" in want else None,
        "test": explain(X_test, "test") if "test" in want else None,
        "sales": explain(X_sales, "sales") if "sales" in want else None,
        "universe": explain(X_univ, "universe") if "universe" in want else None,
    }

make_shap_table

make_shap_table(expl, list_keys, list_vars, list_keys_sale=None, include_pred=True)

Convert a shap explanation into a dataframe breaking down the full contribution to value

Parameters:

Name	Type	Description	Default
expl	Explanation	Output of your _xgboost_shap (values: (n,m), base_values: scalar or (n,)).	required
list_keys	list[str]	Primary keys in the same row order as X_to_explain	required
list_vars	list[str]	Feature names in the order used for training (your canonical order).	required
list_keys_sale	list[str] | None	Optional. Transaction keys in the same row order as X_to_explain. Default is None.	None
include_pred	bool	Optional. Add a column that reconstructs the model output on the explained scale: base_value + sum(shap_values across features). Default is True.	True

Returns:

Type	Description
DataFrame

Source code in openavmkit/shap_analysis.py

def make_shap_table(
    expl: shap.Explanation,
    list_keys: list[str],
    list_vars: list[str],
    list_keys_sale: list[str] = None,
    include_pred: bool = True
) -> pd.DataFrame:
    """
    Convert a shap explanation into a dataframe breaking down the full contribution to value

    Parameters
    ----------
    expl : shap.Explanation
        Output of your _xgboost_shap (values: (n,m), base_values: scalar or (n,)).
    list_keys : list[str]
        Primary keys in the same row order as X_to_explain
    list_vars : list[str]
        Feature names in the order used for training (your canonical order).
    list_keys_sale : list[str] | None
        Optional. Transaction keys in the same row order as X_to_explain. Default is None.
    include_pred : bool
        Optional. Add a column that reconstructs the model output on the explained scale:
        base_value + sum(shap_values across features). Default is True.

    Returns
    -------
    pd.DataFrame
    """
    # 1) Validate / normalize SHAP values shape (expect regression/binary: (n, m))
    vals = expl.values
    if isinstance(vals, list):
        raise ValueError("Got a list of SHAP arrays (likely multiclass). Handle per-class tables separately.")
    vals = np.asarray(vals)
    if vals.ndim != 2:
        raise ValueError(f"Expected 2D SHAP values (n_samples, n_features), got shape {vals.shape}.")

    n, m = vals.shape

    # 2) Base values: scalar or per-row
    base = expl.base_values
    if np.isscalar(base):
        base_arr = np.full((n,), float(base))
    else:
        base = np.asarray(base)
        if base.ndim == 0:
            base_arr = np.full((n,), float(base))
        elif base.ndim == 1 and base.shape[0] == n:
            base_arr = base.astype(float)
        else:
            raise ValueError(f"Unexpected base_values shape {base.shape}. For multiclass, build per-class tables.")

    # 3) Build feature DF in the *training* column order
    # expl.feature_names comes from X_to_explain; align to canonical list_vars
    if expl.feature_names is None:
        # assume expl.values columns already match list_vars
        feature_cols = list_vars
    else:
        # ensure all requested vars exist
        existing = list(expl.feature_names)
        missing = [c for c in list_vars if c not in existing]
        if missing:
            raise ValueError(f"These list_vars are missing from explanation features: {missing}")
        feature_cols = list_vars  # enforce this order

    df_features = pd.DataFrame(vals, columns=expl.feature_names)
    df_features = df_features[feature_cols]  # reorder

    # 4) Keys up front (robust expansion)
    if len(list_keys) != n:
        raise ValueError(f"list_keys length {len(list_keys)} != number of rows {n}")
    if list_keys_sale is not None and len(list_keys_sale) != n:
        raise ValueError(f"list_keys_sale length {len(list_keys_sale)} != number of rows {n}")

    if list_keys_sale is not None:
        df_keys = pd.DataFrame({"key": list_keys, "key_sale": list_keys_sale})
    else:
        df_keys = pd.DataFrame({"key": list_keys})

    # 5) Base value column (between keys and features)
    df_base = pd.DataFrame({"base_value": base_arr})

    # 6) Optional reconstructed prediction on the explained scale
    # (raw margin for classifiers unless you used model_output="probability")
    if include_pred:
        pred = base_arr + df_features.sum(axis=1).to_numpy()
        df_pred = pd.DataFrame({"contribution_sum": pred})
        df = pd.concat([df_keys, df_base, df_features, df_pred], axis=1)
    else:
        df = pd.concat([df_keys, df_base, df_features], axis=1)

    return df

ngboost_internals_ok

ngboost_internals_ok(regressor)

Return True if an NGBRegressor exposes the additive-ensemble internals SHAP needs.

NGBoost's exact tree-SHAP decomposition relies on per-stage base_models, scalings, col_idxs, learning_rate and init_params. This guards against future NGBoost layout changes — callers should skip SHAP (rather than crash) when it returns False.

Source code in openavmkit/shap_analysis.py

def ngboost_internals_ok(regressor) -> bool:
    """Return True if an NGBRegressor exposes the additive-ensemble internals SHAP needs.

    NGBoost's exact tree-SHAP decomposition relies on per-stage ``base_models``,
    ``scalings``, ``col_idxs``, ``learning_rate`` and ``init_params``. This guards
    against future NGBoost layout changes — callers should skip SHAP (rather than
    crash) when it returns False.
    """
    if not all(hasattr(regressor, a) for a in _NGBOOST_REQUIRED_ATTRS):
        return False
    bms = getattr(regressor, "base_models", None)
    if not bms or len(bms) == 0:
        return False
    # Each stage holds one fitted learner per distribution parameter (loc, logscale).
    if len(bms[0]) < 2:
        return False
    return True

plot_full_beeswarm

plot_full_beeswarm(explanation, title='SHAP Beeswarm', save_path=None, save_kwargs=None, wrap_width=20)

Plot a full SHAP beeswarm for a tree-based model with wrapped feature names.

This function wraps long feature names, auto-scales figure size to the number of features, and renders a beeswarm plot with rotated, smaller y-axis labels.

Parameters:

Name	Type	Description	Default
explanation	Explanation	SHAP Explanation object with values, base_values, data, and feature_names.	required
title	str	Title of the plot. Defaults to "SHAP Beeswarm".	'SHAP Beeswarm'
wrap_width	int	Maximum character width for feature name wrapping. Defaults to 20.	20
save_path	str	If provided, save the figure to this path (format inferred from extension). e.g., 'beeswarm.png', 'beeswarm.pdf', 'figs/beeswarm.svg'.	None
save_kwargs	dict	Extra kwargs passed to fig.savefig (e.g., {'dpi': 300, 'bbox_inches': 'tight', 'transparent': True}).	None

Source code in openavmkit/shap_analysis.py

def plot_full_beeswarm(
    explanation: shap.Explanation, 
    title: str = "SHAP Beeswarm", 
    save_path: str | None = None,
    save_kwargs: dict | None = None,
    wrap_width: int = 20
) -> None:
    """
    Plot a full SHAP beeswarm for a tree-based model with wrapped feature names.

    This function wraps long feature names, auto-scales figure size to the number of
    features, and renders a beeswarm plot with rotated, smaller y-axis labels.

    Parameters
    ----------
    explanation : shap.Explanation
        SHAP Explanation object with `values`, `base_values`, `data`, and `feature_names`.
    title : str, optional
        Title of the plot. Defaults to "SHAP Beeswarm".
    wrap_width : int, optional
        Maximum character width for feature name wrapping. Defaults to 20.
    save_path : str, optional
        If provided, save the figure to this path (format inferred from extension).
        e.g., 'beeswarm.png', 'beeswarm.pdf', 'figs/beeswarm.svg'.
    save_kwargs : dict, optional
        Extra kwargs passed to `fig.savefig` (e.g., {'dpi': 300, 'bbox_inches': 'tight',
        'transparent': True}).
    """
    if save_kwargs is None:
        save_kwargs = {"dpi": 300, "bbox_inches": "tight"}

    feature_names = explanation.feature_names
    if not feature_names:
        feature_names = []

    # Wrap feature names
    wrapped_names = [
        "\n".join(textwrap.wrap(fn, width=wrap_width))
        for fn in feature_names
    ]
    expl_wrapped = shap.Explanation(
        values=explanation.values,
        base_values=explanation.base_values,
        data=explanation.data,
        feature_names=wrapped_names,
    )

    # Determine figure size based on # features
    n_feats = len(wrapped_names)
    width = max(12, 0.3 * n_feats)
    height = max(6, 0.3 * n_feats)
    fig, ax = plt.subplots(figsize=(width, height), constrained_layout=True)

    # Draw the beeswarm (max_display defaults to all features here)
    shap.plots.beeswarm(expl_wrapped, max_display=n_feats, show=False)

    # Title + tweak y-labels
    ax.set_title(title)
    plt.setp(ax.get_yticklabels(), rotation=0, ha="right", fontsize=8)

    # Save if requested
    if save_path is not None:
        fig.savefig(save_path, **save_kwargs)

    plt.show()