#!/usr/bin/env python3
"""
Publication-ready end-to-end pipeline for:
Predicting Past-Year Major Depressive Episode (IRAMDEYR) using NSDUH 2023
Key design choices in this version
----------------------------------
1. Feature selection is skipped entirely; existing Tier 1 / 2 / 3 feature sets are used.
2. SVM is removed entirely.
3. All figures are saved as SVG by default; optional PDP panels also save as SVG.
4. Combined ROC / PR / calibration / decision-curve figures are generated per tier
with all models overlaid and CV mean ± SD shading.
5. Threshold selection is done on TRAINING DATA ONLY using out-of-fold probabilities.
The test set is used once, only for final locked evaluation.
6. Baseline table includes weighted descriptive summaries and approximate weighted
association p-values based on weighted GLM / weighted LR tests (with robust SEs).
7. Calibration metrics are added:
- Brier score
- Calibration intercept
- Calibration slope
8. Bootstrap CIs are added for final operating-point metrics (fully weighted):
- Accuracy, Sensitivity, Specificity, PPV, NPV, AUC, AP, F1, F2, Brier
9. A primary summary table is generated with one locked result row per model per tier.
Example
-------
python pipeline_publication_pubready_v2.py \
--raw-data data/raw/NSDUH_2023_Tab.txt \
--run-all --run-shap
"""
from __future__ import annotations
import argparse
import json
import re
import sys
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple
import joblib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
accuracy_score,
average_precision_score,
brier_score_loss,
confusion_matrix,
fbeta_score,
precision_recall_curve,
precision_score,
recall_score,
roc_auc_score,
roc_curve,
)
from sklearn.model_selection import (
GridSearchCV,
RandomizedSearchCV,
RepeatedStratifiedKFold,
StratifiedKFold,
train_test_split,
)
from sklearn.utils import resample
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler
warnings.filterwarnings("ignore")
# Optional dependencies
try:
from xgboost import XGBClassifier
HAS_XGB = True
except Exception:
HAS_XGB = False
try:
import lightgbm as lgb
HAS_LGBM = True
except Exception:
HAS_LGBM = False
try:
import shap
HAS_SHAP = True
except Exception:
HAS_SHAP = False
try:
from scipy.stats import randint as sp_randint, uniform as sp_uniform
from scipy.stats import chi2
HAS_SCIPY = True
except Exception:
HAS_SCIPY = False
try:
import statsmodels.api as sm
HAS_STATSMODELS = True
except Exception:
HAS_STATSMODELS = False
# ---------------------------------------------------------------------
# Global constants
# ---------------------------------------------------------------------
RANDOM_STATE = 42
np.random.seed(RANDOM_STATE)
TARGET = "IRAMDEYR"
WEIGHT = "ANALWT2_C"
SEARCH_CV_FOLDS = 5
REPORT_CV_FOLDS = 5
REPORT_CV_REPEATS = 5
N_BOOT = 1000
CALIBRATION_BINS = 10
PDP_JITTER = 0.08
PDP_CONT_GRID = 25
PDP_CONT_BOOT = 200
# Existing tier lists from prior feature-selection work
TIER_1 = [
"CATAG3",
"IRSEX",
"IRMARIT_ABS",
"IRPINC3_ABS",
"SVYRDUDANY",
"WRKDRGHLP_ABS",
]
TIER_2 = TIER_1 + [
"IRSUICTHNK",
"KSSLR6MAX_ABS",
"RCVYMHPRB",
"IRIMPGOUT_ABS",
]
TIER_3 = TIER_2 + [
"ILLEMFLAG",
"WHODASDASC_ABS",
"MJCMOTHYR",
"MJSMKYR",
"ILLEMMON",
"COCLALCUSE_ABS",
"DIFOBTCRK",
"MHNTENFCV",
"MHNTINSCV",
]
TIERS = {
"Tier_1_Basic": TIER_1,
"Tier_2_Clinical": TIER_2,
"Tier_3_Personalized": TIER_3,
}
# ---------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------
@dataclass
class Paths:
project_root: Path
raw_data: Path
data_dir: Path
cleaned_dir: Path
prepared_dir: Path
results_dir: Path
figures_dir: Path
models_dir: Path
tiers_dir: Path
@property
def cleaned_pkl(self) -> Path:
return self.cleaned_dir / "NSDUH_2023_clean.pkl"
@property
def prepared_pkl(self) -> Path:
return self.prepared_dir / "NSDUH_2023_prepared.pkl"
def make_paths(project_root: Path, raw_data: Path) -> Paths:
data_dir = project_root / "data"
cleaned_dir = data_dir / "cleaned"
prepared_dir = data_dir / "prepared"
results_dir = project_root / "results"
figures_dir = results_dir / "figures"
models_dir = project_root / "models"
tiers_dir = project_root / "tiers"
for p in [data_dir, cleaned_dir, prepared_dir, results_dir, figures_dir, models_dir, tiers_dir]:
p.mkdir(parents=True, exist_ok=True)
return Paths(
project_root=project_root,
raw_data=raw_data,
data_dir=data_dir,
cleaned_dir=cleaned_dir,
prepared_dir=prepared_dir,
results_dir=results_dir,
figures_dir=figures_dir,
models_dir=models_dir,
tiers_dir=tiers_dir,
)
# ---------------------------------------------------------------------
# Plot config
# ---------------------------------------------------------------------
def configure_publication_plots() -> None:
plt.rcParams.update(
{
"figure.dpi": 140,
"font.size": 11,
"axes.titlesize": 13,
"axes.labelsize": 11,
"legend.fontsize": 9,
"xtick.labelsize": 10,
"ytick.labelsize": 10,
"axes.spines.top": False,
"axes.spines.right": False,
"savefig.format": "svg",
"svg.fonttype": "path",
}
)
def export_tier_csvs(paths: Paths) -> None:
for tier_name, feats in TIERS.items():
out = paths.tiers_dir / f"{tier_name.lower()}_features.csv"
pd.DataFrame({"feature": feats}).to_csv(out, index=False)
FEATURE_LABELS = {
"CATAG3": "Age group",
"IRSEX": "Sex",
"IRMARIT_ABS": "Marital status",
"IRPINC3_ABS": "Family income",
"SVYRDUDANY": "Any drug use disorder severity",
"WRKDRGHLP_ABS": "Workplace assistance access",
"IRSUICTHNK": "Serious thoughts of suicide",
"KSSLR6MAX_ABS": "Psychological distress severity",
"RCVYMHPRB": "Received mental health treatment/counseling",
"IRIMPGOUT_ABS": "Difficulty going out alone",
"ILLEMFLAG": "Any illegal drug use",
"WHODASDASC_ABS": "WHODAS disability score",
"MJCMOTHYR": "Marijuana use (past year)",
"MJSMKYR": "Marijuana smoking (past year)",
"ILLEMMON": "Illegal drug use other than Marijuana (past month)",
"COCLALCUSE_ABS": "Alcohol use frequency because of COVID",
"DIFOBTCRK": "Difficulty obtaining crack",
"MHNTENFCV": "Mental health treatment not covered by employer",
"MHNTINSCV": "Mental health treatment not covered by insurance",
}
LEVEL_LABELS = {
"CATAG3": {
2: "Age 18–25",
3: "Age 26–34",
4: "Age 35–49",
5: "Age 50+",
"2": "Age 18–25",
"3": "Age 26–34",
"4": "Age 35–49",
"5": "Age 50+",
},
"IRSEX": {
1: "Male",
2: "Female",
"1": "Male",
"2": "Female",
},
"SVYRDUDANY": {
4: "No disorder",
1: "Mild",
2: "Moderate",
3: "Severe",
"4": "No disorder",
"1": "Mild",
"2": "Moderate",
"3": "Severe",
},
}
# ---------------------------------------------------------------------
# Data cleaning and preparation
# ---------------------------------------------------------------------
def clean_nsduh(raw_path: Path, cleaned_path: Path) -> pd.DataFrame:
df = pd.read_csv(raw_path, sep="\t", low_memory=False)
df = df[df["AGE3"] >= 4].copy()
nan_codes = [
94, 97, 98, 85,
994, 997, 998, 985,
9994, 9997, 9998, 9985,
99, 999, 9999,
89, 989, 9989,
]
df = df.replace(nan_codes, np.nan)
df = df.dropna(subset=[TARGET]).copy()
drop_patterns = [
r"^CASEID", r"^QUESTID", r"^QUESTID2", r"^PANEL", r"^VERSION",
r"^INTV", r"^FILE", r"^YEAR",
r"^WT", r"WGT", r"WEIGHT", r"ANALWT(?!2_C$)",
r"^VESTR", r"^VEREP", r"^ESTRAT",
r"_A$", r"_E$", r"_ORIG$", r"_R$", r"_RC$",
]
cols_to_drop = [c for c in df.columns if any(re.search(pattern, c) for pattern in drop_patterns)]
cols_to_drop = [c for c in cols_to_drop if c != WEIGHT]
cols_to_drop += [c for c in ["WTANSWER", "WTPOUND2", "FILEDATE"] if c in df.columns]
cols_to_drop = sorted(set(cols_to_drop))
df = df.drop(columns=cols_to_drop, errors="ignore")
na_pct = df.isna().mean() * 100
cols_over_99 = na_pct[na_pct > 99].index.tolist()
df = df.drop(columns=cols_over_99, errors="ignore")
leakage_vars = [
"IIAMDEYR", "AMDEYR", "AMDELT", "IRAMDELT", "IIAMDELT",
"AMDEIMP", "IRAMDEIMP", "IIAMDEIMP", "AMDETXRX",
"ADSMMDEA",
"D_MDEA1", "D_MDEA2", "D_MDEA3", "D_MDEA4", "D_MDEA5",
"D_MDEA6", "D_MDEA7", "D_MDEA8", "D_MDEA9",
"AD_MDEA1", "AD_MDEA2", "AD_MDEA3", "AD_MDEA4",
"AD_MDEA5", "AD_MDEA6", "AD_MDEA7", "AD_MDEA8",
"AD_MDEA11", "AD_MDEA21", "AD_MDEA31", "AD_MDEA41",
"AD_MDEA51", "AD_MDEA61", "AD_MDEA71", "AD_MDEA81", "AD_MDEA91",
"ADPB2WK",
"ADDPREV", "ADDSCEV", "ADLOSEV", "ADLSI2WK", "ADDPR2WK",
"ADWRHRS", "ADWRDST", "ADWRCHR", "ADWRIMP", "ADDPPROB",
"ARXMDEYR", "ATXMDEYR", "AHLTMDE",
]
df = df.drop(columns=[c for c in leakage_vars if c in df.columns], errors="ignore")
df[TARGET] = pd.to_numeric(df[TARGET], errors="coerce")
df = df[df[TARGET].isin([0, 1])].copy()
df[TARGET] = df[TARGET].astype(int)
df.to_pickle(cleaned_path)
return df
def _weighted_mean(x: np.ndarray, w: Optional[np.ndarray] = None) -> float:
x = np.asarray(x, dtype=float)
if w is None:
return float(np.mean(x))
w = np.asarray(w, dtype=float)
mask = np.isfinite(x) & np.isfinite(w)
if not np.any(mask):
return np.nan
denom = np.sum(w[mask])
if denom <= 0:
return np.nan
return float(np.sum(x[mask] * w[mask]) / denom)
def _series_is_continuous_numeric(s: pd.Series, min_unique: int = 8) -> bool:
# Explicitly exclude strings and categoricals from being treated as continuous
if isinstance(s.dtype, pd.CategoricalDtype) or pd.api.types.is_object_dtype(s):
return False
s_num = pd.to_numeric(s, errors="coerce")
return s_num.notna().sum() >= 20 and s_num.nunique(dropna=True) >= min_unique
def _series_is_discrete(s: pd.Series, max_unique: int = 20) -> bool:
s_nonmissing = s.dropna()
if s_nonmissing.empty:
return False
if isinstance(s_nonmissing.dtype, pd.CategoricalDtype) or pd.api.types.is_object_dtype(s_nonmissing):
return s_nonmissing.nunique(dropna=True) <= max_unique
s_num = pd.to_numeric(s_nonmissing, errors="coerce")
return s_num.notna().sum() == len(s_nonmissing) and s_num.nunique(dropna=True) <= max_unique
def _ordered_categories_for_plot(s: pd.Series) -> List:
if isinstance(s.dtype, pd.CategoricalDtype):
return [c for c in s.cat.categories if pd.notna(c)]
vals = s.dropna().unique().tolist()
try:
return sorted(vals)
except Exception:
return vals
def _discrete_tick_label(v) -> str:
if isinstance(v, (int, np.integer)):
return str(int(v))
if isinstance(v, (float, np.floating)) and float(v).is_integer():
return str(int(v))
return str(v)
def plot_discrete_effect(
out_svg: Path,
model,
X_ref: pd.DataFrame,
feature: str,
weights: Optional[np.ndarray] = None,
label_func=None,
) -> bool:
s = X_ref[feature]
categories = _ordered_categories_for_plot(s)
if len(categories) < 2:
return False
X_base = X_ref.copy()
p_base = model.predict_proba(X_base)[:, 1]
effects_all = []
x_all = []
freqs = []
mean_effects = [] # Added to store the actual PDP average
for i, cat in enumerate(categories):
X_tmp = X_ref.copy()
if isinstance(X_tmp[feature].dtype, pd.CategoricalDtype):
X_tmp[feature] = pd.Categorical([cat] * len(X_tmp), categories=X_tmp[feature].cat.categories, ordered=X_tmp[feature].cat.ordered)
else:
X_tmp[feature] = cat
p_cf = model.predict_proba(X_tmp)[:, 1]
eff = p_cf - p_base
effects_all.append(eff)
x_all.append(np.full(len(eff), i, dtype=float))
if weights is None:
freqs.append(float((s == cat).mean()))
mean_effects.append(float(np.mean(eff)))
else:
mask = (s == cat).to_numpy()
freqs.append(_weighted_mean(mask.astype(float), weights))
mean_effects.append(_weighted_mean(eff, weights))
effects_all = np.concatenate(effects_all)
x_all = np.concatenate(x_all)
freqs = np.asarray(freqs, dtype=float)
jitter = np.random.normal(0, PDP_JITTER, size=len(x_all))
x_scatter = x_all + jitter
fig, ax1 = plt.subplots(figsize=(7.6, 5.0))
ax1.bar(np.arange(len(categories)), freqs, edgecolor="black", alpha=0.25, color="lightgray")
ax1.set_ylabel("Category frequency" if weights is None else "Weighted category frequency")
ax1.set_xticks(np.arange(len(categories)))
ax1.set_xticklabels([_discrete_tick_label(c) for c in categories], rotation=45, ha="right")
ax1.set_xlabel(label_func(feature) if label_func else feature)
ax2 = ax1.twinx()
# Plot ICE (individual points) faintly
ax2.scatter(x_scatter, effects_all, s=10, alpha=0.15, color="steelblue", zorder=1)
# Plot actual PDP (mean effect line) prominently
ax2.plot(np.arange(len(categories)), mean_effects, color="darkred", marker="o", markersize=8, linewidth=2.5, zorder=2, label="Mean Marginal Effect")
ax2.axhline(0, linestyle="--", linewidth=1, color="black")
ax2.set_ylabel("Δ Predicted Risk")
ax2.legend(loc="upper left", frameon=False)
plt.title(f"Partial Dependence (Discrete): {label_func(feature) if label_func else feature}")
plt.tight_layout()
plt.savefig(out_svg, bbox_inches="tight")
plt.close()
return True
def plot_continuous_effect(
out_svg: Path,
model,
X_ref: pd.DataFrame,
weights: Optional[np.ndarray],
feature: str,
grid_points: int = PDP_CONT_GRID,
boot_iters: int = PDP_CONT_BOOT,
label_func=None,
) -> bool:
x = pd.to_numeric(X_ref[feature], errors="coerce").values
mask = np.isfinite(x)
if mask.sum() < 20:
return False
X_ref2 = X_ref.loc[mask].copy()
x = x[mask]
w_ref = None if weights is None else np.asarray(weights, dtype=float)[mask]
orig_dtype = X_ref2[feature].dtype
qs = np.linspace(0.05, 0.95, grid_points)
grid = np.unique(np.quantile(x, qs))
if len(grid) < 2:
return False
n_samples = len(X_ref2)
pred_matrix = np.zeros((n_samples, len(grid)))
for j, v in enumerate(grid):
X_tmp = X_ref2.copy()
X_tmp[feature] = float(v)
try:
X_tmp[feature] = X_tmp[feature].astype(orig_dtype)
except:
pass
pred_matrix[:, j] = model.predict_proba(X_tmp)[:, 1]
mean = np.array([_weighted_mean(pred_matrix[:, j], w_ref) for j in range(len(grid))])
# Create figure explicitly
fig, ax = plt.subplots(figsize=(7.6, 5.0))
ax.plot(grid, mean, linewidth=2.5, color="darkred", label="Mean Marginal Effect")
if boot_iters and boot_iters > 1:
ys = []
for i in range(boot_iters):
idx = resample(np.arange(n_samples), replace=True, random_state=42 + i)
wb = None if w_ref is None else w_ref[idx]
y_boot = np.array([_weighted_mean(pred_matrix[idx, j], wb) for j in range(len(grid))])
ys.append(y_boot)
ys = np.asarray(ys, dtype=float)
lo = np.percentile(ys, 2.5, axis=0)
hi = np.percentile(ys, 97.5, axis=0)
ax.fill_between(grid, lo, hi, color="darkred", alpha=0.2, label="95% CI")
ax.set_xlabel(label_func(feature) if label_func else feature)
ax.set_ylabel("Predicted Risk" + (" (Weighted)" if weights is not None else ""))
ax.set_title(f"Partial Dependence: {label_func(feature) if label_func else feature}")
ax.legend(frameon=False, loc="best")
ax.grid(True, alpha=0.25)
# Critical Fix: Use fig.savefig instead of plt.savefig
plt.tight_layout()
fig.savefig(out_svg, bbox_inches="tight", facecolor='white')
plt.close(fig) # Close specific figure
return True
def run_partial_dependence_plots(
df: pd.DataFrame,
tier_name: str,
paths: Paths,
features_to_plot: List[str],
model_name: str = "RF",
) -> None:
print(f"\n[PDP] Running PDP-like plots for {model_name} / {tier_name}")
prefix_map = {
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}
if model_name not in prefix_map:
raise ValueError(f"Unsupported model for PDP-like plots: {model_name}")
if tier_name not in TIERS:
raise ValueError(f"Unknown tier: {tier_name}")
model_path = paths.models_dir / f"{prefix_map[model_name]}_{tier_name.lower()}.joblib"
if not model_path.exists():
print(f"[WARN] Model not found: {model_path}")
return
tier_features = [c for c in TIERS[tier_name] if c in df.columns]
requested_features = [f for f in features_to_plot if f in tier_features]
if not requested_features:
print(f"[WARN] No requested PDP features are present in {tier_name}.")
return
missing_cols = [c for c in TIERS[tier_name] if c not in df.columns]
if missing_cols:
raise KeyError(f"Prepared data is missing required tier columns: {missing_cols}")
model = joblib.load(model_path)
X_ref = df[tier_features].copy()
weights = None
if WEIGHT in df.columns:
weights = pd.to_numeric(df[WEIGHT], errors="coerce").fillna(1.0).astype(float).values
out_dir = paths.figures_dir / "pdp_like" / model_name / tier_name
out_dir.mkdir(parents=True, exist_ok=True)
made_any = False
for feature in requested_features:
s = X_ref[feature]
out_svg = out_dir / f"pdp_like_{model_name.lower()}_{tier_name.lower()}_{feature}.svg"
try:
ok = False
if _series_is_continuous_numeric(s):
ok = plot_continuous_effect(
out_svg=out_svg,
model=model,
X_ref=X_ref,
weights=weights,
feature=feature,
grid_points=PDP_CONT_GRID,
boot_iters=PDP_CONT_BOOT,
label_func=pretty_feature_name,
)
elif _series_is_discrete(s):
ok = plot_discrete_effect(
out_svg=out_svg,
model=model,
X_ref=X_ref,
feature=feature,
weights=weights,
label_func=pretty_feature_name,
)
else:
# fallback: try treating anything else as a categorical/discrete variable
ok = plot_discrete_effect(
out_svg=out_svg,
model=model,
X_ref=X_ref,
feature=feature,
weights=weights,
label_func=pretty_feature_name,
)
if ok:
made_any = True
print(f"[PDP] Saved: {out_svg}")
else:
print(f"[WARN] Could not create PDP-like plot for '{feature}' in {tier_name}")
except Exception as e:
print(f"[WARN] Skipping PDP-like for feature '{feature}' in {tier_name}: {e}")
if not made_any:
print(f"[WARN] No PDP-like plots were generated for {model_name} / {tier_name}")
def prepare_features(cleaned_df: pd.DataFrame, prepared_path: Path) -> pd.DataFrame:
df = cleaned_df.copy()
features_to_remove = ["ADRX12MO", "IRMHTRXMED", "ADWRENRG", "ADWRSLNO", "ADPSHMGT"]
df = df.drop(columns=[c for c in features_to_remove if c in df.columns], errors="ignore")
raw_to_recode = {
"AGE3": "CATAG3",
"CAMHRCVR": "RCVYMHPRB",
"MJSKNYR": "MJCMOTHYR",
"DIFGETCRK": "DIFOBTCRK",
"MHTUNENFCV": "MHNTENFCV",
"MHTUNINSCV": "MHNTINSCV",
}
for raw_col, derived_col in raw_to_recode.items():
if derived_col not in df.columns and raw_col in df.columns:
df[derived_col] = df[raw_col]
if raw_col in df.columns and raw_col != derived_col:
df = df.drop(columns=[raw_col], errors="ignore")
if "IRMARIT" in df.columns:
map_ = {1: "Married", 2: "Widowed-Divorced-or-Separated", 3: "Widowed-Divorced-or-Separated", 4: "Never Married"}
order = ["Married", "Widowed-Divorced-or-Separated", "Never Married"]
df["IRMARIT_ABS"] = pd.Categorical(df["IRMARIT"].map(map_), categories=order, ordered=True)
df = df.drop(columns=["IRMARIT"])
if "IRPINC3" in df.columns:
map_ = {
1: "a_<10k",
2: "b_10k-to-40k",
3: "b_10k-to-40k",
4: "b_10k-to-40k",
5: "c_50k-to-75k",
6: "c_50k-to-75k",
7: "d_75k+",
}
order = ["a_<10k", "b_10k-to-40k", "c_50k-to-75k", "d_75k+"]
df["IRPINC3_ABS"] = pd.Categorical(df["IRPINC3"].map(map_), categories=order, ordered=True)
df = df.drop(columns=["IRPINC3"])
if "WRKDRGHLP" in df.columns:
map_ = {
1: "Having access to assistance program at work",
2: "Not having access to assistance program at work",
}
order = [
"Having access to assistance program at work",
"Not having access to assistance program at work",
"No need or do not know about assistance program at work",
]
df["WRKDRGHLP_ABS"] = df["WRKDRGHLP"].map(map_).fillna("No need or do not know about assistance program at work")
df["WRKDRGHLP_ABS"] = pd.Categorical(df["WRKDRGHLP_ABS"], categories=order, ordered=True)
df = df.drop(columns=["WRKDRGHLP"])
if "KSSLR6MAX" in df.columns:
bins = [-1, 6, 12, 18, 24]
labels = ["0-6", "07-12", "13-18", "19-24"]
df["KSSLR6MAX_ABS"] = pd.cut(df["KSSLR6MAX"], bins=bins, labels=labels)
df["KSSLR6MAX_ABS"] = pd.Categorical(df["KSSLR6MAX_ABS"], categories=labels, ordered=True)
df = df.drop(columns=["KSSLR6MAX"])
if "IRIMPGOUT" in df.columns:
map_ = {
1: "b_No difficulty",
2: "c_Mild difficulty",
3: "d_Moderate difficulty",
4: "e_Severe difficulty",
5: "a_Not-Going-Out-Alone-Or-No-Difficulty",
99: "a_Not-Going-Out-Alone-Or-No-Difficulty",
}
order = [
"a_Not-Going-Out-Alone-Or-No-Difficulty",
"b_No difficulty",
"c_Mild difficulty",
"d_Moderate difficulty",
"e_Severe difficulty",
]
df["IRIMPGOUT_ABS"] = df["IRIMPGOUT"].map(map_).fillna("a_Not-Going-Out-Alone-Or-No-Difficulty")
df["IRIMPGOUT_ABS"] = pd.Categorical(df["IRIMPGOUT_ABS"], categories=order, ordered=True)
df = df.drop(columns=["IRIMPGOUT"])
if "WHODASDASC" in df.columns:
map_ = {0: "0", 1: "1-2", 2: "1-2", 3: "3-4", 4: "3-4", 5: "5-6", 6: "5-6", 7: "7-8", 8: "7-8"}
order = ["0", "1-2", "3-4", "5-6", "7-8"]
df["WHODASDASC_ABS"] = pd.Categorical(df["WHODASDASC"].map(map_), categories=order, ordered=True)
df = df.drop(columns=["WHODASDASC"])
if "COCLALCUSE" in df.columns:
raw_map = {
1: "Drink much less or little less",
2: "Drink about the same",
3: "Drink a little more or much more",
}
label_map = {
"Unknown/No PY Alcohol Use": "a_Unknown/No PY Alcohol Use",
"Drink much less or little less": "b_Drink much less or little less",
"Drink about the same": "c_Drink about the same",
"Drink a little more or much more": "d_Drink a little more or much more",
}
order = [
"a_Unknown/No PY Alcohol Use",
"b_Drink much less or little less",
"c_Drink about the same",
"d_Drink a little more or much more",
]
df["COCLALCUSE_ABS"] = df["COCLALCUSE"].map(raw_map).fillna("Unknown/No PY Alcohol Use")
df["COCLALCUSE_ABS"] = pd.Categorical(df["COCLALCUSE_ABS"].map(label_map), categories=order, ordered=True)
df = df.drop(columns=["COCLALCUSE"])
df.to_pickle(prepared_path)
return df
def pretty_feature_name(feature_name: str) -> str:
if feature_name in FEATURE_LABELS:
return FEATURE_LABELS[feature_name]
# one-hot encoded style names: BASE_LEVEL
for base in sorted(FEATURE_LABELS.keys(), key=len, reverse=True):
prefix = f"{base}_"
if feature_name.startswith(prefix):
raw_level = feature_name[len(prefix):]
pretty_base = FEATURE_LABELS[base]
pretty_level = LEVEL_LABELS.get(base, {}).get(raw_level, raw_level)
return f"{pretty_base}: {pretty_level}"
return feature_name
def pretty_code_meaning(var: str, level) -> str:
return LEVEL_LABELS.get(var, {}).get(level, LEVEL_LABELS.get(var, {}).get(str(level), str(level)))
# ---------------------------------------------------------------------
# Baseline characteristics with approximate weighted p-values
# ---------------------------------------------------------------------
def weighted_mean_sd(series: pd.Series, weights: pd.Series) -> Tuple[float, float]:
x = series.astype(float).to_numpy()
w = weights.astype(float).to_numpy()
mean = np.average(x, weights=w)
var = np.average((x - mean) ** 2, weights=w)
return float(mean), float(np.sqrt(var))
def weighted_percent(df: pd.DataFrame, var: str, target_value: Optional[int] = None) -> pd.Series:
d = df.copy()
if target_value is not None:
d = d[d[TARGET] == target_value].copy()
grp = d.groupby(var, dropna=False)[WEIGHT].sum()
return (grp / grp.sum() * 100).round(1)
def format_p_value(p: float) -> str:
if pd.isna(p):
return ""
if p < 0.001:
return "<0.001"
return f"{p:.3f}"
def approximate_weighted_p_continuous(df: pd.DataFrame, var: str) -> float:
if not (HAS_STATSMODELS and HAS_SCIPY):
return np.nan
d = df[[TARGET, var, WEIGHT]].dropna().copy()
if d.empty or d[var].nunique() < 2 or d[TARGET].nunique() < 2:
return np.nan
try:
X = sm.add_constant(d[[var]].astype(float))
y = d[TARGET].astype(float)
model = sm.GLM(y, X, family=sm.families.Binomial(), var_weights=d[WEIGHT].astype(float))
# Adding HC1 robust standard errors
fit = model.fit(cov_type='HC1')
return float(fit.pvalues.get(var, np.nan))
except Exception:
return np.nan
def approximate_weighted_p_categorical(df: pd.DataFrame, var: str) -> float:
if not (HAS_STATSMODELS and HAS_SCIPY):
return np.nan
d = df[[TARGET, var, WEIGHT]].dropna().copy()
if d.empty or d[var].nunique() < 2 or d[TARGET].nunique() < 2:
return np.nan
try:
X_full = pd.get_dummies(d[var].astype(str), drop_first=True)
if X_full.shape[1] == 0:
return np.nan
X_full = sm.add_constant(X_full.astype(float), has_constant="add")
X_null = sm.add_constant(pd.DataFrame(index=d.index), has_constant="add")
y = d[TARGET].astype(float)
w = d[WEIGHT].astype(float)
# Adding HC1 robust standard errors
fit_full = sm.GLM(y, X_full, family=sm.families.Binomial(), var_weights=w).fit(cov_type='HC1')
fit_null = sm.GLM(y, X_null, family=sm.families.Binomial(), var_weights=w).fit(cov_type='HC1')
lr = 2 * (fit_full.llf - fit_null.llf)
df_diff = fit_full.df_model - fit_null.df_model
p = 1 - chi2.cdf(lr, df_diff)
return float(p)
except Exception:
return np.nan
def build_baseline_characteristics_table(paths: Paths, tier_name: str = "Tier_3_Personalized") -> pd.DataFrame:
"""
Build a full baseline/descriptive table for all features in the requested tier.
Uses the prepared dataset so the recoded publication features are included.
"""
df = pd.read_pickle(paths.prepared_pkl).copy()
df = ensure_binary_target(df)
if WEIGHT not in df.columns:
df[WEIGHT] = 1.0
feature_list = [f for f in TIERS[tier_name] if f in df.columns]
rows = []
# Sample size row
no_mde_raw = int((df[TARGET] == 0).sum())
mde_raw = int((df[TARGET] == 1).sum())
total_raw = no_mde_raw + mde_raw
no_mde_pct = round(no_mde_raw / total_raw * 100, 1) if total_raw > 0 else np.nan
mde_pct = round(mde_raw / total_raw * 100, 1) if total_raw > 0 else np.nan
rows.append([
"Total participants",
"Total respondents in analytic sample",
f"{no_mde_raw:,} ({no_mde_pct}%)",
f"{mde_raw:,} ({mde_pct}%)",
"",
"",
])
for var in feature_list:
series = df[var]
is_categorical = (
str(series.dtype) == "category"
or series.dtype == "object"
or series.nunique(dropna=True) <= 10
)
if is_categorical:
no_mde = weighted_percent(df, var, 0)
mde = weighted_percent(df, var, 1)
p_value = approximate_weighted_p_categorical(df, var)
# Keep the original level types for lookup
if str(series.dtype) == "category":
levels = [x for x in series.cat.categories if pd.notna(x)]
else:
levels = [x for x in series.dropna().unique()]
try:
levels = sorted(levels)
except Exception:
levels = list(levels)
for i, level in enumerate(levels):
no_val = no_mde.get(level, np.nan)
mde_val = mde.get(level, np.nan)
level_label = str(level)
rows.append([
pretty_feature_name(f"{var}_{level_label}"),
pretty_code_meaning(var, level),
f"{no_val:.1f} %" if pd.notna(no_val) else "",
f"{mde_val:.1f} %" if pd.notna(mde_val) else "",
"Weighted LR test (HC1 robust SE)" if i == 0 else "",
format_p_value(p_value) if i == 0 else "",
])
else:
d0 = df.loc[df[TARGET] == 0, [var, WEIGHT]].dropna()
d1 = df.loc[df[TARGET] == 1, [var, WEIGHT]].dropna()
mean0, sd0 = weighted_mean_sd(d0[var], d0[WEIGHT]) if len(d0) else (np.nan, np.nan)
mean1, sd1 = weighted_mean_sd(d1[var], d1[WEIGHT]) if len(d1) else (np.nan, np.nan)
p_value = approximate_weighted_p_continuous(df, var)
rows.append([
pretty_feature_name(var),
"",
f"{mean0:.2f} ± {sd0:.2f}",
f"{mean1:.2f} ± {sd1:.2f}",
"Weighted GLM (HC1 robust SE)",
format_p_value(p_value),
])
table = pd.DataFrame(
rows,
columns=[
"Feature (code)",
"Code Meaning",
"No MDE (Weighted %) or Mean ± SD",
"MDE (Weighted %) or Mean ± SD",
"Statistical Test",
"P-value",
],
)
out_name = f"table1_{tier_name}_Baseline_Characteristics_with_pvalues.csv"
table.to_csv(paths.results_dir / out_name, index=False)
note = (
"P-values are approximate weighted model-based association tests using frequency weights. "
"They utilize HC1 robust standard errors but are not true complex-survey design-based tests."
)
pd.DataFrame({"note": [note]}).to_csv(
paths.results_dir / f"table1_{tier_name}_Baseline_note.csv",
index=False,
)
return table
# ---------------------------------------------------------------------
# Feature selection intentionally skipped
# ---------------------------------------------------------------------
def run_feature_selection(*args, **kwargs):
print("[INFO] Feature selection is skipped in this version of the pipeline.")
print("[INFO] Using pre-defined Tier_1 / Tier_2 / Tier_3 feature lists only.")
return {}
# ---------------------------------------------------------------------
# Modeling helpers
# ---------------------------------------------------------------------
def ensure_binary_target(df: pd.DataFrame) -> pd.DataFrame:
out = df.copy()
out[TARGET] = pd.to_numeric(out[TARGET], errors="coerce")
out = out[out[TARGET].isin([0, 1])].copy()
out[TARGET] = out[TARGET].astype(int)
return out
def get_prepared_data(paths: Paths) -> pd.DataFrame:
if not paths.prepared_pkl.exists():
raise FileNotFoundError(f"Prepared data not found: {paths.prepared_pkl}")
df = pd.read_pickle(paths.prepared_pkl)
return ensure_binary_target(df)
def split_final_model_data(df: pd.DataFrame):
y = df[TARGET].astype(int)
if WEIGHT in df.columns:
w = pd.to_numeric(df[WEIGHT], errors="coerce").fillna(1.0).astype(float)
else:
w = pd.Series(np.ones(len(df)), index=df.index, dtype=float)
X = df.drop(columns=[TARGET] + ([WEIGHT] if WEIGHT in df.columns else []))
return train_test_split(X, y, w, test_size=0.2, stratify=y, random_state=RANDOM_STATE)
def _select_existing_features(df_columns: Iterable[str], features: List[str], tier_name: str) -> List[str]:
cols = list(df_columns)
existing = [f for f in features if f in cols]
missing = [f for f in features if f not in cols]
if missing:
print(f"[WARN] {tier_name}: {len(missing)} features missing from prepared data: {missing}")
if not existing:
raise ValueError(f"{tier_name}: none of the requested tier features were found.")
return existing
def fit_with_optional_sample_weight(model, X, y, sample_weight):
try:
model.fit(X, y, clf__sample_weight=sample_weight)
return model
except Exception:
pass
try:
model.fit(X, y, sample_weight=sample_weight)
return model
except Exception:
pass
model.fit(X, y)
return model
def fit_model_for_family(model, model_name: str, X, y, sample_weight):
return fit_with_optional_sample_weight(model, X, y, sample_weight)
def predict_proba_with_model(model, X, model_name: str) -> np.ndarray:
return model.predict_proba(X)[:, 1]
def confusion_elements(y_true: np.ndarray, preds: np.ndarray, sample_weight: Optional[np.ndarray] = None) -> Tuple[float, float, float, float]:
tn, fp, fn, tp = confusion_matrix(y_true, preds, sample_weight=sample_weight).ravel()
return float(tn), float(fp), float(fn), float(tp)
def evaluate_model(y_true: pd.Series, preds: np.ndarray, probs: np.ndarray, sample_weight: Optional[np.ndarray] = None) -> Dict[str, float]:
tn, fp, fn, tp = confusion_elements(np.asarray(y_true), preds, sample_weight)
specificity = tn / (tn + fp) if (tn + fp) > 0 else np.nan
sensitivity = recall_score(y_true, preds, sample_weight=sample_weight, zero_division=0)
ppv = precision_score(y_true, preds, sample_weight=sample_weight, zero_division=0)
npv = tn / (tn + fn) if (tn + fn) > 0 else np.nan
return {
"Accuracy": accuracy_score(y_true, preds, sample_weight=sample_weight),
"Sensitivity": sensitivity,
"Specificity": specificity,
"PPV": ppv,
"NPV": npv,
"AUC": roc_auc_score(y_true, probs, sample_weight=sample_weight),
"AveragePrecision": average_precision_score(y_true, probs, sample_weight=sample_weight),
"F1": fbeta_score(y_true, preds, beta=1, sample_weight=sample_weight, zero_division=0),
"F2": fbeta_score(y_true, preds, beta=2, sample_weight=sample_weight, zero_division=0),
}
def evaluate_at_threshold(y_true: pd.Series, probs: np.ndarray, threshold: float, sample_weight: Optional[np.ndarray] = None) -> Dict[str, float]:
preds = (probs >= threshold).astype(int)
base = evaluate_model(y_true, preds, probs, sample_weight)
base["threshold"] = threshold
return base
def threshold_sweep(y_true: pd.Series, probs: np.ndarray, sample_weight: Optional[np.ndarray] = None) -> pd.DataFrame:
thresholds = np.arange(0.10, 0.91, 0.05)
rows = []
for t in thresholds:
rows.append(evaluate_at_threshold(y_true, probs, float(t), sample_weight))
return pd.DataFrame(rows)
def choose_screening_threshold(metrics_df: pd.DataFrame) -> pd.Series:
screening_candidates = metrics_df[
(metrics_df["Sensitivity"] >= 0.90) &
(metrics_df["NPV"] >= 0.95)
]
if screening_candidates.empty:
return metrics_df.sort_values(
["Sensitivity", "NPV", "Specificity", "F2"],
ascending=[False, False, False, False],
).iloc[0]
return screening_candidates.sort_values(
["Specificity", "F2"],
ascending=[False, False],
).iloc[0]
def choose_locked_threshold_from_training_cv(
best_model,
model_name: str,
X_train: pd.DataFrame,
y_train: pd.Series,
w_train: pd.Series,
) -> Tuple[float, pd.DataFrame]:
skf = StratifiedKFold(n_splits=REPORT_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
oof_probs = np.zeros(len(X_train), dtype=float)
for tr_idx, val_idx in skf.split(X_train, y_train):
Xtr = X_train.iloc[tr_idx].copy()
Xva = X_train.iloc[val_idx].copy()
ytr = y_train.iloc[tr_idx]
wtr = w_train.iloc[tr_idx]
model_clone = clone(best_model)
Xtr_fit, Xva_fit = Xtr, Xva
fit_model_for_family(model_clone, model_name, Xtr_fit, ytr, wtr)
oof_probs[val_idx] = predict_proba_with_model(model_clone, Xva_fit, model_name)
threshold_df = threshold_sweep(
y_train.reset_index(drop=True),
oof_probs,
sample_weight=w_train.reset_index(drop=True)
)
best_threshold_row = choose_screening_threshold(threshold_df)
return float(best_threshold_row["threshold"]), threshold_df
def bootstrap_metric_cis(
y_true: pd.Series,
probs: np.ndarray,
threshold: float,
sample_weight: Optional[np.ndarray] = None,
n_boot: int = N_BOOT,
) -> Dict[str, Tuple[float, float]]:
metrics_store: Dict[str, List[float]] = {
"Accuracy": [],
"Sensitivity": [],
"Specificity": [],
"PPV": [],
"NPV": [],
"AUC": [],
"AveragePrecision": [],
"F1": [],
"F2": [],
"Brier": [],
}
y = np.asarray(y_true).astype(int)
probs = np.asarray(probs)
w = np.asarray(sample_weight) if sample_weight is not None else np.ones_like(y, dtype=float)
for _ in range(n_boot):
idx = np.random.randint(0, len(y), len(y))
y_bs = y[idx]
p_bs = probs[idx]
w_bs = w[idx]
if len(np.unique(y_bs)) < 2:
continue
preds_bs = (p_bs >= threshold).astype(int)
base = evaluate_model(y_bs, preds_bs, p_bs, sample_weight=w_bs)
base["Brier"] = brier_score_loss(y_bs, p_bs, sample_weight=w_bs)
for k in metrics_store:
metrics_store[k].append(float(base[k]))
out: Dict[str, Tuple[float, float]] = {}
for k, vals in metrics_store.items():
if len(vals) == 0:
out[k] = (np.nan, np.nan)
else:
out[k] = tuple(np.percentile(vals, [2.5, 97.5]).tolist())
return out
def calibration_metrics(y_true: pd.Series, probs: np.ndarray) -> Dict[str, float]:
out = {"Brier": brier_score_loss(y_true, probs), "Calib_Intercept": np.nan, "Calib_Slope": np.nan}
if not HAS_STATSMODELS:
return out
try:
y = pd.Series(y_true).astype(float).reset_index(drop=True)
p = np.clip(pd.Series(probs).astype(float).reset_index(drop=True), 1e-6, 1 - 1e-6)
lp = np.log(p / (1 - p))
# intercept: y ~ 1 with offset=lp
X_int = np.ones((len(y), 1))
fit_int = sm.GLM(y, X_int, family=sm.families.Binomial(), offset=lp).fit()
out["Calib_Intercept"] = float(fit_int.params[0])
# slope: y ~ lp
X_slope = sm.add_constant(lp)
fit_slope = sm.GLM(y, X_slope, family=sm.families.Binomial()).fit()
out["Calib_Slope"] = float(fit_slope.params[1])
except Exception:
pass
return out
# ---------------------------------------------------------------------
# Model builders
# ---------------------------------------------------------------------
def build_logreg_search(X: pd.DataFrame) -> GridSearchCV:
categorical_features = X.select_dtypes(include=["object", "category"]).columns.tolist()
numeric_features = [c for c in X.columns if c not in categorical_features]
preprocessor = ColumnTransformer(
transformers=[
("num", Pipeline([
("imputer", SimpleImputer(strategy="median")),
("scaler", StandardScaler()),
]), numeric_features),
("cat", Pipeline([
("imputer", SimpleImputer(strategy="most_frequent")),
("onehot", OneHotEncoder(handle_unknown="ignore")),
]), categorical_features),
],
remainder="drop",
)
pipe = Pipeline([
("preprocessor", preprocessor),
("clf", LogisticRegression(max_iter=2000, solver="liblinear", random_state=RANDOM_STATE)),
])
param_grid = {
"clf__C": [0.1, 0.5, 1.0, 2.0],
"clf__class_weight": [None, "balanced"],
}
cv = StratifiedKFold(n_splits=SEARCH_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
return GridSearchCV(pipe, param_grid=param_grid, scoring="recall", cv=cv, n_jobs=-1, refit=True)
def build_rf_search(X: pd.DataFrame) -> RandomizedSearchCV:
categorical_features = X.select_dtypes(include=["object", "category"]).columns.tolist()
numeric_features = [c for c in X.columns if c not in categorical_features]
preprocessor = ColumnTransformer(
transformers=[
("num", SimpleImputer(strategy="median"), numeric_features),
("cat", Pipeline([
("imputer", SimpleImputer(strategy="most_frequent")),
("ordinal", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)),
]), categorical_features),
],
remainder="drop",
)
pipe = Pipeline([
("preprocessor", preprocessor),
("clf", RandomForestClassifier(random_state=RANDOM_STATE, n_jobs=-1)),
])
param_dist = {
"clf__n_estimators": [300, 500],
"clf__max_depth": [None, 6, 10],
"clf__min_samples_split": [2, 5, 10],
"clf__min_samples_leaf": [1, 2, 4],
"clf__max_features": ["sqrt", 0.5],
"clf__class_weight": [None, "balanced"],
}
cv = StratifiedKFold(n_splits=SEARCH_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
return RandomizedSearchCV(
pipe,
param_distributions=param_dist,
n_iter=15,
scoring="recall",
cv=cv,
n_jobs=-1,
random_state=RANDOM_STATE,
refit=True,
verbose=1,
)
def build_xgb_search(X: pd.DataFrame, y_train: pd.Series) -> RandomizedSearchCV:
if not HAS_XGB:
raise ImportError("xgboost is not installed.")
categorical_features = X.select_dtypes(include=["object", "category"]).columns.tolist()
numeric_features = [c for c in X.columns if c not in categorical_features]
preprocessor = ColumnTransformer(
transformers=[
("num", SimpleImputer(strategy="median"), numeric_features),
("cat", Pipeline([
("imputer", SimpleImputer(strategy="most_frequent")),
("ordinal", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)),
]), categorical_features),
],
remainder="drop",
)
pos = int(y_train.sum())
neg = int(len(y_train) - pos)
scale_pos_weight = neg / max(pos, 1)
pipe = Pipeline([
("preprocessor", preprocessor),
("clf", XGBClassifier(
objective="binary:logistic",
eval_metric="logloss",
tree_method="hist",
random_state=RANDOM_STATE,
n_jobs=-1,
verbosity=0,
)),
])
param_dist = {
"clf__n_estimators": [200, 400, 600],
"clf__max_depth": [3, 4, 5],
"clf__learning_rate": [0.03, 0.05, 0.1],
"clf__subsample": [0.8, 1.0],
"clf__colsample_bytree": [0.8, 1.0],
"clf__min_child_weight": [1, 3],
"clf__scale_pos_weight": [1.0, scale_pos_weight],
"clf__reg_alpha": [0.0, 0.01, 0.1],
"clf__reg_lambda": [1.0, 2.0, 5.0],
}
cv = StratifiedKFold(n_splits=SEARCH_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
return RandomizedSearchCV(
pipe,
param_distributions=param_dist,
n_iter=15,
scoring="recall",
cv=cv,
n_jobs=-1,
random_state=RANDOM_STATE,
refit=True,
verbose=1,
)
def build_lgbm_search(X: pd.DataFrame, y_train: pd.Series) -> RandomizedSearchCV:
if not HAS_LGBM:
raise ImportError("lightgbm is not installed.")
categorical_features = X.select_dtypes(include=["object", "category"]).columns.tolist()
numeric_features = [c for c in X.columns if c not in categorical_features]
preprocessor = ColumnTransformer(
transformers=[
("num", SimpleImputer(strategy="median"), numeric_features),
("cat", Pipeline([
("imputer", SimpleImputer(strategy="constant", fill_value="Missing")),
("ordinal", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)),
]), categorical_features),
],
remainder="drop",
)
preprocessor.set_output(transform="pandas")
pos = int(y_train.sum())
neg = int(len(y_train) - pos)
scale_pos_weight = neg / max(pos, 1)
pipe = Pipeline([
("preprocessor", preprocessor),
("clf", lgb.LGBMClassifier(
objective="binary",
boosting_type="gbdt",
class_weight="balanced",
random_state=RANDOM_STATE,
n_jobs=-1,
verbosity=-1,
max_bin=255,
)),
])
param_dist = {
"clf__num_leaves": sp_randint(15, 64),
"clf__max_depth": [-1, 3, 4, 5, 6, 8],
"clf__learning_rate": sp_uniform(0.01, 0.04),
"clf__n_estimators": sp_randint(200, 800),
"clf__min_child_samples": sp_randint(5, 20),
"clf__subsample": sp_uniform(0.7, 0.3),
"clf__colsample_bytree": sp_uniform(0.7, 0.3),
"clf__reg_alpha": [0.0, 0.01, 0.1, 1.0],
"clf__reg_lambda": [0.0, 0.01, 0.1, 1.0],
"clf__scale_pos_weight": [1.0, scale_pos_weight],
}
cv = StratifiedKFold(n_splits=SEARCH_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
return RandomizedSearchCV(
estimator=pipe,
param_distributions=param_dist,
n_iter=15,
scoring="recall",
cv=cv,
n_jobs=-1,
random_state=RANDOM_STATE,
refit=True,
)
# ---------------------------------------------------------------------
# Curves for publication figures
# ---------------------------------------------------------------------
def net_benefit(y_true: np.ndarray, probs: np.ndarray, thresholds: np.ndarray) -> np.ndarray:
y_true = np.asarray(y_true).astype(int)
probs = np.asarray(probs)
n = len(y_true)
out = []
for pt in thresholds:
pred_pos = probs >= pt
tp = np.sum((pred_pos == 1) & (y_true == 1))
fp = np.sum((pred_pos == 1) & (y_true == 0))
nb = (tp / n) - (fp / n) * (pt / (1 - pt))
out.append(nb)
return np.asarray(out)
def collect_cv_curve_payload(best_model, model_name: str, X_train: pd.DataFrame, y_train: pd.Series, w_train: pd.Series) -> Dict[str, np.ndarray]:
rkf = RepeatedStratifiedKFold(
n_splits=REPORT_CV_FOLDS,
n_repeats=REPORT_CV_REPEATS,
random_state=RANDOM_STATE,
)
fpr_grid = np.linspace(0.0, 1.0, 200)
recall_grid = np.linspace(0.0, 1.0, 200)
cal_bins = np.linspace(0.0, 1.0, CALIBRATION_BINS + 1)
cal_centers = (cal_bins[:-1] + cal_bins[1:]) / 2
dca_thresholds = np.linspace(0.01, 0.99, 99)
roc_curves, pr_curves, cal_curves, dca_curves = [], [], [], []
aucs, aps, prevalences = [], [], []
for tr_idx, val_idx in rkf.split(X_train, y_train):
Xtr = X_train.iloc[tr_idx].copy()
Xva = X_train.iloc[val_idx].copy()
ytr = y_train.iloc[tr_idx]
yva = y_train.iloc[val_idx]
wtr = w_train.iloc[tr_idx]
model_clone = clone(best_model)
Xtr_fit, Xva_fit = Xtr, Xva
fit_model_for_family(model_clone, model_name, Xtr_fit, ytr, wtr)
probs = predict_proba_with_model(model_clone, Xva_fit, model_name)
fpr, tpr, _ = roc_curve(yva, probs)
tpr_interp = np.interp(fpr_grid, fpr, tpr)
tpr_interp[0] = 0.0
roc_curves.append(tpr_interp)
aucs.append(roc_auc_score(yva, probs))
precision, recall, _ = precision_recall_curve(yva, probs)
recall_sorted = recall[::-1]
precision_sorted = precision[::-1]
pr_interp = np.interp(recall_grid, recall_sorted, precision_sorted)
pr_curves.append(pr_interp)
aps.append(average_precision_score(yva, probs))
probs_clipped = np.clip(probs, 1e-6, 1 - 1e-6)
bin_ids = np.digitize(probs_clipped, cal_bins, right=True) - 1
obs_rates = []
for b in range(CALIBRATION_BINS):
mask = bin_ids == b
if np.any(mask):
obs_rates.append(np.mean(yva[mask]))
else:
obs_rates.append(np.nan)
cal_curves.append(obs_rates)
dca_curves.append(net_benefit(yva.to_numpy(), probs, dca_thresholds))
prevalences.append(float(np.mean(yva)))
roc_curves = np.asarray(roc_curves)
pr_curves = np.asarray(pr_curves)
cal_curves = np.asarray(cal_curves, dtype=float)
dca_curves = np.asarray(dca_curves)
return {
"fpr_grid": fpr_grid,
"roc_mean": np.nanmean(roc_curves, axis=0),
"roc_std": np.nanstd(roc_curves, axis=0),
"auc_mean": float(np.nanmean(aucs)),
"auc_std": float(np.nanstd(aucs)),
"recall_grid": recall_grid,
"pr_mean": np.nanmean(pr_curves, axis=0),
"pr_std": np.nanstd(pr_curves, axis=0),
"ap_mean": float(np.nanmean(aps)),
"ap_std": float(np.nanstd(aps)),
"cal_x": cal_centers,
"cal_mean": np.nanmean(cal_curves, axis=0),
"cal_std": np.nanstd(cal_curves, axis=0),
"dca_thresholds": dca_thresholds,
"dca_mean": np.nanmean(dca_curves, axis=0),
"dca_std": np.nanstd(dca_curves, axis=0),
"prevalence_mean": float(np.nanmean(prevalences)),
}
def save_combined_tier_figures(curve_payloads: Dict[str, Dict[str, Dict[str, np.ndarray]]], figures_dir: Path) -> None:
for tier_name, model_payloads in curve_payloads.items():
safe_tier = tier_name.lower()
# ROC
plt.figure(figsize=(7.0, 6.0))
for model_name, payload in model_payloads.items():
x = payload["fpr_grid"]
y = payload["roc_mean"]
s = payload["roc_std"]
plt.plot(x, y, lw=2, label=f"{model_name} (AUC {payload['auc_mean']:.3f}±{payload['auc_std']:.3f})")
plt.fill_between(x, np.maximum(0, y - s), np.minimum(1, y + s), alpha=0.15)
plt.plot([0, 1], [0, 1], linestyle="--", lw=1, color="black")
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.title(f"ROC Comparison — {tier_name}")
plt.legend(frameon=False, loc="lower right")
plt.tight_layout()
plt.savefig(figures_dir / f"{safe_tier}_combined_roc.svg", format="svg", bbox_inches="tight")
plt.close()
# PR
plt.figure(figsize=(7.0, 6.0))
for model_name, payload in model_payloads.items():
x = payload["recall_grid"]
y = payload["pr_mean"]
s = payload["pr_std"]
plt.plot(x, y, lw=2, label=f"{model_name} (AP {payload['ap_mean']:.3f}±{payload['ap_std']:.3f})")
plt.fill_between(x, np.maximum(0, y - s), np.minimum(1, y + s), alpha=0.15)
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.title(f"Precision–Recall Comparison — {tier_name}")
plt.legend(frameon=False, loc="best")
plt.tight_layout()
plt.savefig(figures_dir / f"{safe_tier}_combined_pr.svg", format="svg", bbox_inches="tight")
plt.close()
# Calibration
plt.figure(figsize=(7.0, 6.0))
for model_name, payload in model_payloads.items():
x = payload["cal_x"]
y = payload["cal_mean"]
s = payload["cal_std"]
plt.plot(x, y, marker="o", lw=2, label=model_name)
plt.fill_between(x, np.maximum(0, y - s), np.minimum(1, y + s), alpha=0.15)
plt.plot([0, 1], [0, 1], linestyle="--", lw=1, color="black")
plt.xlabel("Predicted Probability")
plt.ylabel("Observed Event Rate")
plt.title(f"Calibration Comparison — {tier_name}")
plt.legend(frameon=False, loc="best")
plt.tight_layout()
plt.savefig(figures_dir / f"{safe_tier}_combined_calibration.svg", format="svg", bbox_inches="tight")
plt.close()
# Decision curve
first_payload = next(iter(model_payloads.values()))
thresholds = first_payload["dca_thresholds"]
prevalence = first_payload["prevalence_mean"]
plt.figure(figsize=(7.0, 6.0))
for model_name, payload in model_payloads.items():
x = payload["dca_thresholds"]
y = payload["dca_mean"]
s = payload["dca_std"]
plt.plot(x, y, lw=2, label=model_name)
plt.fill_between(x, y - s, y + s, alpha=0.15)
treat_all = prevalence - (1 - prevalence) * (thresholds / (1 - thresholds))
treat_none = np.zeros_like(thresholds)
plt.plot(thresholds, treat_all, linestyle="--", lw=1, label="Treat All")
plt.plot(thresholds, treat_none, linestyle=":", lw=1, label="Treat None")
plt.xlabel("Threshold Probability")
plt.ylabel("Net Benefit")
plt.title(f"Decision Curve Analysis — {tier_name}")
plt.legend(frameon=False, loc="best")
plt.tight_layout()
plt.savefig(figures_dir / f"{safe_tier}_combined_decision_curve.svg", format="svg", bbox_inches="tight")
plt.close()
# ---------------------------------------------------------------------
# Final model training / evaluation
# ---------------------------------------------------------------------
def run_cv_table(best_model, model_name: str, X_train: pd.DataFrame, y_train: pd.Series, w_train: pd.Series, metrics_order: List[str]) -> pd.DataFrame:
skf = StratifiedKFold(n_splits=REPORT_CV_FOLDS, shuffle=True, random_state=RANDOM_STATE)
fold_rows = []
for fold_i, (tr_idx, val_idx) in enumerate(skf.split(X_train, y_train), start=1):
Xtr = X_train.iloc[tr_idx].copy()
Xval = X_train.iloc[val_idx].copy()
ytr = y_train.iloc[tr_idx]
yval = y_train.iloc[val_idx]
wtr = w_train.iloc[tr_idx]
model_clone = clone(best_model)
Xtr_fit, Xval_fit = Xtr, Xval
fit_model_for_family(model_clone, model_name, Xtr_fit, ytr, wtr)
preds = model_clone.predict(Xval_fit)
probs = predict_proba_with_model(model_clone, Xval_fit, model_name)
fold_metrics = evaluate_model(yval, preds, probs)
fold_metrics["Fold"] = f"Fold {fold_i}"
fold_rows.append(fold_metrics)
fold_df = pd.DataFrame(fold_rows).set_index("Fold")[metrics_order].T
fold_df["Mean"] = fold_df.mean(axis=1)
fold_df["Std Dev"] = fold_df.std(axis=1)
fold_df = fold_df.reset_index().rename(columns={"index": "Metric"})
return fold_df
def run_model_family(model_name: str, df: pd.DataFrame, paths: Paths) -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, Dict[str, np.ndarray]]]:
X_train, X_test, y_train, y_test, w_train, w_test = split_final_model_data(df)
tier_results = []
tier_cv_tables = []
tier_curve_payloads = {}
metrics_order = ["Accuracy", "Sensitivity", "Specificity", "PPV", "NPV", "AUC", "AveragePrecision", "F1", "F2"]
for tier_name, features in TIERS.items():
print(f"\n===== {model_name}: {tier_name} =====")
tier_features = _select_existing_features(X_train.columns, features, tier_name)
Xtr = X_train[tier_features].copy()
Xte = X_test[tier_features].copy()
if model_name == "LogReg":
search = build_logreg_search(Xtr)
elif model_name == "RF":
search = build_rf_search(Xtr)
elif model_name == "XGB":
search = build_xgb_search(Xtr, y_train)
elif model_name == "LGBM":
search = build_lgbm_search(Xtr, y_train)
else:
raise ValueError(f"Unknown model: {model_name}")
Xtr_fit, Xte_fit = Xtr, Xte
fit_model_for_family(search, model_name, Xtr_fit, y_train, w_train)
best_model = search.best_estimator_ if hasattr(search, "best_estimator_") else search
model_file_prefix = {
"LogReg": "logreg",
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}[model_name]
joblib.dump(best_model, paths.models_dir / f"{model_file_prefix}_{tier_name.lower()}.joblib")
locked_threshold, train_threshold_df = choose_locked_threshold_from_training_cv(
best_model=best_model,
model_name=model_name,
X_train=Xtr,
y_train=y_train,
w_train=w_train,
)
train_threshold_df.insert(0, "tier", tier_name)
train_threshold_df.insert(0, "model", model_name)
train_threshold_df.to_csv(
paths.results_dir / f"{model_name.lower()}_{tier_name.lower()}_training_threshold_sweep.csv",
index=False,
)
# Final locked evaluation on test set (Weighted metrics)
probs_test = predict_proba_with_model(best_model, Xte_fit, model_name)
preds_default = (probs_test >= 0.50).astype(int)
preds_locked = (probs_test >= locked_threshold).astype(int)
metrics_default = evaluate_model(y_test, preds_default, probs_test, sample_weight=w_test)
metrics_locked = evaluate_model(y_test, preds_locked, probs_test, sample_weight=w_test)
cal = calibration_metrics(y_test, probs_test)
metric_cis = bootstrap_metric_cis(y_test, probs_test, threshold=locked_threshold, sample_weight=w_test, n_boot=N_BOOT)
row_default = {
"selection": "Default_0.50",
"tier": tier_name,
"model": model_name,
"tuned_for": "recall",
"n_features": len(tier_features),
"threshold": 0.50,
**metrics_default,
"Brier": cal["Brier"],
"Calib_Intercept": cal["Calib_Intercept"],
"Calib_Slope": cal["Calib_Slope"],
"Best_Params": json.dumps(getattr(search, "best_params_", {}), default=str),
}
row_locked = {
"selection": "Locked_TrainCV_Threshold",
"tier": tier_name,
"model": model_name,
"tuned_for": "recall",
"n_features": len(tier_features),
"threshold": locked_threshold,
**metrics_locked,
"Brier": cal["Brier"],
"Calib_Intercept": cal["Calib_Intercept"],
"Calib_Slope": cal["Calib_Slope"],
"Accuracy_CI_Low": metric_cis["Accuracy"][0],
"Accuracy_CI_High": metric_cis["Accuracy"][1],
"Sensitivity_CI_Low": metric_cis["Sensitivity"][0],
"Sensitivity_CI_High": metric_cis["Sensitivity"][1],
"Specificity_CI_Low": metric_cis["Specificity"][0],
"Specificity_CI_High": metric_cis["Specificity"][1],
"PPV_CI_Low": metric_cis["PPV"][0],
"PPV_CI_High": metric_cis["PPV"][1],
"NPV_CI_Low": metric_cis["NPV"][0],
"NPV_CI_High": metric_cis["NPV"][1],
"AUC_CI_Low": metric_cis["AUC"][0],
"AUC_CI_High": metric_cis["AUC"][1],
"AveragePrecision_CI_Low": metric_cis["AveragePrecision"][0],
"AveragePrecision_CI_High": metric_cis["AveragePrecision"][1],
"F1_CI_Low": metric_cis["F1"][0],
"F1_CI_High": metric_cis["F1"][1],
"F2_CI_Low": metric_cis["F2"][0],
"F2_CI_High": metric_cis["F2"][1],
"Brier_CI_Low": metric_cis["Brier"][0],
"Brier_CI_High": metric_cis["Brier"][1],
"Best_Params": json.dumps(getattr(search, "best_params_", {}), default=str),
}
tier_results.extend([row_default, row_locked])
cv_table = run_cv_table(best_model, model_name, Xtr, y_train, w_train, metrics_order)
cv_table.insert(0, "Tier", tier_name)
cv_table.insert(0, "Model", model_name)
tier_cv_tables.append(cv_table)
tier_curve_payloads[tier_name] = collect_cv_curve_payload(best_model, model_name, Xtr, y_train, w_train)
results_df = pd.DataFrame(tier_results)
cv_results_df = pd.concat(tier_cv_tables, ignore_index=True)
model_to_prefix = {
"LogReg": "logreg",
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}
prefix = model_to_prefix[model_name]
results_df.to_csv(paths.results_dir / f"{prefix}_two_rows_per_tier.csv", index=False)
cv_results_df.to_csv(paths.results_dir / f"{prefix}_{SEARCH_CV_FOLDS}fold_threshold0.50_by_tier.csv", index=False)
primary_subset = results_df[results_df["selection"] == "Locked_TrainCV_Threshold"].copy()
primary_subset.to_csv(paths.results_dir / f"{prefix}_primary_locked_results.csv", index=False)
return results_df, cv_results_df, tier_curve_payloads
# ---------------------------------------------------------------------
# Comparison tables
# ---------------------------------------------------------------------
def create_model_comparison_tables(paths: Paths, included_models: List[str]) -> None:
prefix_map = {
"LogReg": "logreg",
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}
metrics_keep = ["Accuracy", "Sensitivity", "Specificity", "PPV", "NPV", "AUC", "AveragePrecision", "F1", "F2"]
all_two_row_dfs = []
all_cv_dfs = []
for model in included_models:
prefix = prefix_map[model]
two_row_path = paths.results_dir / f"{prefix}_two_rows_per_tier.csv"
cv_path = paths.results_dir / f"{prefix}_{SEARCH_CV_FOLDS}fold_threshold0.50_by_tier.csv"
if two_row_path.exists():
all_two_row_dfs.append(pd.read_csv(two_row_path))
if cv_path.exists():
all_cv_dfs.append(pd.read_csv(cv_path))
if all_two_row_dfs:
all_models = pd.concat(all_two_row_dfs, ignore_index=True)
all_models.to_csv(paths.results_dir / "all_models_by_tier_thresholds.csv", index=False)
means = (
all_models[all_models["selection"] == "Default_0.50"]
.sort_values(["model", "tier"])
.loc[:, ["model", "tier", "n_features", "Accuracy", "Sensitivity", "Specificity", "PPV", "NPV", "AUC", "AveragePrecision", "F1", "F2", "Brier", "Calib_Intercept", "Calib_Slope"]]
.reset_index(drop=True)
)
means.to_csv(paths.results_dir / "model_comparison_table_means.csv", index=False)
primary_results = (
all_models[all_models["selection"] == "Locked_TrainCV_Threshold"]
.sort_values(["tier", "model"])
.loc[:, [
"tier", "model", "n_features", "threshold",
"Accuracy", "Accuracy_CI_Low", "Accuracy_CI_High",
"Sensitivity", "Sensitivity_CI_Low", "Sensitivity_CI_High",
"Specificity", "Specificity_CI_Low", "Specificity_CI_High",
"PPV", "PPV_CI_Low", "PPV_CI_High",
"NPV", "NPV_CI_Low", "NPV_CI_High",
"AUC", "AUC_CI_Low", "AUC_CI_High",
"AveragePrecision", "AveragePrecision_CI_Low", "AveragePrecision_CI_High",
"F1", "F1_CI_Low", "F1_CI_High",
"F2", "F2_CI_Low", "F2_CI_High",
"Brier", "Brier_CI_Low", "Brier_CI_High",
"Calib_Intercept", "Calib_Slope",
]]
.reset_index(drop=True)
)
primary_results.to_csv(paths.results_dir / "primary_results_locked_thresholds.csv", index=False)
if all_cv_dfs:
cv_all = pd.concat(all_cv_dfs, ignore_index=True)
cv_all = cv_all[cv_all["Metric"].isin(metrics_keep)].copy()
cv_all.to_csv(paths.results_dir / "all_models_comparison.csv", index=False)
# ---------------------------------------------------------------------
# SHAP
# ---------------------------------------------------------------------
def run_shap_for_tree_models(df: pd.DataFrame, paths: Paths, models_to_use: Optional[List[str]] = None) -> None:
if not HAS_SHAP:
print("[WARN] shap is not installed; skipping SHAP analysis.")
return
if models_to_use is None:
models_to_use = ["LGBM", "XGB"]
X_train, X_test, y_train, y_test, w_train, w_test = split_final_model_data(df)
all_importances = []
for model_name in models_to_use:
for tier_name, features in TIERS.items():
tier_features = _select_existing_features(X_test.columns, features, tier_name)
model_file_prefix = {
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}[model_name]
model_path = paths.models_dir / f"{model_file_prefix}_{tier_name.lower()}.joblib"
if not model_path.exists():
print(f"[WARN] Missing model for SHAP: {model_path}")
continue
model = joblib.load(model_path)
if hasattr(model, "best_estimator_"):
model = model.best_estimator_
Xte = X_test[tier_features].copy()
clf = model
transformed = Xte
feature_names = tier_features
if hasattr(model, "named_steps"):
pre = model.named_steps.get("preprocessor")
clf = model.named_steps.get("clf")
if pre is not None:
transformed = pre.transform(Xte)
try:
feature_names = pre.get_feature_names_out().tolist()
except Exception:
feature_names = [f"feature_{i}" for i in range(transformed.shape[1])]
try:
explainer = shap.TreeExplainer(clf)
shap_values = explainer.shap_values(transformed)
if isinstance(shap_values, list):
shap_values = shap_values[-1]
except Exception as e:
print(f"[WARN] SHAP failed for {model_name} / {tier_name}: {e}")
continue
plt.figure(figsize=(9, 6))
shap.summary_plot(shap_values, transformed, feature_names=feature_names, plot_type="bar", show=False)
plt.title(f"SHAP Feature Importance — {model_name} — {tier_name}")
plt.tight_layout()
plt.savefig(paths.figures_dir / f"{model_name.lower()}_shap_importance_{tier_name}.svg", format="svg", bbox_inches="tight")
plt.close()
plt.figure(figsize=(9, 6))
shap.summary_plot(shap_values, transformed, feature_names=feature_names, show=False)
plt.title(f"SHAP Summary — {model_name} — {tier_name}")
plt.tight_layout()
plt.savefig(paths.figures_dir / f"{model_name.lower()}_shap_beeswarm_{tier_name}.svg", format="svg", bbox_inches="tight")
plt.close()
mean_abs = np.abs(np.asarray(shap_values)).mean(axis=0)
shap_importance = pd.DataFrame({"feature": feature_names, "mean_abs_shap": mean_abs}).sort_values("mean_abs_shap", ascending=False)
shap_importance["model"] = model_name
shap_importance["tier"] = tier_name
shap_importance["rank"] = np.arange(1, len(shap_importance) + 1)
all_importances.append(shap_importance)
if all_importances:
out = pd.concat(all_importances, ignore_index=True)
out = out[["model", "tier", "rank", "feature", "mean_abs_shap"]]
out.to_csv(paths.results_dir / "tree_model_shap_importance.csv", index=False)
def run_shap_dependence_panels_for_tier3(
df: pd.DataFrame,
paths: Paths,
models_to_use: Optional[List[str]] = None,
max_display: int = 8,
sample_n: int = 5000,
) -> None:
if not HAS_SHAP:
print("[WARN] shap is not installed; skipping SHAP dependence panels.")
return
if models_to_use is None:
models_to_use = ["RF", "XGB", "LGBM"]
X_train, X_test, y_train, y_test, w_train, w_test = split_final_model_data(df)
tier_name = "Tier_3_Personalized"
tier_features = _select_existing_features(X_test.columns, TIERS[tier_name], tier_name)
for model_name in models_to_use:
model_file_prefix = {
"RF": "rf",
"XGB": "xgb",
"LGBM": "lightgbm",
}[model_name]
model_path = paths.models_dir / f"{model_file_prefix}_{tier_name.lower()}.joblib"
if not model_path.exists():
print(f"[WARN] Missing model for SHAP dependence: {model_path}")
continue
print(f"[SHAP] Dependence plots for {model_name} / {tier_name}")
model = joblib.load(model_path)
if hasattr(model, "best_estimator_"):
model = model.best_estimator_
Xte = X_test[tier_features].copy()
if len(Xte) > sample_n:
rng = np.random.default_rng(RANDOM_STATE)
idx = rng.choice(len(Xte), size=sample_n, replace=False)
Xte = Xte.iloc[idx].copy()
clf = model
transformed = Xte
feature_names = tier_features
if hasattr(model, "named_steps"):
pre = model.named_steps.get("preprocessor")
clf = model.named_steps.get("clf")
if pre is not None:
transformed = pre.transform(Xte)
try:
feature_names = pre.get_feature_names_out().tolist()
except Exception:
feature_names = [f"feature_{i}" for i in range(transformed.shape[1])]
if hasattr(transformed, "toarray"):
transformed_for_shap = transformed.toarray()
else:
transformed_for_shap = np.asarray(transformed)
try:
explainer = shap.TreeExplainer(clf)
shap_values = explainer.shap_values(transformed_for_shap)
if isinstance(shap_values, list):
shap_values = shap_values[-1]
shap_values = np.asarray(shap_values)
if shap_values.ndim == 3:
shap_values = shap_values[:, :, -1]
elif shap_values.ndim != 2:
raise ValueError(f"Unexpected SHAP shape: {shap_values.shape}")
except Exception as e:
print(f"[WARN] SHAP dependence failed for {model_name}: {e}")
continue
mean_abs = np.abs(shap_values).mean(axis=0)
top_idx = np.argsort(mean_abs)[::-1][:max_display].astype(int).tolist()
top_features = [feature_names[i] for i in top_idx]
plot_X = pd.DataFrame(transformed_for_shap, columns=feature_names)
fig, axes = plt.subplots(2, 4, figsize=(16, 10))
axes = axes.flatten()
for ax_i, feat_name in enumerate(top_features):
plt.sca(axes[ax_i])
try:
shap.dependence_plot(
feat_name,
shap_values,
plot_X,
interaction_index="auto",
ax=axes[ax_i],
show=False,
)
axes[ax_i].set_title(pretty_feature_name(feat_name), fontsize=10)
except Exception as e:
axes[ax_i].text(
0.5, 0.5,
f"Plot failed:\n{feat_name}",
ha="center", va="center", fontsize=10
)
axes[ax_i].set_axis_off()
for j in range(len(top_features), len(axes)):
axes[j].set_axis_off()
fig.suptitle(f"SHAP Dependence Plots — Top {len(top_features)} Features — {model_name} — {tier_name}", fontsize=14)
plt.tight_layout(rect=[0, 0, 1, 0.97])
plt.savefig(
paths.figures_dir / f"{model_name.lower()}_tier3_shap_dependence_top{max_display}.svg",
format="svg",
bbox_inches="tight",
)
plt.close()
top_table = pd.DataFrame({
"feature": feature_names,
"feature_label": [pretty_feature_name(f) for f in feature_names],
"mean_abs_shap": mean_abs
}).sort_values("mean_abs_shap", ascending=False).reset_index(drop=True)
top_table.to_csv(
paths.results_dir / f"{model_name.lower()}_tier3_shap_top_features.csv",
index=False,
)
# ---------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------
def run_pipeline(
paths: Paths,
run_cleaning: bool,
run_preparation: bool,
run_baseline: bool,
run_models: bool,
run_shap_flag: bool,
run_pdp: bool = False,
pdp_model_name: str = "RF",
pdp_tier_name: str = "Tier_3_Personalized",
pdp_features: Optional[List[str]] = None,
) -> None:
configure_publication_plots()
export_tier_csvs(paths)
if run_cleaning:
print("\n[1/4] Cleaning raw data...")
clean_nsduh(paths.raw_data, paths.cleaned_pkl)
print(f"Saved cleaned data -> {paths.cleaned_pkl}")
if run_preparation:
print("\n[2/4] Preparing features...")
if not paths.cleaned_pkl.exists():
raise FileNotFoundError("Run cleaning first or provide cleaned pickle.")
cleaned_df = pd.read_pickle(paths.cleaned_pkl)
prepare_features(cleaned_df, paths.prepared_pkl)
print(f"Saved prepared data -> {paths.prepared_pkl}")
if run_baseline:
print("\n[3/4] Building baseline characteristics table with approximate weighted p-values...")
build_baseline_characteristics_table(paths, tier_name="Tier_3_Personalized")
print(f"Saved -> {paths.results_dir / 'table1_Tier_3_Personalized_Baseline_Characteristics_with_pvalues.csv'}")
if run_models:
print("\n[4/4] Training final models by tier using existing tier definitions...")
prepared_df = get_prepared_data(paths)
model_order = ["LogReg", "RF"]
if HAS_XGB:
model_order.append("XGB")
else:
print("[WARN] xgboost not installed; skipping XGB.")
if HAS_LGBM and HAS_SCIPY:
model_order.append("LGBM")
else:
print("[WARN] lightgbm and/or scipy not installed; skipping LGBM.")
curve_payloads_all = {tier_name: {} for tier_name in TIERS.keys()}
for model_name in model_order:
_, _, model_curve_payloads = run_model_family(model_name, prepared_df, paths)
for tier_name, payload in model_curve_payloads.items():
curve_payloads_all[tier_name][model_name] = payload
create_model_comparison_tables(paths, model_order)
save_combined_tier_figures(curve_payloads_all, paths.figures_dir)
print(f"Saved model comparison tables -> {paths.results_dir}")
print(f"Saved combined publication figures -> {paths.figures_dir}")
if run_shap_flag:
print("\n[SHAP] Running SHAP for tree models...")
run_shap_for_tree_models(prepared_df, paths)
print("\n[SHAP] Running Tier 3 SHAP dependence panels...")
run_shap_dependence_panels_for_tier3(
prepared_df,
paths,
models_to_use=["RF", "XGB", "LGBM"],
max_display=8,
sample_n=5000,
)
if run_pdp:
pdp_features_to_use = pdp_features or TIERS[pdp_tier_name]
pdp_models_to_run: List[str] = []
for candidate in ["RF", "XGB", "LGBM"]:
model_path = paths.models_dir / ({"RF": "rf", "XGB": "xgb", "LGBM": "lightgbm"}[candidate] + f"_{pdp_tier_name.lower()}.joblib")
if model_path.exists():
pdp_models_to_run.append(candidate)
if not pdp_models_to_run and pdp_model_name:
pdp_models_to_run = [pdp_model_name]
for pdp_model in pdp_models_to_run:
run_partial_dependence_plots(
df=prepared_df,
tier_name=pdp_tier_name,
paths=paths,
features_to_plot=pdp_features_to_use,
model_name=pdp_model,
)
print("\nPipeline complete.")
# ---------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Consolidated NSDUH MDE publication pipeline (feature selection skipped, no SVM, SVG figures only)."
)
parser.add_argument(
"--project-root",
type=str,
default=".",
help="Project root directory where data/, results/, models/, tiers/ live.",
)
parser.add_argument(
"--raw-data",
type=str,
default="data/raw/NSDUH_2023_Tab.txt",
help="Path to raw NSDUH tab-delimited file.",
)
parser.add_argument("--run-all", action="store_true", help="Run cleaning, preparation, baseline table, final modeling, and optional SHAP.")
parser.add_argument("--run-cleaning", action="store_true")
parser.add_argument("--run-preparation", action="store_true")
parser.add_argument("--run-baseline", action="store_true")
parser.add_argument("--run-models", action="store_true")
parser.add_argument("--run-shap", action="store_true")
parser.add_argument("--run-pdp", action="store_true", help="Generate partial dependence plots after model fitting.")
parser.add_argument("--pdp-model", type=str, default="RF", choices=["RF", "XGB", "LGBM"], help="Fallback model to use for PDP generation if auto-detection finds no saved model files.")
parser.add_argument("--pdp-tier", type=str, default="Tier_3_Personalized", choices=list(TIERS.keys()), help="Tier to use for PDP generation.")
parser.add_argument("--pdp-features", nargs="*", default=None, help="Optional list of feature names to plot for PDP. Defaults to a curated Tier 3 set.")
return parser.parse_args()
def main() -> None:
args = parse_args()
project_root = Path(args.project_root).resolve()
raw_data = Path(args.raw_data).resolve() if Path(args.raw_data).is_absolute() else (project_root / args.raw_data).resolve()
paths = make_paths(project_root, raw_data)
if args.run_all:
run_cleaning = True
run_preparation = True
run_baseline = True
run_models = True
run_shap_flag = args.run_shap
run_pdp = args.run_pdp
else:
run_cleaning = args.run_cleaning
run_preparation = args.run_preparation
run_baseline = args.run_baseline
run_models = args.run_models
run_shap_flag = args.run_shap
run_pdp = args.run_pdp
if not any([run_cleaning, run_preparation, run_baseline, run_models]):
print("No stage selected. Use --run-all or one or more stage flags.")
sys.exit(1)
if run_cleaning or run_baseline:
if not raw_data.exists():
raise FileNotFoundError(f"Raw data file not found: {raw_data}")
run_pipeline(
paths=paths,
run_cleaning=run_cleaning,
run_preparation=run_preparation,
run_baseline=run_baseline,
run_models=run_models,
run_shap_flag=run_shap_flag,
run_pdp=run_pdp,
pdp_model_name=args.pdp_model,
pdp_tier_name=args.pdp_tier,
pdp_features=args.pdp_features,
)
if __name__ == "__main__":
main()