benchmark_quick.py

#!/usr/bin/env python3
"""
Quick benchmark: HyperOpt-GBT (Python) vs XGBoost vs LightGBM vs CatBoost

Runs on synthetic datasets to validate accuracy and speed of the core
innovations: GOSS, quantile sketch binning, and histogram-based splits.

Usage:
    pip install hyperopt-gbt xgboost lightgbm catboost scikit-learn
    python benchmark_quick.py
"""

import time
import warnings
import numpy as np
from sklearn.datasets import make_classification, make_regression
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, root_mean_squared_error

warnings.filterwarnings("ignore")

# ── Helpers ──────────────────────────────────────────────────────────────────

def timer(func):
    """Time a callable, return (result, elapsed_seconds)."""
    t0 = time.perf_counter()
    result = func()
    return result, time.perf_counter() - t0


def print_table(headers, rows):
    """Pretty-print a markdown-style table."""
    widths = [max(len(h), max((len(str(r[i])) for r in rows), default=0))
              for i, h in enumerate(headers)]
    fmt = " | ".join(f"{{:<{w}}}" for w in widths)
    sep = "-|-".join("-" * w for w in widths)
    print(fmt.format(*headers))
    print(sep)
    for row in rows:
        print(fmt.format(*[str(c) for c in row]))
    print()


# ── Optional imports (graceful degradation) ──────────────────────────────────

def try_import(name):
    try:
        return __import__(name)
    except ImportError:
        return None

xgb = try_import("xgboost")
lgb = try_import("lightgbm")
cb  = try_import("catboost")

# Always import our library
from hyperopt_gbt import HyperOptGradientBoostedClassifier, HyperOptGradientBoostedRegressor

# Try Rust backend
try:
    import rust_gbt as rgbt
    HAS_RUST = True
except ImportError:
    HAS_RUST = False


# =============================================================================
# BENCHMARK 1: Large-scale binary classification
# =============================================================================

