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)

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,
):
    """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...")

    return {
        "train": explain(X_train, "train"),
        "test": explain(X_test, "test"),
        "sales": explain(X_sales, "sales"),
        "universe": explain(X_univ, "universe"),
    }

get_full_model_shaps

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

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

Returns:

Type	Description
dict	A dict containing shap.Explanation objects 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
):
    """
    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.

    Returns
    -------
    dict
        A dict containing shap.Explanation objects 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
        )

    # 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
        )

    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...")

    shap_sales = _shap_explain(model_type, tree_explainer, X_sales, cat_data=cat_data, approximate=approximate, verbose=verbose, label="sales")
    shap_train = _shap_explain(model_type, tree_explainer, X_train, cat_data=cat_data, approximate=approximate, verbose=verbose, label="train")
    shap_test  = _shap_explain(model_type, tree_explainer, X_test,  cat_data=cat_data, approximate=approximate, verbose=verbose, label="test")
    shap_univ  = _shap_explain(model_type, tree_explainer, X_univ,  cat_data=cat_data, approximate=approximate, verbose=verbose, label="universe")

    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)

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,
):
    """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}...")

    return {
        "train": explain(X_train, "train"),
        "test": explain(X_test, "test"),
        "sales": explain(X_sales, "sales"),
        "universe": explain(X_univ, "universe"),
    }

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()