imbalance_report.py

#!/usr/bin/env python3
"""Profile class imbalance in a CSV target and recommend a handling strategy, cheapest first.

Reports class counts and shares, the imbalance ratio, minority count, how many minority examples
each CV fold would get, and the severity:
  mild      minority >= 20%                     usually nothing special beyond good metrics
  moderate  5% - 20%                            metrics + stratified CV + threshold tuning
  severe    0.5% - 5%                           + class weights or focal-style losses; PR-AUC; capacity metrics
  extreme   < 0.5% or < 100 minority examples   + consider anomaly detection, collecting labels, rules
Then the strategy ladder (each step only if the previous is not enough):
  1 metrics (PR-AUC, recall@K, cost) and a natural-prevalence test set
  2 stratified / grouped CV with enough minority rows per fold
  3 threshold from capacity or costs (threshold_tuner.py)
  4 class weights (class_weight="balanced", scale_pos_weight) - recalibrate if probabilities are used
  5 resampling inside an imblearn Pipeline (random undersampling of negatives for speed,
    SMOTE / ADASYN rarely better than weights for GBMs)
  6 different framing (anomaly detection, two-stage models) or more labels
and, with --code, an imblearn/scikit-learn snippet for steps 4-5.

Usage:
    python imbalance_report.py data.csv --target is_fraud
    python imbalance_report.py data.csv --target churned --folds 5 --model gbm --code
    python imbalance_report.py data.csv --target segment --json        # multiclass

Exit codes: 0 report printed, 2 usage or read error.
"""

import argparse
import csv
import json
import sys
from collections import Counter

LADDER = [
    ("metrics", "Use PR-AUC / recall at the operating point / expected cost; keep the test set at the natural prevalence"),
    ("cv", "Stratified (or StratifiedGroup) K-fold so every fold has enough minority rows; repeat CV when folds are small"),
    ("threshold", "Pick the threshold from capacity or costs on validation data (threshold_tuner.py), not 0.5"),
    ("weights", "Class weights (class_weight='balanced', scale_pos_weight, auto_class_weights) - cheap, no data copies; "
                "probabilities become biased: recalibrate or prior-correct"),
    ("resample", "Resample only the training folds inside an imblearn Pipeline; undersampling negatives speeds training; "
                 "SMOTE helps mainly linear / kNN models on dense numeric data"),
    ("reframe", "Too few positives to learn from: anomaly detection, rules + label collection, or a two-stage model"),
]


