Data leakage — the column that gives the game away

from pathlib import Path
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

REL = 'public/datasets/decision-trees/hr-employee-attrition.csv'
here = Path.cwd()
LOCAL = next((p / REL for p in [here, *here.parents] if (p / REL).exists()), None)
CSV = LOCAL or 'https://raw.githubusercontent.com/Dr-Shashank-S-Sharma/expedify-ai-courses/main/datasets/decision-trees/hr-employee-attrition.csv'

df = pd.read_csv(CSV)
y = (df.Attrition == "Yes").astype(int)
WITHOUT_ID = pd.get_dummies(df.drop(columns=["Attrition", "EmployeeID"]), drop_first=True)
WITH_ID = pd.get_dummies(df.drop(columns=["Attrition"]), drop_first=True)

tr, te = train_test_split(range(len(df)), test_size=0.35, random_state=7, stratify=y)
ytr, yte = y.iloc[tr], y.iloc[te]

def fit(X, depth=4):
    return DecisionTreeClassifier(max_depth=depth, random_state=7).fit(X.iloc[tr], ytr)

def report(X, depth=4):
    t = fit(X, depth)
    pred = t.predict(X.iloc[te])
    return {
        "train": round(float(t.score(X.iloc[tr], ytr)), 3),
        "test": round(float(t.score(X.iloc[te], yte)), 3),
        "caught": int(((pred == 1) & (yte == 1)).sum()),
        "importances": dict(sorted(zip(X.columns, t.feature_importances_), key=lambda r: -r[1])),
    }

print(f"{len(tr)} people to learn from, {len(te)} held back, {int(yte.sum())} of whom left")

for d in (3, 4, 6, 10):
    imp = report(WITH_ID, d)["importances"].get("EmployeeID", 0.0)
    print(f"depth {d:>2}   share of the tree's decisions resting on EmployeeID: {imp:.1%}")

record("id_importance_3", f"{report(WITH_ID, 3)['importances'].get('EmployeeID', 0):.1%}")
record("id_importance_10", f"{report(WITH_ID, 10)['importances'].get('EmployeeID', 0):.1%}")

def gini(v):
    if len(v) == 0:
        return 0.0
    p = v.mean()
    return 1 - p * p - (1 - p) ** 2

def best_gain(col, target):
    """The best a single question about this column could possibly do."""
    order = np.argsort(col, kind="stable")
    c, t = np.asarray(col)[order], np.asarray(target)[order]
    n, g0, best = len(t), gini(np.asarray(target)), 0.0
    for i in range(1, n):
        if c[i] == c[i - 1]:
            continue
        best = max(best, g0 - (i * gini(t[:i]) + (n - i) * gini(t[i:])) / n)
    return best

rng = np.random.default_rng(11)
noise = []
for k in (2, 4, 10, 50, 200, 975):
    gains = [best_gain(rng.integers(0, k, len(tr)), ytr.values) for _ in range(30)]
    noise.append({"distinct": k, "mean": round(float(np.mean(gains)), 4)})
    print(f"a meaningless column with {k:>4} distinct values   best gain {np.mean(gains):.4f}")

record("noise_by_cardinality", noise)
record("noise_2", noise[0]["mean"])
record("noise_975", noise[-1]["mean"])

real = {
    "OverTime": best_gain((df.OverTime == "Yes").astype(int).iloc[tr].values, ytr.values),
    "StockOptionLevel": best_gain(df.StockOptionLevel.iloc[tr].values, ytr.values),
    "EmployeeID": best_gain(df.EmployeeID.iloc[tr].values, ytr.values),
}
for k, v in real.items():
    print(f"\nthe real {k:<18} best gain {v:.4f}")
record("gain_overtime", round(real["OverTime"], 4))
record("gain_stock", round(real["StockOptionLevel"], 4))
record("gain_id", round(real["EmployeeID"], 4))

def plot(ax):
    ax.plot([n["distinct"] for n in noise], [n["mean"] for n in noise],
            marker="o", ms=5, color="#9aa0aa", label="a column that is pure noise")
    ax.set_xscale("log")
    for label, v, c in [("the real OverTime column", real["OverTime"], "#e2574c"),
                        ("the real StockOptionLevel column", real["StockOptionLevel"], "#3b6fd4"),
                        ("the real EmployeeID column", real["EmployeeID"], "#c2871a")]:
        ax.axhline(v, ls="--", lw=1.1, color=c)
        ax.text(975, v + 0.0004, label, ha="right", fontsize=8.5, color=c)
    ax.set_xlabel("Distinct values the column has")
    ax.set_ylabel("Best gain a single question can get")
    ax.set_ylim(0, real["OverTime"] * 1.35)
    ax.legend(frameon=False, fontsize=9, loc="center left")

save_fig("lucky-columns", plot, figsize=(7, 4.0))

