Customer churn prediction in Python, step by step

import pandas as pd
import requests
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import precision_score, recall_score, roc_auc_score
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.pipeline import Pipeline, make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

df = pd.DataFrame(requests.get("https://api.gradient.test/v1/customers").json()["data"])
numeric = ["tenure_months", "orders_last_year", "avg_order_value", "days_since_last_order",
           "discount_share", "email_open_rate", "support_tickets", "trail_club"]
categorical = ["region", "channel", "top_category"]
X_dev, X_test, y_dev, y_test = train_test_split(df[numeric + categorical], df["churned"], test_size=0.25,
                                                random_state=42, stratify=df["churned"])
print(f"development set {len(X_dev)}, test set {len(X_test)} ({int(y_test.sum())} churners), churn rate {y_dev.mean():.0%}")

rule = (X_dev["days_since_last_order"] > 90).astype(int)
print(f"\ntoday's 90-day rule: flags {rule.mean():.0%} of customers, precision {precision_score(y_dev, rule):.2f}, "
      f"recall {recall_score(y_dev, rule):.2f}")
print(f"recency alone as a score: AUC {roc_auc_score(y_dev, X_dev['days_since_last_order']):.3f}\n")

prep = ColumnTransformer([
    ("num", make_pipeline(SimpleImputer(strategy="median", add_indicator=True), StandardScaler()), numeric),
    ("cat", make_pipeline(SimpleImputer(missing_values="", strategy="constant", fill_value="unknown"),
                          OneHotEncoder(handle_unknown="ignore")), categorical),
])
candidates = {
    "logistic regression": LogisticRegression(C=0.1, max_iter=1000),
    "random forest": RandomForestClassifier(n_estimators=200, min_samples_leaf=10, random_state=0),
    "gradient boosting": HistGradientBoostingClassifier(max_depth=2, learning_rate=0.05, random_state=0),
}
for name, model in candidates.items():
    scores = cross_val_score(Pipeline([("prep", prep), ("model", model)]), X_dev, y_dev, cv=5, scoring="roc_auc")
    print(f"{name:<20} cross-validated AUC {scores.mean():.3f} ± {scores.std():.3f}")

model = Pipeline([("prep", prep), ("model", LogisticRegression(C=0.1, max_iter=1000))]).fit(X_dev, y_dev)
p = model.predict_proba(X_test)[:, 1]
y = y_test.to_numpy()
print(f"test AUC: {roc_auc_score(y, p):.3f}\n")

order = np.argsort(-p)
for k in [15, 30, 50]:
    rate = y[order[:k]].mean()
    print(f"top {k:>2} by risk: churn rate {rate:.0%} (lift {rate / y.mean():.1f}x), {int(y[order[:k]].sum())} of {int(y.sum())} churners")

OFFER, VALUE = 15, 80
def realized(contacted):
    return VALUE * int(y[contacted].sum()) - OFFER * int(contacted.sum())
plans = {
    "contact everyone": np.ones(len(y), dtype=bool),
    "today's 90-day rule": X_test["days_since_last_order"].to_numpy() > 90,
    "model: p x $80 > $15": p * VALUE > OFFER,
}
print()
for name, contacted in plans.items():
    print(f"{name:<22} contacts {int(contacted.sum()):>3}   net value on the test set: ${realized(contacted):>5}")

imp = permutation_importance(model, X_test, y_test, scoring="roc_auc", n_repeats=10, random_state=0)
ranked = pd.Series(imp.importances_mean, index=X_test.columns).sort_values(ascending=False)
print("AUC lost when each column is shuffled:")
print(ranked.head(5).round(3).to_string())

p = model.predict_proba(X_test)[:, 1]
table = pd.DataFrame({"predicted": p, "churned": y_test.to_numpy()})
table["band"] = pd.cut(table["predicted"], [0, 0.1, 0.2, 0.4, 1.0])
print("\ncalibration on the test set:")
print(table.groupby("band", observed=True).agg(customers=("churned", "size"),
      predicted=("predicted", "mean"), actual=("churned", "mean")).round(2).to_string())

def lift_at_k(y_true, probs, k):
    y_true, probs = np.asarray(y_true), np.asarray(probs)
    top = np.argsort(-probs)[:k]
    return float(y_true[top].mean() / y_true.mean())

def choose_contacts(probs, offer_cost, value):
    return np.asarray(probs) * value > offer_cost

def campaign_value(y_true, contacted, offer_cost, value):
    y_true, contacted = np.asarray(y_true), np.asarray(contacted, dtype=bool)
    return int(value * y_true[contacted].sum() - offer_cost * contacted.sum())