tanml 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

tanml/checks/explainability/shap_check.py

@@ -1,55 +1,287 @@
+# tanml/checks/explainability/shap_check.py
 from tanml.checks.base import BaseCheck
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
 import shap
+import numpy as np
 import pandas as pd
-import matplotlib.pyplot as plt
 import traceback
+import warnings
 from pathlib import Path
 from datetime import datetime
+from scipy import sparse as sp
+
+def _safe_numeric_cast_df(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert columns to numeric where possible; leave others unchanged."""
+    out = {}
+    for c in df.columns:
+        s = df[c]
+        try:
+            out[c] = pd.to_numeric(s)  # no 'errors' kwarg
+        except Exception:
+            out[c] = s
+    return pd.DataFrame(out, index=df.index)
 
 
 class SHAPCheck(BaseCheck):
+    """
+    SHAP for regression + binary classification (no multiclass).
+
+    Config under rule_config["explainability"]["shap"]:
+      - enabled: bool (default True, also checked under rule_config["SHAPCheck"]["enabled"])
+      - task: "auto" | "classification" | "regression"          (default "auto")
+      - algorithm: "auto" | "tree" | "linear" | "kernel" | "permutation" (default "auto")
+      - model_output: "auto" | "raw" | "log_odds" | "probability"  (tree-only hint; default "auto")
+      - background_strategy: "sample" | "kmeans"                 (default "sample")
+      - background_sample_size: int                              (default 100)
+      - test_sample_size: int                                    (default 200)
+      - max_display: int                                         (default 20)
+      - seed: int                                                (default 42)
+      - out_dir: str (optional)                                  (preferred save folder)
+    """
+
     def __init__(self, model, X_train, X_test, y_train, y_test, rule_config=None, cleaned_df=None):
         super().__init__(model, X_train, X_test, y_train, y_test, rule_config, cleaned_data=cleaned_df)
         self.cleaned_df = cleaned_df
 
-    def run(self):
-        result = {}
+    # -------------------------- helpers --------------------------
 
+    @staticmethod
+    def _to_df(X, names=None):
+        if isinstance(X, pd.DataFrame):
+            return X
+        if sp.issparse(X):
+            df = pd.DataFrame.sparse.from_spmatrix(X)
+        else:
+            df = pd.DataFrame(np.asarray(X))
+        if names is not None and len(names) == df.shape[1]:
+            df.columns = list(names)
+        return df
+
+    @staticmethod
+    def _task(y, forced="auto"):
+        if forced and forced != "auto":
+            return forced
         try:
-            expl_cfg = self.rule_config.get("explainability", {})
-            bg_n = expl_cfg.get("background_sample_size", 100)
-            test_n = expl_cfg.get("test_sample_size", 200)
+            yv = y.iloc[:, 0] if isinstance(y, pd.DataFrame) else (y if isinstance(y, pd.Series) else pd.Series(y))
+            uniq = pd.Series(yv).dropna().unique()
+            return "classification" if len(uniq) <= 2 else "regression"
+        except Exception:
+            return "regression"
+
+    @staticmethod
+    def _pos_cls_idx(model, X_one):
+        """Return index of the positive class (1/True if available, else max-label) for binary classification."""
+        try:
+            if hasattr(model, "classes_") and len(model.classes_) == 2:
+                classes = list(model.classes_)
+                for pos in (1, True):
+                    if pos in classes:
+                        return classes.index(pos)
+                return classes.index(max(classes))
+            if hasattr(model, "predict_proba"):
+                proba = model.predict_proba(X_one)
+                return 1 if proba.shape[1] == 2 else 0
+        except Exception:
+            pass
+        return 1
+
+    @staticmethod
+    def _looks_like_tree(m):
+        mod = type(m).__module__.lower()
+        name = type(m).__name__.lower()
+        return (
+            "xgboost" in mod
+            or "lightgbm" in mod
+            or "catboost" in mod
+            or "sklearn.ensemble" in mod
+            or "sklearn.tree" in mod
+            or "randomforest" in name
+            or "gradientboost" in name
+            or "extratrees" in name
+            or "decisiontree" in name
+        )
+
+    @staticmethod
+    def _looks_like_linear(m):
+        mod = type(m).__module__.lower()
+        name = type(m).__name__.lower()
+        return (
+            "sklearn.linear_model" in mod
+            or "logistic" in name
+            or "linear" in name
+            or "ridge" in name
+            or "lasso" in name
+            or "elastic" in name
+        )
+
+    def _predict_fn(self, is_cls: bool, pos_idx: int | None):
+        """
+        Vectorized prediction function for permutation/kernel explainers.
+        Returns positive-class probability when classification is detected and predict_proba is available.
+        """
+        if is_cls and hasattr(self.model, "predict_proba"):
+            def f(X):
+                p = self.model.predict_proba(X)
+                i = 1 if (p.ndim == 2 and p.shape[1] == 2) else (pos_idx or 0)
+                return p[:, i]
+            return f
+        return self.model.predict
+
+    def _explainer(self, algorithm, background, model_output_hint, is_cls, pos_idx):
+        """
+        Choose fastest viable explainer:
+          - tree → TreeExplainer (with interventional perturbation)
+          - linear → LinearExplainer
+          - default/auto → PermutationExplainer (avoid slow Kernel, unless explicitly requested)
+        """
+        m = self.model
+        alg = (algorithm or "auto").lower()
+
+        # Prefer fast paths in auto
+        if alg == "tree" or (alg == "auto" and self._looks_like_tree(m)):
+            mo = None if model_output_hint == "auto" else model_output_hint
+            expl = shap.TreeExplainer(
+                m, data=background, feature_perturbation="interventional", model_output=mo
+            )
+            return expl, "tree"
 
-            X_sample = self.X_test[:test_n]
-            background = shap.utils.sample(self.X_train, bg_n, random_state=42)
+        if alg == "linear" or (alg == "auto" and self._looks_like_linear(m)):
+            return shap.LinearExplainer(m, background), "linear"
 
-            X_sample = pd.DataFrame(X_sample)
-            background = pd.DataFrame(background)
+        if alg == "permutation" or alg == "auto":
+            fn = self._predict_fn(is_cls, pos_idx)
+            return shap.explainers.Permutation(fn, background, max_evals=2000), "perm"
 
-            explainer = shap.Explainer(self.model, background)
-            shap_exp = explainer(X_sample)
+        # Only use Kernel if explicitly requested
+        if alg == "kernel":
+            fn = self._predict_fn(is_cls, pos_idx)
+            return shap.KernelExplainer(fn, background), "kernel"
 
-            if shap_exp.values.ndim == 3:
-                shap_exp.values = shap_exp.values[:, :, 1]
-                shap_exp.base_values = shap_exp.base_values[:, 1]
+        # Fallback (should not hit)
+        fn = self._predict_fn(is_cls, pos_idx)
+        return shap.explainers.Permutation(fn, background, max_evals=2000), "perm"
 
-            segment = self.rule_config.get("meta", {}).get("segment", "global")
+    # ---------------------------- main ----------------------------
+
+    def run(self):
+        out = {}
+        try:
+            warnings.filterwarnings("ignore", category=UserWarning)
+
+            # -------- read config (note: under explainability.shap) ----------
+            exp_cfg = (self.rule_config or {}).get("explainability", {}) or {}
+            cfg = exp_cfg.get("shap", {}) if isinstance(exp_cfg, dict) else {}
+            seed = int(cfg.get("seed", 42))
+            bg_n = int(cfg.get("background_sample_size", 100))
+            test_n = int(cfg.get("test_sample_size", 200))
+            task_forced = (cfg.get("task") or "auto").lower()
+            algorithm = (cfg.get("algorithm") or "auto").lower()
+            bg_strategy = (cfg.get("background_strategy") or "sample").lower()
+            model_output_hint = (cfg.get("model_output") or "auto").lower()
+            max_display = int(cfg.get("max_display", 20))
+
+            # -------- resolve output directory + timestamp ----------
+            out_dir_opt = cfg.get("out_dir")
+            options_dir = ((self.rule_config or {}).get("options") or {}).get("save_artifacts_dir")
+            # prefer explicit shap.out_dir, then global artifacts dir, then local fallback
+            outdir = Path(out_dir_opt or options_dir or (Path(__file__).resolve().parents[2] / "tmp_report_assets"))
+            outdir.mkdir(parents=True, exist_ok=True)
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-            output_path = Path(f"reports/images/shap_summary_{segment}_{ts}.png")
-            output_path.parent.mkdir(parents=True, exist_ok=True)
 
-            plt.figure(figsize=(8, 6))
-            shap.plots.beeswarm(shap_exp, show=False)
-            plt.savefig(output_path, bbox_inches="tight")
+            # -------- materialize dataframes & coerce dtypes ----------
+            feature_names = list(self.X_train.columns) if isinstance(self.X_train, pd.DataFrame) else None
+            Xtr = self._to_df(self.X_train, feature_names)
+            Xte = self._to_df(self.X_test, feature_names)
+            if Xtr.empty or Xte.empty:
+                raise ValueError("Empty X_train or X_test for SHAP.")
+
+            # Avoid slow implicit object->float conversions during plotting
+            Xtr = _safe_numeric_cast_df(Xtr)
+            Xte = _safe_numeric_cast_df(Xte)
+
+
+            # -------- task resolution & sanity for binary classification ----------
+            task = self._task(self.y_train, forced=task_forced)
+            is_cls = (task == "classification")
+            if is_cls:
+                yv = self.y_train if isinstance(self.y_train, (pd.Series, pd.DataFrame)) else pd.Series(self.y_train)
+                if len(pd.Series(yv).dropna().unique()) > 2:
+                    raise ValueError("Binary classification only: y_train has >2 classes.")
+
+            # positive-class index hint (used for permutation/kernel predict function)
+            pos_idx_hint = self._pos_cls_idx(self.model, Xte.iloc[:1].values)
+
+            # -------- background selection ----------
+            if bg_strategy == "kmeans" and len(Xtr) > bg_n and not sp.issparse(self.X_train):
+                background = shap.kmeans(Xtr, bg_n, seed=seed)
+            else:
+                background = shap.utils.sample(Xtr, bg_n, random_state=seed)
+
+            # -------- slice test rows to explain ----------
+            Xs = Xte.head(test_n)
+
+            # -------- choose explainer & compute SHAP once ----------
+            explainer, kind = self._explainer(algorithm, background, model_output_hint, is_cls, pos_idx_hint)
+            if kind == "tree":
+                sv = explainer(Xs, check_additivity=False)  # big speedup, visually identical plots
+            else:
+                sv = explainer(Xs)
+
+            bg_shape = background.shape if hasattr(background, "shape") else None
+            print(f"SHAP explainer={type(explainer).__name__} kind={kind} Xs={Xs.shape} "
+                  f"bg={'kmeans' if bg_shape is None else bg_shape}")
+
+            # -------- squeeze to 2-D for binary cls (if needed) ----------
+            vals = sv.values
+            if hasattr(vals, "ndim") and vals.ndim == 3:
+                pos_idx = self._pos_cls_idx(self.model, Xs.iloc[:1].values)
+                sv.values = vals[:, :, pos_idx]
+                if isinstance(sv.base_values, np.ndarray) and sv.base_values.ndim == 2:
+                    sv.base_values = sv.base_values[:, pos_idx]
+            else:
+                pos_idx = None if task == "regression" else self._pos_cls_idx(self.model, Xs.iloc[:1].values)
+
+            # -------- save plots ----------
+            segment = "global"
+
+            beeswarm_path = outdir / f"shap_beeswarm_{segment}_{ts}.png"
+            plt.figure(figsize=(9, 6))
+            shap.plots.beeswarm(sv, max_display=max_display, show=False)
+            plt.tight_layout()
+            plt.savefig(beeswarm_path, bbox_inches="tight", dpi=120, transparent=False)
             plt.close()
 
-            print(f"✅ SHAP plot saved at: {output_path}")
-            result["shap_plot_path"] = str(output_path)
-            result["status"] = "SHAP plot generated successfully"
+            bar_path = outdir / f"shap_bar_{segment}_{ts}.png"
+            plt.figure(figsize=(9, 6))
+            shap.plots.bar(sv, max_display=max_display, show=False)
+            plt.tight_layout()
+            plt.savefig(bar_path, bbox_inches="tight", dpi=120, transparent=False)
+            plt.close()
 
-        except Exception:
-            err = traceback.format_exc()
-            print(f"⚠️ SHAPCheck failed:\n{err}")
-            result["status"] = f"SHAP plot failed:\n{err}"
+            # -------- top features ----------
+            # mean absolute SHAP across rows
+            mean_abs = np.abs(sv.values).mean(axis=0)
+            idx = np.argsort(mean_abs)[::-1][:max_display]
+            cols = list(Xs.columns)
+            top = [{"feature": cols[i] if i < len(cols) else f"f{i}", "mean_abs_shap": float(mean_abs[i])} for i in idx]
+            top_list_pairs = [[d["feature"], d["mean_abs_shap"]] for d in top]  # compat with old report builder
 
-        return result
+            out.update({
+                "status": "ok",
+                "task": task,
+                "positive_class_index": pos_idx if task == "classification" else None,
+                # new + old keys (backward-compatible)
+                "plots": {"beeswarm": str(beeswarm_path), "bar": str(bar_path)},
+                "images": {"beeswarm": str(beeswarm_path), "bar": str(bar_path)},
+                "top_features": top,
+                "shap_top_features": top_list_pairs,
+            })
+            print(f"✅ SHAP saved: {beeswarm_path}, {bar_path}")
+
+        except Exception:
+            out["status"] = "error: " + traceback.format_exc()
+        return out

tanml/checks/logit_stats.py

@@ -1,59 +1,191 @@
 # tanml/checks/logit_stats.py
-
+import numpy as np
+import pandas as pd
 import statsmodels.api as sm
+from contextlib import suppress
+
+# ---------- helpers: winsorize + standardize (no sklearn) ----------
+def _winsorize_df(df: pd.DataFrame, low_q=0.005, high_q=0.995) -> pd.DataFrame:
+    with suppress(Exception):
+        q_low = df.quantile(low_q)
+        q_hi = df.quantile(high_q)
+        df = df.clip(lower=q_low, upper=q_hi, axis=1)
+    return df
+
+def _standardize_df(df: pd.DataFrame) -> pd.DataFrame:
+    mu = df.mean()
+    sd = df.std(ddof=0).replace(0, np.nan)
+    z = (df - mu) / sd
+    return z.fillna(0.0)
+
+def _prep_design_matrix_df(
+    X: pd.DataFrame,
+    ref_columns: pd.Index | None = None,
+    add_const: bool = True
+) -> pd.DataFrame:
+    """
+    Consistent design prep for statsmodels:
+      - ensure DataFrame
+      - one-hot encode categoricals (drop_first=True)
+      - align to ref_columns (if provided; excludes 'const')
+      - coerce to numeric, replace inf, fillna(0)
+      - drop zero-variance columns
+      - winsorize tails, standardize
+      - optionally add intercept
+    """
+    Xd = pd.DataFrame(X).copy()
+    Xd = pd.get_dummies(Xd, drop_first=True)
+
+    if ref_columns is not None:
+        ref_wo_const = [c for c in ref_columns if c != "const"]
+        Xd = Xd.reindex(columns=ref_wo_const, fill_value=0.0)
+
+    # numeric + clean
+    for c in Xd.columns:
+        Xd[c] = pd.to_numeric(Xd[c], errors="coerce")
+    Xd = Xd.replace([np.inf, -np.inf], np.nan).fillna(0.0)
 
-class LogisticStatsCheck:
-    def __init__(self, model, X, y, config):
-        self.model = model
-        self.X = X
-        self.y = y
-        self.config = config or {}
+    # drop zero-variance columns (except 'const' which isn't present yet)
+    var = Xd.var(numeric_only=True)
+    drop_cols = list(var.index[var == 0])
+    if drop_cols:
+        Xd = Xd.drop(columns=drop_cols, errors="ignore")
 
-    def run(self):
+    # stabilize logits: winsorize + standardize
+    Xd = _winsorize_df(Xd, 0.005, 0.995)
+    Xd = _standardize_df(Xd)
+
+    if add_const:
+        Xd = sm.add_constant(Xd, has_constant="add")
+    return Xd
+
+def _encode_binary_target(y) -> pd.Series:
+    yb = pd.Series(y)
+    uniq = pd.unique(yb.dropna())
+    try:
+        cats = pd.Series(uniq).astype("category").cat.categories
+        if set(map(int, cats)) == {0, 1}:
+            return yb.astype(int)
+    except Exception:
+        pass
+    if len(uniq) != 2:
+        raise ValueError("Logit requires a binary target with exactly two classes.")
+    # map majority -> 0, minority -> 1 (stable ordering)
+    counts = yb.value_counts().sort_values(ascending=False).index.tolist()
+    return yb.map({counts[0]: 0, counts[1]: 1}).astype(int)
+
+def _fit_stage1_ridge(yb: pd.Series, Xd: pd.DataFrame, alpha_grid=(1.0, 0.3, 0.1, 0.03, 0.01)):
+    """
+    Stage-1: ridge-regularized logit for stability. Returns best result and alpha.
+    """
+    best = None
+    for a in alpha_grid:
         try:
-            # 1) Add constant and fit the statsmodels Logit
-            Xc = sm.add_constant(self.X, has_constant='add')
-            res = sm.Logit(self.y, Xc).fit(disp=False)
-
-            # 2) Extract coefficient table
-            coef    = res.params
-            stderr  = res.bse
-            zscore  = coef / stderr
-            pvals   = res.pvalues
-
-            table = []
-            for feat in coef.index:
-                label = "Intercept" if feat.lower() == "const" else feat
-                table.append({
-                    "feature":     label,
-                    "coefficient": float(coef[feat]),
-                    "std_error":   float(stderr[feat]),
-                    "z_score":     float(zscore[feat]),
-                    "p_value":     float(pvals[feat]),
-                })
-
-            # 3) Fit statistics
-            fit = {
-                "log_lik":   float(res.llf),
-                "aic":       float(res.aic),
-                "bic":       float(res.bic),
-                "pseudo_r2": float(res.prsquared),
-            }
-
-            # 4) Full summary text
-            summary = res.summary().as_text()
-
-            return {
-                "table":   table,
-                "fit":     fit,
-                "summary": summary,
-                "object":  res
-            }
-
-        except Exception as e:
-            return {
-                "table":   [],
-                "fit":     {},
-                "summary": f"LogisticStatsCheck failed: {e}",
-                "object":  None
-            }
+            res = sm.Logit(yb, Xd).fit_regularized(alpha=a, L1_wt=0.0, maxiter=2000, disp=False)
+            coef = res.params.values
+            if not np.all(np.isfinite(coef)):
+                continue
+            score = float(getattr(res, "llf", -np.inf))
+            cand = {"alpha": a, "res": res, "score": score}
+            if best is None or cand["score"] > best["score"] + 1e-9 or \
+               (abs(cand["score"] - best["score"]) <= 1e-9 and a < best["alpha"]):
+                best = cand
+        except Exception:
+            continue
+    if best is None:
+        # stronger ridge last-ditch
+        res = sm.Logit(yb, Xd).fit_regularized(alpha=3.0, L1_wt=0.0, maxiter=2000, disp=False)
+        best = {"alpha": 3.0, "res": res, "score": float(getattr(res, "llf", -np.inf))}
+    return best
+
+def _fit_stage2_inference(yb: pd.Series, Xd: pd.DataFrame, start_params: np.ndarray):
+    """
+    Stage-2: try unpenalized MLE (for classic p-values). If it fails, fall back to GLM Binomial with robust SEs.
+    Returns (model_tag, results).
+    """
+    # Unpenalized MLE with good starts
+    try:
+        mle = sm.Logit(yb, Xd).fit(start_params=start_params, method="lbfgs", maxiter=5000, disp=False)
+        return ("logit_mle", mle)
+    except Exception:
+        pass
+
+    # Robust GLM fallback
+    glm = sm.GLM(yb, Xd, family=sm.families.Binomial())
+    glm_res = glm.fit(cov_type="HC3")  # robust/sandwich SEs
+    return ("glm_robust", glm_res)
+
+def compute_logit_stats(X: pd.DataFrame, y) -> dict:
+    """
+    Train-only stats summary (Sections a & b):
+      - robust design prep (winsorize + z-score)
+      - Stage-1 ridge to stabilize
+      - Stage-2 unpenalized MLE for inference (fallback to GLM robust)
+      - returns summary text and a tidy coef table
+    """
+    # design + target
+    Xd = _prep_design_matrix_df(X, ref_columns=None, add_const=True)
+    yb = _encode_binary_target(y)
+
+    # Stage-1: ridge (stability)
+    stg1 = _fit_stage1_ridge(yb, Xd, alpha_grid=(1.0, 0.3, 0.1, 0.03, 0.01))
+    start = stg1["res"].params.values
+
+    # Stage-2: inference
+    model_tag, res = _fit_stage2_inference(yb, Xd, start_params=start)
+
+    # summary text
+    with suppress(Exception):
+        # summary2() is prettier when available
+        summary_text = res.summary2().as_text()
+    if 'summary_text' not in locals():
+        summary_text = str(res.summary())
+
+    # coef table
+    params = res.params
+    bse = res.bse
+    # z or t values depending on object; label as z
+    if hasattr(res, "tvalues"):
+        stat_vals = res.tvalues
+    elif hasattr(res, "zvalues"):
+        stat_vals = res.zvalues
+    else:
+        stat_vals = params / bse.replace(0, np.nan)
+
+    pvals = res.pvalues
+    ci = res.conf_int(alpha=0.05)
+    ci.columns = ["ci_low", "ci_high"]
+
+    coef_df = pd.DataFrame({
+        "feature": params.index,
+        "coef": params.values,
+        "std err": bse.values,
+        "z": stat_vals.values,
+        "P>|z|": pvals.values,
+        "ci_low": ci["ci_low"].values,
+        "ci_high": ci["ci_high"].values,
+    })
+
+    # const first
+    if "const" in coef_df["feature"].values:
+        coef_df = pd.concat(
+            [coef_df.loc[coef_df["feature"] == "const"],
+             coef_df.loc[coef_df["feature"] != "const"].sort_values("feature")],
+            ignore_index=True
+        )
+
+    # round to 6 decimals so small effects don't show up as 0.000000
+    for c in ["coef","std err","z","P>|z|","ci_low","ci_high"]:
+        coef_df[c] = pd.to_numeric(coef_df[c], errors="coerce").round(6)
+
+    # annotate header line to reflect fallback if used
+    if model_tag == "glm_robust":
+        summary_text = "Results: GLM Binomial (robust SE)\n" + summary_text
+    else:
+        summary_text = "Results: Logit (unpenalized MLE)\n" + summary_text
+
+    return {
+        "summary_text": summary_text,
+        "coef_table_headers": ["feature","coef","std err","z","P>|z|","ci_low","ci_high"],
+        "coef_table_rows": coef_df.to_dict(orient="records"),
+    }