record("noise_50", [n["mean"] for n in noise if n["distinct"] == 50][0])
print(f"real StockOptionLevel      {real['StockOptionLevel']:.4f}")
print(f"noise with 50 values       {[n['mean'] for n in noise if n['distinct'] == 50][0]:.4f}")
print(f"real EmployeeID            {real['EmployeeID']:.4f}")
print(f"noise with 4 values        {[n['mean'] for n in noise if n['distinct'] == 4][0]:.4f}")

# CONSTRUCTED, not in the file. True for 90% of leavers and 4% of stayers,
# which is roughly how an exit-interview flag behaves in a real HR system.
rng2 = np.random.default_rng(5)
leak = np.where(y == 1, rng2.random(len(y)) < 0.90, rng2.random(len(y)) < 0.04).astype(int)

WITH_LEAK = WITH_ID.copy()
WITH_LEAK["ExitInterviewBooked"] = leak

clean, leaky = report(WITH_ID), report(WITH_LEAK)
print(f"{'':34} {'on what it saw':>15} {'on strangers':>14} {'leavers found':>15}")
for label, r in (("without the invented column", clean), ("with it", leaky)):
    print(f"{label:<34} {r['train']:>15.3f} {r['test']:>14.3f} {r['caught']:>10} of {int(yte.sum())}")

record("clean_train", f"{clean['train']:.1%}")
record("clean_test", f"{clean['test']:.1%}")
record("clean_caught", clean["caught"])
record("clean_gap", f"{clean['train'] - clean['test']:.1%}")
record("leaky_train", f"{leaky['train']:.1%}")
record("leaky_test", f"{leaky['test']:.1%}")
record("leaky_caught", leaky["caught"])
record("leaky_gap", f"{abs(leaky['train'] - leaky['test']):.1%}")
record("leak_importance", f"{leaky['importances']['ExitInterviewBooked']:.1%}")
record("n_test_leavers", int(yte.sum()))

def plot2(ax):
    labels = ["honest tree\n(no invented column)", "leaky tree\n(with it)"]
    trains = [clean["train"], leaky["train"]]
    tests = [clean["test"], leaky["test"]]
    xs = range(2)
    ax.bar([x - 0.19 for x in xs], trains, width=0.36, color="#9aa0aa", label="on what it saw")
    ax.bar([x + 0.19 for x in xs], tests, width=0.36, color="#e2574c", label="on strangers")
    for x, (a, b) in enumerate(zip(trains, tests)):
        ax.text(x - 0.19, a + 0.006, f"{a:.3f}", ha="center", fontsize=9)
        ax.text(x + 0.19, b + 0.006, f"{b:.3f}", ha="center", fontsize=9)
        ax.text(x, 0.60, f"gap {abs(a - b):.3f}", ha="center", fontsize=9.5, color="#77777f")
    ax.set_xticks(list(xs))
    ax.set_xticklabels(labels, fontsize=9)
    ax.set_ylim(0.55, 1.02)
    ax.set_ylabel("Share called right")
    ax.legend(frameon=False, fontsize=9, loc="upper left")
    ax.grid(axis="x", visible=False)

save_fig("the-leak-has-no-tell", plot2, figsize=(7, 3.8))

top = [(c, v) for c, v in leaky["importances"].items() if v > 0][:5]
for c, v in top:
    print(f"{c:<24} {v:>6.1%}  {'#' * round(v * 50)}")

record("leak_top_columns", [{"column": c, "share": f"{v:.1%}"} for c, v in top])

LABEL = {"ExitInterviewBooked": "Exit interview booked?", "OverTime_Yes": "Works overtime?",
         "JobSatisfaction": "Job satisfaction", "YearsAtCompany": "Years at company",
         "MonthlyIncome": "Monthly income", "Age": "Age", "EmployeeID": "Employee number",
         "StockOptionLevel": "Stock option level", "TotalWorkingYears": "Total working years",
         "DistanceFromHome": "Distance from home", "Education": "Education level"}

def to_nodes(clf, names, node=0):
    """sklearn's arrays -> the shape the lesson draws. sklearn sends
    `feature <= threshold` LEFT, so a yes/no column has its NO side on the
    left and the branches are swapped to ask the question the human way."""
    t = clf.tree_
    stayed, gone = t.value[node][0]
    n = int(t.n_node_samples[node])
    note = f"{n} people - {round(gone / (stayed + gone) * n)} left"
    if t.children_left[node] == -1:
        return {"leaf": "leaves" if gone >= stayed else "stays", "note": note}
    col, thr = names[t.feature[node]], t.threshold[node]
    lo = to_nodes(clf, names, t.children_left[node])
    hi = to_nodes(clf, names, t.children_right[node])
    if set(np.unique(WITH_LEAK[col])) <= {0, 1, True, False}:
        return {"q": LABEL.get(col, col), "note": note, "yes": hi, "no": lo}
    whole = bool(pd.Series(WITH_LEAK[col]).dropna().mod(1).eq(0).all())
    step = np.ceil(thr) if whole else thr
    return {"q": f"{LABEL.get(col, col)} under {step:,.0f}?", "note": note, "yes": lo, "no": hi}

record("leaky_tree", to_nodes(fit(WITH_LEAK, 2), list(WITH_LEAK.columns)))