Class Imbalance

"""One model, one set of scores, four different class mixes.

The model is trained once and never touched again. All that changes is how many
negatives are kept in the evaluation set — so the model's ability to rank is
identical in every row, by construction.

Watch the three metrics disagree about what happened:
  majority_acc    what you get by predicting "no" every time — it goes UP
  roc_auc         flat, because it is a property of the ranking alone
  avg_precision   falls by half, because finding the positives really does get
                  harder as they get rarer

Printed as a markdown table, in identifiers only, so one artifact serves both
the English and the Portuguese page.
"""

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

N, DIM, TRAIN = 60_000, 10, 20_000

rng = np.random.default_rng(11)
direction = rng.normal(size=DIM)
X = rng.normal(size=(N, DIM))
risk = X @ direction + 1.0 * rng.normal(size=N)
y = (risk > np.quantile(risk, 0.70)).astype(int)

fitted = make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000)).fit(X[:TRAIN], y[:TRAIN])
scores = fitted.predict_proba(X[TRAIN:])[:, 1]
held = y[TRAIN:]

positives, negatives = np.flatnonzero(held == 1), np.flatnonzero(held == 0)
draw = np.random.default_rng(0)

print("| `prevalence` | `majority_class_acc` | `roc_auc` | `avg_precision` |")
print("|---:|---:|---:|---:|")
for target in (0.30, 0.10, 0.02, 0.005):
    keep = min(len(positives), int(target * len(negatives) / (1 - target)))
    rows = np.concatenate([draw.choice(positives, keep, replace=False), negatives])
    subset, p = held[rows], scores[rows]
    print(f"| {subset.mean():.3f} | {max(subset.mean(), 1 - subset.mean()):.3f} "
          f"| {roc_auc_score(subset, p):.3f} | **{average_precision_score(subset, p):.3f}** |")

"""Four ways to "handle imbalance", and what each one actually changes.

A 2% positive rate, one model family, four treatments. Read the table by
column rather than by row:

  roc_auc         barely moves — which is the first thing to notice about it
  avg_precision   barely moves either: none of these improves the ranking
  f1@0.5          moves a lot, and downward for the resampled variants
  f1@best         back to level once the threshold is tuned for each
  brier, mean_pred  wrecked by weighting and by oversampling, and restored by
                  the prior correction on the last row

Logistic regression rather than a network, because `class_weight` is available
and the run is deterministic. For a network the equivalent of class weights is
`pos_weight` in `BCEWithLogitsLoss`, with exactly the same consequence for the
probabilities it outputs.

Printed as a markdown table, in identifiers only, so one artifact serves both
the English and the Portuguese page.
"""

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (average_precision_score, brier_score_loss,
                             f1_score, roc_auc_score)
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

N, DIM, POSITIVE = 20_000, 10, 0.02

rng = np.random.default_rng(11)
direction = rng.normal(size=DIM)
X = rng.normal(size=(N, DIM))
risk = X @ direction + 1.0 * rng.normal(size=N)
y = (risk > np.quantile(risk, 1 - POSITIVE)).astype(int)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=0
)

model = lambda **kw: make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000, **kw))


def row(name, p):
    grid = np.quantile(p, np.linspace(0.90, 0.9999, 300))
    best = max(f1_score(y_test, p >= t) for t in grid)
    print(f"| `{name}` | {roc_auc_score(y_test, p):.3f} | **{average_precision_score(y_test, p):.3f}** "
          f"| {f1_score(y_test, p >= 0.5):.3f} | **{best:.3f}** "
          f"| {brier_score_loss(y_test, p):.4f} | {p.mean():.3f} |")


print("| `strategy` | `roc_auc` | `avg_precision` | `f1@0.5` | `f1@best` | `brier` | `mean_pred` |")
print("|---|---:|---:|---:|---:|---:|---:|")

plain = model().fit(X_train, y_train).predict_proba(X_test)[:, 1]
row("baseline", plain)
row("threshold_tuned", plain)                 # the same model — only the cut moves

weighted = model(class_weight="balanced").fit(X_train, y_train)
row("class_weight_balanced", weighted.predict_proba(X_test)[:, 1])

minority, majority = np.flatnonzero(y_train == 1), np.flatnonzero(y_train == 0)
copies = np.random.default_rng(0).choice(minority, len(majority) - len(minority), replace=True)
index = np.concatenate([np.arange(len(y_train)), copies])
resampled = model().fit(X_train[index], y_train[index])
p_over = resampled.predict_proba(X_test)[:, 1]
row("random_oversample", p_over)

# Undo the base rate the resampling invented (Elkan / King & Zeng): shift the odds
# by the ratio of the true prior to the training prior.
prior_train, prior_true = y_train[index].mean(), y_train.mean()
odds = p_over / (1 - p_over) * (prior_true / (1 - prior_true)) * ((1 - prior_train) / prior_train)
row("oversample_then_corrected", odds / (1 + odds))