def main():
    ap = argparse.ArgumentParser(description="Profile class imbalance and recommend a strategy.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("--target", required=True)
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--model", choices=["linear", "gbm", "rf", "knn"], default="gbm")
    ap.add_argument("--code", action="store_true")
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args()
    try:
        with (sys.stdin if a.csv == "-" else open(a.csv, newline="", encoding="utf-8-sig")) as fh:
            rd = csv.DictReader(fh)
            if a.target not in (rd.fieldnames or []):
                raise ValueError(f"target {a.target!r} not found")
            counts = Counter((r[a.target] or "").strip() for r in rd if (r[a.target] or "").strip())
    except (OSError, ValueError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    if len(counts) < 2:
        print("error: need at least two classes", file=sys.stderr)
        return 2
    n = sum(counts.values())
    minority, m_count = min(counts.items(), key=lambda kv: kv[1])
    majority, M_count = max(counts.items(), key=lambda kv: kv[1])
    share = m_count / n
    per_fold = m_count / a.folds
    if share < 0.005 or m_count < 100:
        sev, steps = "extreme", ["metrics", "cv", "threshold", "weights", "reframe"]
    elif share < 0.05:
        sev, steps = "severe", ["metrics", "cv", "threshold", "weights", "resample"]
    elif share < 0.2:
        sev, steps = "moderate", ["metrics", "cv", "threshold"]
    else:
        sev, steps = "mild", ["metrics", "cv"]
    notes = []
    if per_fold < 30:
        notes.append(f"only ~{per_fold:.0f} minority rows per fold: use fewer folds, RepeatedStratifiedKFold, and report "
                     "confidence intervals")
    if len(counts) > 2:
        rare = [c for c, k in counts.items() if k < 50]
        if rare:
            notes.append(f"{len(rare)} class(es) with < 50 rows ({', '.join(rare[:5])}): merge, collect more, or evaluate "
                         "with per-class recall and macro averages")
    if a.model == "gbm" and "resample" in steps:
        notes.append("for gradient boosting, class weights or plain unweighted training + threshold usually match or "
                     "beat SMOTE; undersample negatives only to cut training time (then prior-correct)")
    if a.model in ("linear", "knn") and "resample" in steps:
        notes.append("linear / kNN models can benefit from SMOTE on dense numeric features; compare with "
                     "class_weight='balanced' on the same folds")
    out = {"rows": n, "classes": dict(counts.most_common()), "minority": minority, "minority_count": m_count,
           "minority_share": share, "imbalance_ratio": M_count / m_count, "minority_per_fold": per_fold,
           "severity": sev, "ladder": [{"step": k, "action": v} for k, v in LADDER if k in steps], "notes": notes}
    if a.code:
        spw = round(M_count / m_count, 2)
        # sampling_strategy = minority/majority ratio after undersampling; must exceed the current ratio
        ratio_now = m_count / M_count
        target = round(min(1.0, max(0.1, 2 * ratio_now)), 3)
        print(f'''# Imbalance handling for {a.target!r}: minority {minority!r} = {share:.2%} ({sev})
from imblearn.pipeline import Pipeline                 # pip install imbalanced-learn; samplers run on training folds only
from imblearn.under_sampling import RandomUnderSampler
from imblearn.over_sampling import SMOTE
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_score

cv = StratifiedKFold(n_splits={a.folds}, shuffle=True, random_state=0)

# Step 4: class weights (no resampling)
weighted = Pipeline([("prep", prep), ("model", LogisticRegression(class_weight="balanced", max_iter=5000))])
# LightGBM: LGBMClassifier(scale_pos_weight={spw})   XGBoost: XGBClassifier(scale_pos_weight={spw})
# CatBoost: CatBoostClassifier(auto_class_weights="Balanced")

# Step 5: resampling inside the pipeline (never before the split)
under = Pipeline([("prep", prep), ("under", RandomUnderSampler(sampling_strategy={target}, random_state=0)),   # now {ratio_now:.3f}
                  ("model", LogisticRegression(max_iter=5000))])
smote = Pipeline([("prep", prep), ("smote", SMOTE(k_neighbors=5, random_state=0)),
                  ("model", LogisticRegression(max_iter=5000))])

for name, pipe in {{"weighted": weighted, "under": under, "smote": smote}}.items():
    print(name, cross_val_score(pipe, X, y, cv=cv, scoring="average_precision").mean().round(4))
# Undersampling changes probabilities: training positive rate becomes {target / (1 + target):.4f};
# correct with threshold_tuner.py --train-rate {target / (1 + target):.4f} --true-rate {share:.4f}
''')
        return 0
    if a.json:
        print(json.dumps(out, indent=2))
        return 0
    print(f"Target {a.target!r}: {len(counts)} classes, {n} rows")
    for c, k in counts.most_common(10):
        print(f"  {c!s:<16} {k:>8}  {k / n:7.2%}")
    print(f"minority {minority!r}: {m_count} rows ({share:.2%}), imbalance ratio {M_count / m_count:.1f}:1, "
          f"~{per_fold:.0f} per fold with {a.folds} folds -> severity: {sev.upper()}")
    print("\nStrategy (stop at the first step that meets the goal):")
    for i, st in enumerate(out["ladder"], 1):
        print(f"  {i}. {st['action']}")
    for nt in notes:
        print(f"note: {nt}")
    return 0


if __name__ == "__main__":
    sys.exit(main())