""" Ablation Studies Module ======================== Evaluates the impact of: 1. Feature engineering strategies (which feature groups matter) 2. Spatial encoding methods (spatial lag, Fourier, graph, clusters) 3. Clustering granularity (K=5,10,20,50 for geographic clusters) 4. Transformer embeddings vs. pure GBDT 5. Calibration methods (Platt vs. Isotonic vs. Temperature vs. Ensemble) """ import numpy as np import pandas as pd from sklearn.model_selection import StratifiedKFold from sklearn.metrics import roc_auc_score import json from copy import deepcopy from hybrid_model import ( HybridModel, evaluate_model, prepare_data, FEATURE_GROUPS, get_feature_columns, train_xgboost, train_lightgbm ) from calibration import ( PlattScaling, TemperatureScaling, IsotonicCalibration, EnsembleCalibration, expected_calibration_error, calibration_report ) from sklearn.preprocessing import QuantileTransformer def run_ablation_feature_groups(df, target_col='is_hotspot', n_splits=3): """Ablation 1: Impact of each feature group. Tests: all features, drop-one-group, only-one-group. """ print("\n" + "="*80) print("ABLATION 1: Feature Group Impact Analysis") print("="*80) all_feature_cols = get_feature_columns(df, include_spatial=False, include_fourier=False, include_graph=False, include_clusters=False) X_all = df[all_feature_cols].values y = df[target_col].values.astype(int) qt = QuantileTransformer(output_distribution='normal', random_state=42) results = {} # Baseline: all feature groups print("\n--- All Base Features ---") metrics = _cross_validate(X_all, y, qt, n_splits) results['all_features'] = metrics print(f" AUC: {metrics['auc_roc']:.4f} | F1: {metrics['f1_macro']:.4f} | Acc: {metrics['accuracy']:.4f}") # Drop-one-group analysis print("\n--- Drop-One-Group Analysis ---") for group_name, group_cols in FEATURE_GROUPS.items(): remaining_cols = [c for c in all_feature_cols if c not in group_cols] if len(remaining_cols) == 0: continue X_drop = df[remaining_cols].values metrics = _cross_validate(X_drop, y, qt, n_splits) results[f'drop_{group_name}'] = metrics delta_auc = results['all_features']['auc_roc'] - metrics['auc_roc'] print(f" Drop {group_name:15s}: AUC={metrics['auc_roc']:.4f} " f"(Δ={delta_auc:+.4f}) | F1={metrics['f1_macro']:.4f}") # Only-one-group analysis print("\n--- Single Feature Group Performance ---") for group_name, group_cols in FEATURE_GROUPS.items(): valid_cols = [c for c in group_cols if c in df.columns] if len(valid_cols) == 0: continue X_single = df[valid_cols].values metrics = _cross_validate(X_single, y, qt, n_splits) results[f'only_{group_name}'] = metrics print(f" Only {group_name:15s}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}") return results def run_ablation_spatial_encoding(df, target_col='is_hotspot', n_splits=3): """Ablation 2: Impact of spatial encoding methods. Tests: no spatial, only spatial lag, only Fourier, only graph, only clusters, all spatial. """ print("\n" + "="*80) print("ABLATION 2: Spatial Encoding Method Impact") print("="*80) y = df[target_col].values.astype(int) qt = QuantileTransformer(output_distribution='normal', random_state=42) results = {} configs = { 'no_spatial': dict(include_spatial=False, include_fourier=False, include_graph=False, include_clusters=False), 'spatial_lag_only': dict(include_spatial=True, include_fourier=False, include_graph=False, include_clusters=False), 'fourier_only': dict(include_spatial=False, include_fourier=True, include_graph=False, include_clusters=False), 'graph_only': dict(include_spatial=False, include_fourier=False, include_graph=True, include_clusters=False), 'clusters_only': dict(include_spatial=False, include_fourier=False, include_graph=False, include_clusters=True), 'all_spatial': dict(include_spatial=True, include_fourier=True, include_graph=True, include_clusters=True), } for config_name, config in configs.items(): feature_cols = get_feature_columns(df, **config) X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values metrics = _cross_validate(X, y, qt, n_splits) results[config_name] = metrics print(f" {config_name:25s}: AUC={metrics['auc_roc']:.4f} | " f"F1={metrics['f1_macro']:.4f} | Features={len(feature_cols)}") return results def run_ablation_clustering_granularity(df, target_col='is_hotspot', n_splits=3): """Ablation 3: Impact of clustering granularity (K=5,10,20,50).""" print("\n" + "="*80) print("ABLATION 3: Clustering Granularity Impact") print("="*80) y = df[target_col].values.astype(int) qt = QuantileTransformer(output_distribution='normal', random_state=42) results = {} base_features = get_feature_columns(df, include_spatial=True, include_fourier=True, include_graph=True, include_clusters=False) for k in [5, 10, 20, 50]: cluster_cols = [c for c in df.columns if f'kmeans_{k}' in c] feature_cols = base_features + cluster_cols feature_cols = [f for f in feature_cols if f in df.columns] X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values metrics = _cross_validate(X, y, qt, n_splits) results[f'k={k}'] = metrics print(f" K={k:3d}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}") # All cluster granularities combined cluster_cols = [c for c in df.columns if 'kmeans' in c or 'dbscan' in c or 'dist_to_center' in c] feature_cols = base_features + cluster_cols feature_cols = list(dict.fromkeys([f for f in feature_cols if f in df.columns])) X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values metrics = _cross_validate(X, y, qt, n_splits) results['all_granularities'] = metrics print(f" ALL : AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}") return results def run_ablation_transformer_embeddings(df, target_col='is_hotspot', n_splits=3): """Ablation 4: Transformer embeddings vs. pure GBDT.""" print("\n" + "="*80) print("ABLATION 4: Transformer Embedding Impact") print("="*80) feature_cols = get_feature_columns(df) X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values y = df[target_col].values.astype(int) results = {} # Pure GBDT (no transformer) print("\n Training pure GBDT (no transformer)...") qt = QuantileTransformer(output_distribution='normal', random_state=42) metrics_gbdt = _cross_validate(X, y, qt, n_splits) results['pure_gbdt'] = metrics_gbdt print(f" Pure GBDT: AUC={metrics_gbdt['auc_roc']:.4f} | F1={metrics_gbdt['f1_macro']:.4f}") # Hybrid with representative transformer configs (reduced for tractability) for d_token, n_layers in [(32, 2), (64, 2), (128, 2)]: config_name = f'hybrid_d{d_token}_L{n_layers}' print(f"\n Training {config_name}...") metrics = _cross_validate_hybrid(X, y, n_splits, d_token=d_token, n_layers=n_layers) results[config_name] = metrics print(f" {config_name}: AUC={metrics['auc_roc']:.4f} | F1={metrics['f1_macro']:.4f}") return results def run_ablation_calibration(df, target_col='is_hotspot', n_splits=3): """Ablation 5: Calibration method comparison.""" print("\n" + "="*80) print("ABLATION 5: Calibration Method Comparison") print("="*80) feature_cols = get_feature_columns(df) X = df[feature_cols].replace([np.inf, -np.inf], np.nan).fillna(0).values y = df[target_col].values.astype(int) qt = QuantileTransformer(output_distribution='normal', random_state=42) # Get uncalibrated predictions via CV skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42) uncal_probs = np.zeros(len(y)) for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)): X_train, X_val = X[train_idx], X[val_idx] y_train, y_val = y[train_idx], y[val_idx] X_train_qt = qt.fit_transform(X_train) X_val_qt = qt.transform(X_val) xgb_model = train_xgboost(X_train_qt, y_train, X_val_qt, y_val) lgb_model = train_lightgbm(X_train_qt, y_train, X_val_qt, y_val) xgb_probs = xgb_model.predict_proba(X_val_qt)[:, 1] lgb_probs = lgb_model.predict_proba(X_val_qt)[:, 1] uncal_probs[val_idx] = 0.5 * xgb_probs + 0.5 * lgb_probs # Now test each calibration method using nested CV results = {} # Uncalibrated ece_uncal, _ = expected_calibration_error(y, uncal_probs) results['uncalibrated'] = { 'ECE': ece_uncal, 'AUC': roc_auc_score(y, uncal_probs), 'Brier': float(np.mean((uncal_probs - y)**2)) } print(f" Uncalibrated: ECE={ece_uncal:.4f} | AUC={results['uncalibrated']['AUC']:.4f}") # Test each calibration method for cal_name, CalClass in [('platt', PlattScaling), ('isotonic', IsotonicCalibration), ('temperature', TemperatureScaling)]: cal_probs = np.zeros(len(y)) skf_cal = StratifiedKFold(n_splits=3, shuffle=True, random_state=123) for _, (cal_train, cal_test) in enumerate(skf_cal.split(uncal_probs.reshape(-1, 1), y)): cal = CalClass() cal.fit(uncal_probs[cal_train], y[cal_train]) cal_probs[cal_test] = cal.predict_proba(uncal_probs[cal_test]) cal_probs = np.clip(cal_probs, 1e-7, 1-1e-7) ece_cal, _ = expected_calibration_error(y, cal_probs) results[cal_name] = { 'ECE': ece_cal, 'AUC': roc_auc_score(y, cal_probs), 'Brier': float(np.mean((cal_probs - y)**2)) } print(f" {cal_name:12s}: ECE={ece_cal:.4f} | AUC={results[cal_name]['AUC']:.4f} | " f"Brier={results[cal_name]['Brier']:.4f}") return results def _cross_validate(X, y, qt, n_splits=5): """Cross-validate with XGBoost+LightGBM ensemble.""" from sklearn.metrics import roc_auc_score skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42) all_metrics = [] for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)): X_train, X_val = X[train_idx], X[val_idx] y_train, y_val = y[train_idx], y[val_idx] qt_fold = QuantileTransformer(output_distribution='normal', random_state=42) X_train_qt = qt_fold.fit_transform(X_train) X_val_qt = qt_fold.transform(X_val) xgb_model = train_xgboost(X_train_qt, y_train, X_val_qt, y_val) lgb_model = train_lightgbm(X_train_qt, y_train, X_val_qt, y_val) xgb_probs = xgb_model.predict_proba(X_val_qt)[:, 1] lgb_probs = lgb_model.predict_proba(X_val_qt)[:, 1] probs = 0.5 * xgb_probs + 0.5 * lgb_probs preds = (probs >= 0.5).astype(int) fold_metrics = evaluate_model(y_val, preds, probs) all_metrics.append(fold_metrics) # Average metrics across folds avg_metrics = {} for key in all_metrics[0]: vals = [m[key] for m in all_metrics] avg_metrics[key] = float(np.mean(vals)) avg_metrics[f'{key}_std'] = float(np.std(vals)) return avg_metrics def _cross_validate_hybrid(X, y, n_splits=5, d_token=64, n_layers=2): """Cross-validate with full hybrid model (transformer + GBDT).""" skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42) all_metrics = [] for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)): X_train, X_val = X[train_idx], X[val_idx] y_train, y_val = y[train_idx], y[val_idx] model = HybridModel( task='classification', use_transformer=True, use_calibration=False, d_token=d_token, n_layers=n_layers ) model.fit(X_train, y_train, X_val, y_val) probs = model.predict_proba(X_val) preds = (probs >= 0.5).astype(int) fold_metrics = evaluate_model(y_val, preds, probs) all_metrics.append(fold_metrics) avg_metrics = {} for key in all_metrics[0]: vals = [m[key] for m in all_metrics] avg_metrics[key] = float(np.mean(vals)) avg_metrics[f'{key}_std'] = float(np.std(vals)) return avg_metrics def run_all_ablations(df, target_col='is_hotspot'): """Run all ablation studies and save results.""" all_results = {} all_results['feature_groups'] = run_ablation_feature_groups(df, target_col) all_results['spatial_encoding'] = run_ablation_spatial_encoding(df, target_col) all_results['clustering_granularity'] = run_ablation_clustering_granularity(df, target_col) all_results['transformer_embeddings'] = run_ablation_transformer_embeddings(df, target_col) all_results['calibration'] = run_ablation_calibration(df, target_col) return all_results