Source code for modeva.models.kernel_xgb.api

"""
MoDeVa model wrappers for the GBDT leaf-kernel module.

MoGBDTKernelRegressor — five-head GNW on a fitted XGBoost regressor.
MoGBDTKernelClassifier — five-head GNW on a fitted XGBoost binary classifier.

All interpretability methods return ``ValidationResult`` objects following the
MoDeVa standard.  Use ``result.plot()`` to visualize and ``result.value`` to
access the raw numerical payload.

Default prediction head: ``exact_gbdt`` --- the per-tree hard leaf selector
with leaf-score values that recovers the base GBDT exactly (Theorem 1 of the
GBDT-GNW paper).  Other heads are user-selectable via ``predict(X, head=...)``
or ``predict_proba(X, head=...)``.
"""

import numpy as np
import pandas as pd
from scipy.special import expit as _sigmoid
from sklearn.base import RegressorMixin, ClassifierMixin

from ..base import ModelBaseRegressor, ModelBaseClassifier
from ...utils.results import ValidationResult
from .core import LeafKernelCore, _VALID_HEADS
from . import visualize as viz


class _LeafKernelInterpretMixin:
    """Shared interpretability surface for both regressor and classifier."""

    # ------------------------------------------------------------------
    # Explanation
    # ------------------------------------------------------------------

    def explain_local(self, X, sample_index=0, feature_names=None):
        """Local TKLE evidence ledger for a single query.

        Returns
        -------
        ValidationResult
            ``result.value`` contains the full ledger.  ``result.table`` is a
            top-20-neighbor table.  ``result.plot('weight_bar')`` and
            ``result.plot('neighbor_scatter')`` show the prediction's evidence.
        """
        if feature_names is None:
            feature_names = getattr(self, "feature_names_", None)
        ledger = self._core.explain_local(X, feature_names=feature_names)

        s = int(sample_index)
        order = np.argsort(-ledger["topk_weights"][s])
        top = min(20, len(order))
        sel = order[:top]
        table = pd.DataFrame({
            "Neighbor": [f"n{int(ledger['topk_idx'][s, j])}" for j in sel],
            "Weight": ledger["topk_weights"][s, sel],
            "K(x,x_i)": ledger["topk_K"][s, sel],
            "y_i": ledger["topk_y"][s, sel],
            "f_M(x_i)": ledger["topk_fhat"][s, sel],
        })

        options = viz.viz_local_evidence(ledger, sample_index=s)

        return ValidationResult(
            key="gbdt_kernel_local",
            model=self.name,
            inputs={"sample_index": s, "feature_names": feature_names},
            value=ledger,
            table=table,
            options=options,
        )

    def get_evidence_diagnostics(self, X):
        """Six TKLE diagnostics (neff, Delta_y, Delta_q, G_q, Delta_K, C_cal)
        across a query batch.

        Returns
        -------
        ValidationResult
            ``result.table`` summarizes each diagnostic's distribution;
            ``result.plot('neff')`` etc. show individual histograms.
        """
        ledger = self._core.explain_local(X, feature_names=None)
        diags = {k: np.asarray(ledger[k]) for k in
                 ["neff", "delta_y", "delta_q", "G_q", "delta_K", "C_cal"]}
        table = pd.DataFrame({
            "Diagnostic": list(diags.keys()),
            "Mean": [float(v.mean()) for v in diags.values()],
            "Median": [float(np.median(v)) for v in diags.values()],
            "p10": [float(np.percentile(v, 10)) for v in diags.values()],
            "p90": [float(np.percentile(v, 90)) for v in diags.values()],
        })
        options = viz.viz_evidence_diagnostics(ledger)
        return ValidationResult(
            key="gbdt_kernel_evidence_diagnostics",
            model=self.name,
            inputs=None,
            value=ledger,
            table=table,
            options=options,
        )

    def importance_global(self, X, feature_names=None):
        """Kernel-weighted feature importance: features whose neighbors are
        tightly concentrated (small within-neighborhood dispersion) are
        deemed important under the learned geometry.
        """
        if feature_names is None:
            feature_names = getattr(self, "feature_names_", None)
        imp = self._core.importance_global(X, feature_names=feature_names)
        table = pd.DataFrame({
            "Name": imp["feature_names"],
            "Importance": imp["importance"],
        }).sort_values("Importance", ascending=False).reset_index(drop=True)
        options = viz.viz_importance_global(imp)
        return ValidationResult(
            key="gbdt_kernel_importance",
            model=self.name,
            inputs={"feature_names": feature_names},
            value=imp,
            table=table,
            options=options,
        )

    def get_neighbor_analysis(self, X, sample_index=0):
        """Top-k neighbor table for a single query."""
        na = self._core.get_neighbor_analysis(X, sample_index=sample_index)
        table = pd.DataFrame({
            "Neighbor": [f"n{int(i)}" for i in na["neighbor_idx"]],
            "Weight": na["neighbor_weights"],
            "K(x,x_i)": na["neighbor_K"],
            "y_i": na["neighbor_y"],
            "f_M(x_i)": na["neighbor_fhat"],
        }).sort_values("Weight", ascending=False).reset_index(drop=True)
        options = viz.viz_neighbor_analysis(na)
        return ValidationResult(
            key="gbdt_kernel_neighbors",
            model=self.name,
            inputs={"sample_index": sample_index},
            value=na,
            table=table,
            options=options,
        )

    # ------------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------------

    def nominate_head(self, X_val, y_val):
        """Run the GNW-paper diagnostics and recommend a head."""
        diag = self._core.nominate_head(X_val, y_val)
        head_table = pd.DataFrame({
            "Head": list(diag["head_errors"].keys()),
            "Val error": list(diag["head_errors"].values()),
        }).sort_values("Val error").reset_index(drop=True)
        diag_table = pd.DataFrame([{
            "B_R (region bias)": diag["B_R"],
            "V_R (region variance)": diag["V_R"],
            "df(lambda) (KRR DoF)": diag["df_lambda"],
            "C(10) (spectral conc.)": diag["C_m"],
            "Recommended head": diag["recommended_head"],
        }])
        options = viz.viz_head_diagnostics(diag)
        return ValidationResult(
            key="gbdt_kernel_head_nomination",
            model=self.name,
            inputs=None,
            value=diag,
            table={"errors": head_table, "diagnostics": diag_table},
            options=options,
        )

    def diagnose_weakness(self, top_features=8, n_bins=10):
        """Per-cluster weakness scores + JS feature profile (paper_cluster_gated §4)."""
        diag = self._core.diagnose_weakness(
            top_features=top_features, n_bins=n_bins
        )
        feature_names = getattr(self, "feature_names_", None)
        js_df = diag["js_table"].copy()
        if feature_names is not None:
            js_df["feature"] = [feature_names[int(i)]
                                for i in js_df["feature_index"]]
        else:
            js_df["feature"] = js_df["feature_index"]
        diag["js_table"] = js_df
        options = viz.viz_weakness(diag)
        return ValidationResult(
            key="gbdt_kernel_weakness",
            model=self.name,
            inputs={"n_bins": n_bins, "top_features": top_features},
            value=diag,
            table={"clusters": diag["cluster_table"], "js": js_df},
            options=options,
        )

    # ------------------------------------------------------------------
    # ICL view
    # ------------------------------------------------------------------

    def predict_with_context(self, X_query, context_X, context_y, head="gnw_label"):
        """In-context prediction: use the supplied ``(context_X, context_y)``
        pool as the NW memory instead of the original training set.

        The base kernel geometry is unchanged (it is the fitted GBDT's leaf
        kernel); only the keys/values are replaced by the user-provided
        context.  Useful for "what if I had observed this set of cases?"
        explanations.

        Parameters
        ----------
        X_query : ndarray
            Queries to predict.
        context_X : ndarray
            Context features (acts as the NW memory).
        context_y : ndarray
            Context labels.
        head : {"gnw_label", "gnw_leaf"}, default="gnw_label"
            Only neighbor-based heads support an external context.

        Returns
        -------
        predictions : ndarray
        weights : ndarray of shape (n_query, k)
            Normalized kernel weights against the context pool.
        idx : ndarray of shape (n_query, k)
            Indices into ``context_X`` of the top-k retrieved cases.
        """
        if head not in ("gnw_label", "gnw_leaf"):
            raise ValueError("ICL context only supports neighbor heads "
                             "('gnw_label' or 'gnw_leaf').")
        X_query = np.asarray(X_query, dtype=np.float64)
        context_X = np.asarray(context_X, dtype=np.float64)
        context_y = np.asarray(context_y, dtype=np.float64).ravel()

        # Swap training data temporarily
        saved_X = self._core._X_train
        saved_y = self._core._y_train
        saved_leaf = self._core._leaf_idx_train
        saved_fhat = self._core._fhat_train

        self._core._X_train = context_X
        self._core._y_train = context_y
        self._core._leaf_idx_train = (
            self._core._xgb.apply(context_X).astype(np.int64)
        )
        self._core._fhat_train = self._core._predict_exact_gbdt(context_X)

        try:
            idx, w = self._core._normalized_weights(X_query)
            if head == "gnw_label":
                pred = (w * context_y[idx]).sum(axis=1)
            else:
                pred = (w * self._core._fhat_train[idx]).sum(axis=1)
        finally:
            self._core._X_train = saved_X
            self._core._y_train = saved_y
            self._core._leaf_idx_train = saved_leaf
            self._core._fhat_train = saved_fhat

        return pred, w, idx


