Data Drift Detection - Is My Model Still Valid?

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import ks_2samp, chi2_contingency
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score

np.random.seed(42)

# TRAINING DISTRIBUTION (reference)
N_train = 5000
df_train = pd.DataFrame({
    "age":    np.random.normal(38, 12, N_train).clip(18, 75),
    "income": np.random.exponential(55000, N_train).clip(15000, 200000),
    "credit": np.random.normal(680, 80, N_train).clip(300, 850),
})

# PRODUCTION DATA (month 1: minor drift, month 6: major drift)
N_prod = 1000
df_prod_m1 = pd.DataFrame({
    "age":    np.random.normal(40, 13, N_prod).clip(18, 75),    # slight drift
    "income": np.random.exponential(55000, N_prod).clip(15000, 200000),
    "credit": np.random.normal(680, 80, N_prod).clip(300, 850),
})
df_prod_m6 = pd.DataFrame({
    "age":    np.random.normal(30, 10, N_prod).clip(18, 75),    # major drift (younger)
    "income": np.random.exponential(40000, N_prod).clip(15000, 200000),  # lower income
    "credit": np.random.normal(620, 90, N_prod).clip(300, 850), # lower credit
})

# ━━━━━━━━━━━━━━━━━━━━━━━━━━
# METHOD 1: KS TEST (two-sample)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━
def ks_drift_report(ref: pd.DataFrame, prod: pd.DataFrame, threshold: float = 0.05) -> pd.DataFrame:
    results = []
    for col in ref.columns:
        stat, pval = ks_2samp(ref[col], prod[col])
        drifted    = pval < threshold
        results.append({"Feature": col, "KS_Stat": stat, "P_Value": pval, "Drifted": drifted})
    return pd.DataFrame(results)

print("KS Drift Test -- Month 1 (minor drift):")
print(ks_drift_report(df_train, df_prod_m1).round(4).to_string(index=False))

print("\nKS Drift Test -- Month 6 (major drift):")
print(ks_drift_report(df_train, df_prod_m6).round(4).to_string(index=False))

# ━━━━━━━━━━━━━━━━━━━━━━━━━━
# METHOD 2: PSI (Population Stability Index)
# PSI < 0.1: stable | 0.1-0.2: slight drift | > 0.2: significant drift
# ━━━━━━━━━━━━━━━━━━━━━━━━━━
def compute_psi(reference: np.ndarray, production: np.ndarray, n_bins: int = 10) -> float:
    bins = np.percentile(reference, np.linspace(0, 100, n_bins + 1))
    bins[0] -= 1e-8; bins[-1] += 1e-8  # ensure all values captured

    ref_pct  = np.histogram(reference,  bins=bins)[0] / len(reference)
    prod_pct = np.histogram(production, bins=bins)[0] / len(production)

    # Avoid log(0)
    ref_pct  = np.clip(ref_pct,  1e-6, None)
    prod_pct = np.clip(prod_pct, 1e-6, None)

    psi = np.sum((prod_pct - ref_pct) * np.log(prod_pct / ref_pct))
    return psi

print("\nPSI Scores:")
for col in df_train.columns:
    psi_m1 = compute_psi(df_train[col].values, df_prod_m1[col].values)
    psi_m6 = compute_psi(df_train[col].values, df_prod_m6[col].values)
    status_m1 = "STABLE" if psi_m1 < 0.1 else ("SLIGHT" if psi_m1 < 0.25 else "DRIFT!")
    status_m6 = "STABLE" if psi_m6 < 0.1 else ("SLIGHT" if psi_m6 < 0.25 else "DRIFT!")
    print(f"  {col:10s}: Month 1 PSI={psi_m1:.3f} ({status_m1:7s}) | Month 6 PSI={psi_m6:.3f} ({status_m6})")

# ━━━━━━━━━━━━━━━━━━━━━━━━━━
# DRIFT RESPONSE FRAMEWORK
# ━━━━━━━━━━━━━━━━━━━━━━━━━━
response_plan = {
    "Monitor":      "PSI < 0.10 or KS p > 0.10: no action. Log and watch.",
    "Alert":        "PSI 0.10-0.25 or KS p < 0.10: investigate, trigger shadow evaluation",
    "Retrain":      "PSI > 0.25 or model AUC drops > 3%: retrain on recent data + validation",
    "Emergency":    "AUC drops > 10% or data pipeline failure: roll back to previous version",
}
print("\nDrift Response Framework:")
for severity, action in response_plan.items():
    print(f"  {severity:12s}: {action}")