def benchmark_classification(n_train=80_000, n_test=20_000, n_features=30,
                              n_trees=50, seed=42):
    print("=" * 72)
    print(f"BENCHMARK 1: Binary Classification  ({n_train:,} train / {n_test:,} test / "
          f"{n_features} features / {n_trees} trees)")
    print("=" * 72)

    rng = np.random.RandomState(seed)

    # Nonlinear synthetic data
    X = rng.randn(n_train + n_test, n_features)
    signal = (X[:, 0] * X[:, 1]
              + np.sin(X[:, 2]) * 2
              + (X[:, 3] > 0).astype(float) * 1.5
              + rng.randn(n_train + n_test) * 0.5)
    y = (signal > np.median(signal)).astype(float)

    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]

    rows = []

    # ── HyperOpt-GBT (Python, GOSS) ─────────────────────────────────────────
    clf = HyperOptGradientBoostedClassifier(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        use_goss=True, goss_a=0.2, goss_b=0.1,
        n_bins=255, binning="uniform", random_state=seed,
    )
    _, train_time = timer(lambda: clf.fit(X_train, y_train))
    proba, pred_time = timer(lambda: clf.predict_proba(X_test)[:, 1])
    auc = roc_auc_score(y_test, proba)
    rows.append(["HyperOpt-GBT (GOSS)", f"{auc:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── HyperOpt-GBT (Python, no GOSS) ──────────────────────────────────────
    clf2 = HyperOptGradientBoostedClassifier(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        use_goss=False, n_bins=255, binning="uniform", random_state=seed,
    )
    _, train_time = timer(lambda: clf2.fit(X_train, y_train))
    proba2, pred_time = timer(lambda: clf2.predict_proba(X_test)[:, 1])
    auc2 = roc_auc_score(y_test, proba2)
    rows.append(["HyperOpt-GBT (no GOSS)", f"{auc2:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── HyperOpt-GBT (quantile sketch) ──────────────────────────────────────
    clf3 = HyperOptGradientBoostedClassifier(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        use_goss=True, goss_a=0.2, goss_b=0.1,
        n_bins=255, binning="quantile_sketch", random_state=seed,
    )
    _, train_time = timer(lambda: clf3.fit(X_train, y_train))
    proba3, pred_time = timer(lambda: clf3.predict_proba(X_test)[:, 1])
    auc3 = roc_auc_score(y_test, proba3)
    rows.append(["HyperOpt-GBT (quantile)", f"{auc3:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── Rust backend ─────────────────────────────────────────────────────────
    if HAS_RUST:
        model = rgbt.PyRustGBT()
        _, train_time = timer(lambda: model.fit(
            X_train, y_train, n_estimators=n_trees, learning_rate=0.1,
            max_depth=6, n_bins=255, use_goss=True, goss_a=0.2, goss_b=0.1,
            task="classification", verbose=False,
        ))
        proba_r, pred_time = timer(lambda: model.predict_proba(X_test))
        auc_r = roc_auc_score(y_test, np.asarray(proba_r))
        rows.append(["Rust-GBT (GOSS)", f"{auc_r:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── XGBoost ──────────────────────────────────────────────────────────────
    if xgb:
        xgb_clf = xgb.XGBClassifier(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            tree_method="hist", random_state=seed, verbosity=0,
        )
        _, train_time = timer(lambda: xgb_clf.fit(X_train, y_train))
        proba_x, pred_time = timer(lambda: xgb_clf.predict_proba(X_test)[:, 1])
        auc_x = roc_auc_score(y_test, proba_x)
        rows.append(["XGBoost (hist)", f"{auc_x:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── LightGBM ─────────────────────────────────────────────────────────────
    if lgb:
        lgb_clf = lgb.LGBMClassifier(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            random_state=seed, verbose=-1,
        )
        _, train_time = timer(lambda: lgb_clf.fit(X_train, y_train))
        proba_l, pred_time = timer(lambda: lgb_clf.predict_proba(X_test)[:, 1])
        auc_l = roc_auc_score(y_test, proba_l)
        rows.append(["LightGBM", f"{auc_l:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # ── CatBoost ─────────────────────────────────────────────────────────────
    if cb:
        cb_clf = cb.CatBoostClassifier(
            iterations=n_trees, learning_rate=0.1, depth=6,
            random_seed=seed, verbose=0,
        )
        _, train_time = timer(lambda: cb_clf.fit(X_train, y_train))
        proba_c, pred_time = timer(lambda: cb_clf.predict_proba(X_test)[:, 1])
        auc_c = roc_auc_score(y_test, proba_c)
        rows.append(["CatBoost", f"{auc_c:.4f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    print()
    print_table(["Library", "AUC", "Train Time", "Predict Time"], rows)


# =============================================================================
# BENCHMARK 2: GOSS ablation
# =============================================================================

def benchmark_goss_ablation(n_train=80_000, n_test=20_000, n_features=30,
                            n_trees=50, seed=42):
    print("=" * 72)
    print("BENCHMARK 2: GOSS Ablation")
    print("=" * 72)

    rng = np.random.RandomState(seed)
    X = rng.randn(n_train + n_test, n_features)
    signal = (X[:, 0] * X[:, 1]
              + np.sin(X[:, 2]) * 2
              + (X[:, 3] > 0).astype(float) * 1.5
              + rng.randn(n_train + n_test) * 0.5)
    y = (signal > np.median(signal)).astype(float)
    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]

    configs = [
        ("Full data (no GOSS)", False, 0.0, 0.0, "100%"),
        ("GOSS a=0.3, b=0.1",  True,  0.3, 0.1, "40%"),
        ("GOSS a=0.2, b=0.1",  True,  0.2, 0.1, "30%"),
        ("GOSS a=0.1, b=0.05", True,  0.1, 0.05, "15%"),
    ]

    baseline_time = None
    rows = []

    for name, use_goss, a, b, data_pct in configs:
        clf = HyperOptGradientBoostedClassifier(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            use_goss=use_goss, goss_a=a, goss_b=b,
            n_bins=255, random_state=seed,
        )
        _, train_time = timer(lambda: clf.fit(X_train, y_train))
        proba = clf.predict_proba(X_test)[:, 1]
        auc = roc_auc_score(y_test, proba)

        if baseline_time is None:
            baseline_time = train_time
        speedup = baseline_time / train_time if train_time > 0 else float("inf")

        rows.append([name, data_pct, f"{auc:.4f}", f"{train_time:.2f}s", f"{speedup:.1f}x"])

    print()
    print_table(["Configuration", "Data Used", "AUC", "Train Time", "Speedup"], rows)


# =============================================================================
# BENCHMARK 3: Quantile sketch vs uniform on skewed data
# =============================================================================

def benchmark_quantile_sketch(n_train=40_000, n_test=10_000, n_trees=50, seed=42):
    print("=" * 72)
    print("BENCHMARK 3: Quantile Sketch vs Uniform Binning (Skewed Data)")
    print("=" * 72)

    rng = np.random.RandomState(seed)
    n_total = n_train + n_test
    n_features = 10

    # Create highly skewed features: 85% in [0, 0.5], 15% outliers at ~50-100
    X = np.zeros((n_total, n_features))
    for f in range(n_features):
        mask = rng.rand(n_total) < 0.85
        X[mask, f] = rng.exponential(0.1, mask.sum())
        X[~mask, f] = rng.uniform(50, 100, (~mask).sum())

    # Target depends on the dense region
    signal = X[:, 0] * 3 + np.sin(X[:, 1] * 10) + (X[:, 2] > 0.3).astype(float) * 2
    y = (signal > np.median(signal)).astype(float)

    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]

    rows = []
    for n_bins in [31, 63, 127, 255]:
        # Uniform
        clf_u = HyperOptGradientBoostedClassifier(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            n_bins=n_bins, binning="uniform", use_goss=False, random_state=seed,
        )
        clf_u.fit(X_train, y_train)
        auc_u = roc_auc_score(y_test, clf_u.predict_proba(X_test)[:, 1])

        # Quantile sketch
        clf_q = HyperOptGradientBoostedClassifier(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            n_bins=n_bins, binning="quantile_sketch", use_goss=False, random_state=seed,
        )
        clf_q.fit(X_train, y_train)
        auc_q = roc_auc_score(y_test, clf_q.predict_proba(X_test)[:, 1])

        gain = auc_q - auc_u
        rows.append([str(n_bins), f"{auc_u:.4f}", f"{auc_q:.4f}", f"+{gain:.4f}"])

    print()
    print_table(["Bins", "Uniform AUC", "Quantile AUC", "Gain"], rows)


# =============================================================================
# BENCHMARK 4: Regression (California Housing style)
# =============================================================================

def benchmark_regression(n_train=20_000, n_test=5_000, n_features=8,
                         n_trees=100, seed=42):
    print("=" * 72)
    print(f"BENCHMARK 4: Regression ({n_train:,} train / {n_test:,} test)")
    print("=" * 72)

    X, y = make_regression(
        n_samples=n_train + n_test,
        n_features=n_features,
        n_informative=6,
        noise=10.0,
        random_state=seed,
    )
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=n_test, random_state=seed
    )

    rows = []

    # HyperOpt-GBT
    reg = HyperOptGradientBoostedRegressor(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        use_goss=True, goss_a=0.2, goss_b=0.1,
        n_bins=255, random_state=seed,
    )
    _, train_time = timer(lambda: reg.fit(X_train, y_train))
    pred, pred_time = timer(lambda: reg.predict(X_test))
    rmse = root_mean_squared_error(y_test, pred)
    rows.append(["HyperOpt-GBT (GOSS)", f"{rmse:.2f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    # HyperOpt-GBT quantile
    reg_q = HyperOptGradientBoostedRegressor(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        use_goss=True, goss_a=0.2, goss_b=0.1,
        n_bins=255, binning="quantile_sketch", random_state=seed,
    )
    _, train_time = timer(lambda: reg_q.fit(X_train, y_train))
    pred_q, pred_time = timer(lambda: reg_q.predict(X_test))
    rmse_q = root_mean_squared_error(y_test, pred_q)
    rows.append(["HyperOpt-GBT (quantile)", f"{rmse_q:.2f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    if xgb:
        xgb_reg = xgb.XGBRegressor(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            tree_method="hist", random_state=seed, verbosity=0,
        )
        _, train_time = timer(lambda: xgb_reg.fit(X_train, y_train))
        pred_x, pred_time = timer(lambda: xgb_reg.predict(X_test))
        rmse_x = root_mean_squared_error(y_test, pred_x)
        rows.append(["XGBoost", f"{rmse_x:.2f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    if lgb:
        lgb_reg = lgb.LGBMRegressor(
            n_estimators=n_trees, learning_rate=0.1, max_depth=6,
            random_state=seed, verbose=-1,
        )
        _, train_time = timer(lambda: lgb_reg.fit(X_train, y_train))
        pred_l, pred_time = timer(lambda: lgb_reg.predict(X_test))
        rmse_l = root_mean_squared_error(y_test, pred_l)
        rows.append(["LightGBM", f"{rmse_l:.2f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    if cb:
        cb_reg = cb.CatBoostRegressor(
            iterations=n_trees, learning_rate=0.1, depth=6,
            random_seed=seed, verbose=0,
        )
        _, train_time = timer(lambda: cb_reg.fit(X_train, y_train))
        pred_c, pred_time = timer(lambda: cb_reg.predict(X_test))
        rmse_c = root_mean_squared_error(y_test, pred_c)
        rows.append(["CatBoost", f"{rmse_c:.2f}", f"{train_time:.2f}s", f"{pred_time*1e3:.0f}ms"])

    print()
    print_table(["Library", "RMSE", "Train Time", "Predict Time"], rows)


# =============================================================================
# BENCHMARK 5: Inference engine comparison
# =============================================================================

def benchmark_inference_engines(n_train=20_000, n_test=50_000, n_trees=50, seed=42):
    print("=" * 72)
    print(f"BENCHMARK 5: Inference Engine Comparison ({n_test:,} test samples)")
    print("=" * 72)

    from hyperopt_gbt.inference import (
        compile_inference_engine,
        NaiveEngine,
        FlatTreeEngine,
        BatchedSIMDEngine,
        QuickScorerEngine,
    )

    rng = np.random.RandomState(seed)
    X = rng.randn(n_train + n_test, 20)
    signal = X[:, 0] * X[:, 1] + np.sin(X[:, 2]) * 2 + rng.randn(n_train + n_test) * 0.3
    y = (signal > np.median(signal)).astype(float)
    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]

    clf = HyperOptGradientBoostedClassifier(
        n_estimators=n_trees, learning_rate=0.1, max_depth=6,
        n_bins=255, random_state=seed,
    )
    clf.fit(X_train, y_train)

    # Bin test data
    X_test_binned = clf._transform_to_bins(X_test)

    rows = []
    engines = [
        ("Naive", NaiveEngine(clf.trees_)),
        ("Flat Tree", FlatTreeEngine(clf.trees_, clf.n_bins)),
        ("Batched SIMD", BatchedSIMDEngine(clf.trees_, clf.n_bins)),
        ("QuickScorer", QuickScorerEngine(clf.trees_, clf.n_bins)),
    ]

    for name, engine in engines:
        # Warmup
        _ = engine.predict(X_test_binned[:100])

        _, elapsed = timer(lambda: engine.predict(X_test_binned))
        throughput = n_test / elapsed
        rows.append([name, f"{elapsed*1e3:.1f}ms", f"{throughput:,.0f} samples/s"])

    # sklearn predict for reference
    _, elapsed = timer(lambda: clf.predict_proba(X_test))
    throughput = n_test / elapsed
    rows.append(["sklearn predict_proba", f"{elapsed*1e3:.1f}ms", f"{throughput:,.0f} samples/s"])

    print()
    print_table(["Engine", "Latency", "Throughput"], rows)


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    print()
    print("╔══════════════════════════════════════════════════════════════════════╗")
    print("║          HyperOpt-GBT — Quick Benchmark Suite                      ║")
    print("╠══════════════════════════════════════════════════════════════════════╣")
    print(f"║  Rust backend:   {'AVAILABLE' if HAS_RUST else 'not found (pip install maturin && cd rust_gbt && maturin develop --release)':55s} ║")
    print(f"║  XGBoost:        {'AVAILABLE' if xgb else 'not installed':55s} ║")
    print(f"║  LightGBM:       {'AVAILABLE' if lgb else 'not installed':55s} ║")
    print(f"║  CatBoost:       {'AVAILABLE' if cb else 'not installed':55s} ║")
    print("╚══════════════════════════════════════════════════════════════════════╝")
    print()

    benchmark_classification()
    benchmark_goss_ablation()
    benchmark_quantile_sketch()
    benchmark_regression()
    benchmark_inference_engines()

    print("=" * 72)
    print("All benchmarks complete.")
    print("=" * 72)