# ============================================================================
# Regressor
# ============================================================================


[docs] class MoGBDTKernelRegressor(_LeafKernelInterpretMixin, RegressorMixin, ModelBaseRegressor): """GBDT-as-learned-kernel regressor with five interchangeable heads. Wraps a pre-trained tree ensemble and exposes prediction heads on the induced leaf kernel. Default head ``exact_gbdt`` recovers the base GBDT prediction exactly via the generalized Nadaraya-Watson representation (Theorem 1 of the GBDT-GNW paper). Parameters ---------- base_model : fitted XGBoost regressor (or MoXGBRegressor wrapper) Pre-trained ensemble whose leaf geometry will be re-used. name : str, optional Identifier. Default: ``"GBDTKernel"``. kernel_topk : int, default=200 Number of training neighbors used by the NW heads. kernel_rho : float, default=1.0 Sharpening exponent: ``w_i \\propto K(x, x_i)^rho``. ridge_lambda : float, default=1.0 KRR regularization on the leaf one-hot basis. n_clusters : int, default=8 Number of behavioral cohorts for cluster-gated residual repair. gate_mode : {"defensive", "adaptive"}, default="defensive" nystrom_landmarks : int, default=300 residual_gamma_max : float, default=1.5 gate_tmin : float, default=1.96 shrink_tau : float, default=20.0 random_state : int, default=0 """ def __init__(self, base_model, name=None, kernel_topk=200, kernel_rho=1.0, ridge_lambda=1.0, n_clusters=8, gate_mode="defensive", nystrom_landmarks=300, residual_gamma_max=1.5, gate_tmin=1.96, shrink_tau=20.0, random_state=0): self.base_model = base_model self.name = name or "GBDTKernel" self.kernel_topk = kernel_topk self.kernel_rho = kernel_rho self.ridge_lambda = ridge_lambda self.n_clusters = n_clusters self.gate_mode = gate_mode self.nystrom_landmarks = nystrom_landmarks self.residual_gamma_max = residual_gamma_max self.gate_tmin = gate_tmin self.shrink_tau = shrink_tau self.random_state = random_state self._core = None self.feature_names_ = None
[docs] def fit(self, X, y, sample_weight=None, X_val=None, y_val=None, feature_names=None, verbose=False): X, y, _ = self._validate_fit_inputs(X, y, sample_weight) self.feature_names_ = (list(feature_names) if feature_names is not None else [f"x{j}" for j in range(X.shape[1])]) self._core = LeafKernelCore( base_model=self.base_model, classification=False, kernel_topk=self.kernel_topk, kernel_rho=self.kernel_rho, ridge_lambda=self.ridge_lambda, n_clusters=self.n_clusters, nystrom_landmarks=self.nystrom_landmarks, residual_gamma_max=self.residual_gamma_max, gate_mode=self.gate_mode, gate_tmin=self.gate_tmin, shrink_tau=self.shrink_tau, random_state=self.random_state, ) self._core.fit(X, y, X_val=X_val, y_val=y_val, verbose=verbose) return self
[docs] def predict(self, X, head="exact_gbdt"): """Predict in target space using the chosen head.""" return self._core.predict(X, head=head)
def _predict(self, X): # Required by ModelBase: always returns the exact-GBDT head. return self._core.predict(X, head="exact_gbdt")
# ============================================================================ # Classifier # ============================================================================
[docs] class MoGBDTKernelClassifier(_LeafKernelInterpretMixin, ClassifierMixin, ModelBaseClassifier): """GBDT-as-learned-kernel binary classifier with five heads. Same parameters and head menu as :class:`MoGBDTKernelRegressor`. All predictions are produced in logit space internally; ``predict_proba`` applies the sigmoid. """ def __init__(self, base_model, name=None, kernel_topk=200, kernel_rho=1.0, ridge_lambda=1.0, n_clusters=8, gate_mode="defensive", nystrom_landmarks=300, residual_gamma_max=1.5, gate_tmin=1.96, shrink_tau=20.0, random_state=0): self.base_model = base_model self.name = name or "GBDTKernel-Cls" self.kernel_topk = kernel_topk self.kernel_rho = kernel_rho self.ridge_lambda = ridge_lambda self.n_clusters = n_clusters self.gate_mode = gate_mode self.nystrom_landmarks = nystrom_landmarks self.residual_gamma_max = residual_gamma_max self.gate_tmin = gate_tmin self.shrink_tau = shrink_tau self.random_state = random_state self._core = None self.feature_names_ = None
[docs] def fit(self, X, y, sample_weight=None, X_val=None, y_val=None, feature_names=None, verbose=False): X, y, _ = self._validate_fit_inputs(X, y, sample_weight) self.feature_names_ = (list(feature_names) if feature_names is not None else [f"x{j}" for j in range(X.shape[1])]) self._core = LeafKernelCore( base_model=self.base_model, classification=True, kernel_topk=self.kernel_topk, kernel_rho=self.kernel_rho, ridge_lambda=self.ridge_lambda, n_clusters=self.n_clusters, nystrom_landmarks=self.nystrom_landmarks, residual_gamma_max=self.residual_gamma_max, gate_mode=self.gate_mode, gate_tmin=self.gate_tmin, shrink_tau=self.shrink_tau, random_state=self.random_state, ) self._core.fit(X, y, X_val=X_val, y_val=y_val, verbose=verbose) return self
[docs] def decision_function(self, X, head="exact_gbdt"): """Margin / logit prediction under the chosen head.""" return self._core.predict(X, head=head)
[docs] def predict_proba(self, X, head="exact_gbdt"): """Class probabilities under the chosen head.""" raw = self._core.predict(X, head=head) if head == "gnw_label": # For label head with binary y, raw is already in probability space prob_1 = np.clip(raw, 0.0, 1.0) else: prob_1 = _sigmoid(raw) return np.column_stack([1.0 - prob_1, prob_1])
[docs] def predict(self, X, head="exact_gbdt"): proba = self.predict_proba(X, head=head) return (proba[:, 1] > 0.5).astype(int)
def _predict_proba(self, X): # Required by ModelBase return self.predict_proba(X, head="exact_gbdt")