""" calibration.py -------------- Probability calibration for Dataset A transaction model. Compares: 1. Raw XGBoost probabilities 2. Sigmoid calibration (Platt scaling) 3. Isotonic calibration Uses Validation set ONLY for fitting calibration. Evaluates Brier Score, Log Loss, and Expected Calibration Error (ECE). Selects and freezes the best calibration method for Test evaluation. Outputs: - models/transaction_model/calibrated_model.joblib - data/processed/calibration_report.json """ from __future__ import annotations import json import logging from pathlib import Path from typing import Any import joblib import numpy as np import pandas as pd from sklearn.calibration import CalibratedClassifierCV, calibration_curve from sklearn.metrics import brier_score_loss, log_loss ROOT = Path(__file__).resolve().parents[2] DATA_DIR = ROOT / "data" PROCESSED_DIR = DATA_DIR / "processed" MODELS_DIR = ROOT / "models" logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", ) LOGGER = logging.getLogger("probability-calibration") NUMERIC_FEATURES = [ "amount", "amount_log1p", "hour", "day_of_week", "is_weekend", "customer_txn_count_past", "customer_amount_mean_past", "customer_amount_std_past", "device_txn_count_past", "customer_amount_dev", "identity_available", "missing_p_email", "missing_r_email", "missing_addr1", "missing_device_info" ] CATEGORICAL_FEATURES = [ "ProductCD", "card1", "card2", "card3", "card4", "card5", "card6", "addr1", "addr2", "P_emaildomain", "R_emaildomain", "DeviceType", "DeviceInfo" ] def compute_ece(y_true: np.ndarray, y_prob: np.ndarray, n_bins: int = 10) -> float: """Computes Expected Calibration Error (ECE).""" prob_true, prob_pred = calibration_curve(y_true, y_prob, n_bins=n_bins, strategy="uniform") bin_edges = np.linspace(0, 1, n_bins + 1) ece = 0.0 total_samples = len(y_true) for i in range(n_bins): mask = (y_prob >= bin_edges[i]) & (y_prob < bin_edges[i + 1]) bin_size = np.sum(mask) if bin_size > 0: bin_acc = np.mean(y_true[mask]) bin_conf = np.mean(y_prob[mask]) ece += (bin_size / total_samples) * abs(bin_acc - bin_conf) return float(ece) def evaluate_calibration_metrics(y_true: np.ndarray, y_prob: np.ndarray) -> dict[str, float]: """Calculates Brier Score, Log Loss, and ECE.""" brier = float(brier_score_loss(y_true, y_prob)) loss = float(log_loss(y_true, y_prob)) ece = compute_ece(y_true, y_prob) return { "brier_score": round(brier, 6), "log_loss": round(loss, 6), "ece": round(ece, 6), } def fit_and_evaluate_calibration( dataset_a_path: Path | None = None, ) -> dict[str, Any]: if dataset_a_path is None: dataset_a_path = PROCESSED_DIR / "dataset_a_features.parquet" df_a = pd.read_parquet(dataset_a_path) val_df = df_a[df_a["split"] == "validation"].copy() test_df = df_a[df_a["split"] == "test"].copy() tx_model_path = MODELS_DIR / "transaction_model" / "xgboost_model.joblib" encoder_path = MODELS_DIR / "transaction_model" / "encoder.joblib" if not tx_model_path.exists(): raise FileNotFoundError(f"Base transaction model not found at {tx_model_path}") xgb_tx = joblib.load(tx_model_path) encoder = joblib.load(encoder_path) val_cat = encoder.transform(val_df[CATEGORICAL_FEATURES].astype(str)) X_val = np.hstack([val_df[NUMERIC_FEATURES].values.astype(np.float32), val_cat.astype(np.float32)]) y_val = val_df["isFraud"].values.astype(int) test_cat = encoder.transform(test_df[CATEGORICAL_FEATURES].astype(str)) X_test = np.hstack([test_df[NUMERIC_FEATURES].values.astype(np.float32), test_cat.astype(np.float32)]) y_test = test_df["isFraud"].values.astype(int) # 1. Raw XGBoost probabilities val_prob_raw = xgb_tx.predict_proba(X_val)[:, 1] test_prob_raw = xgb_tx.predict_proba(X_test)[:, 1] raw_val_m = evaluate_calibration_metrics(y_val, val_prob_raw) raw_test_m = evaluate_calibration_metrics(y_test, test_prob_raw) # 2. Sigmoid Calibration (Platt scaling fitted strictly on Validation probabilities) LOGGER.info("Fitting Sigmoid probability calibration on Validation set ...") from sklearn.linear_model import LogisticRegression from sklearn.isotonic import IsotonicRegression cal_sigmoid = LogisticRegression(C=1e5, solver="lbfgs") cal_sigmoid.fit(val_prob_raw.reshape(-1, 1), y_val) val_prob_sig = cal_sigmoid.predict_proba(val_prob_raw.reshape(-1, 1))[:, 1] test_prob_sig = cal_sigmoid.predict_proba(test_prob_raw.reshape(-1, 1))[:, 1] sig_val_m = evaluate_calibration_metrics(y_val, val_prob_sig) sig_test_m = evaluate_calibration_metrics(y_test, test_prob_sig) # 3. Isotonic Calibration (fitted strictly on Validation probabilities) LOGGER.info("Fitting Isotonic probability calibration on Validation set ...") cal_isotonic = IsotonicRegression(out_of_bounds="clip") cal_isotonic.fit(val_prob_raw, y_val) val_prob_iso = cal_isotonic.transform(val_prob_raw) test_prob_iso = cal_isotonic.transform(test_prob_raw) iso_val_m = evaluate_calibration_metrics(y_val, val_prob_iso) iso_test_m = evaluate_calibration_metrics(y_test, test_prob_iso) methods = { "raw": {"val": raw_val_m, "test": raw_test_m, "model": xgb_tx}, "sigmoid": {"val": sig_val_m, "test": sig_test_m, "model": cal_sigmoid}, "isotonic": {"val": iso_val_m, "test": iso_test_m, "model": cal_isotonic}, } # Select best calibration method based on Validation Brier Score best_method = min(methods.keys(), key=lambda m: methods[m]["val"]["brier_score"]) LOGGER.info("Selected best calibration method: %s (Val Brier=%.6f, Val LogLoss=%.6f)", best_method, methods[best_method]["val"]["brier_score"], methods[best_method]["val"]["log_loss"]) # Save calibrated model artifact cal_model_path = MODELS_DIR / "transaction_model" / "calibrated_model.joblib" joblib.dump(methods[best_method]["model"], cal_model_path) report = { "dataset": "Dataset A Transaction Model Calibration", "selected_calibration_method": best_method, "methods": { k: {"validation": v["val"], "test": v["test"]} for k, v in methods.items() }, } json_path = PROCESSED_DIR / "calibration_report.json" with json_path.open("w", encoding="utf-8") as f: json.dump(report, f, indent=2) LOGGER.info("Calibration report saved to %s", json_path) return report if __name__ == "__main__": fit_and_evaluate_calibration()