spb/wp3_uqo Public

1# Author: Simon-Pierre Boucher — contact@spboucher.ai2#3"""Shared model-training helpers."""45import numpy as np6from sklearn.metrics import mean_squared_error, r2_score7from sklearn.model_selection import train_test_split8from xgboost import XGBRegressor910from .config import RANDOM_STATE1112XGB_PARAMS = dict(13    n_estimators=1000,14    max_depth=8,15    learning_rate=0.05,16    subsample=0.8,17    colsample_bytree=0.8,18    min_child_weight=5,19    reg_alpha=0.1,20    reg_lambda=1.0,21    random_state=RANDOM_STATE,22    n_jobs=-1,23    early_stopping_rounds=50,24)252627def _fit_and_score(X, y, idx_train, idx_test):28    """Fit XGBoost on one split (10% of train reserved for early stopping)."""29    X_tr, X_te = X[idx_train], X[idx_test]30    y_tr, y_te = y[idx_train], y[idx_test]3132    X_fit, X_eval, y_fit, y_eval = train_test_split(33        X_tr, y_tr, test_size=0.1, random_state=RANDOM_STATE)3435    model = XGBRegressor(**XGB_PARAMS)36    model.fit(X_fit, y_fit, eval_set=[(X_eval, y_eval)], verbose=0)3738    y_pred = model.predict(X_te)39    return {40        "r2": r2_score(y_te, y_pred),41        "rmse": float(np.sqrt(mean_squared_error(y_te, y_pred))),42        "best_iter": getattr(model, "best_iteration", XGB_PARAMS["n_estimators"]),43    }444546def train_xgb(X, y, idx_train, idx_test, idx_geo_train, idx_geo_test):47    """Train XGBoost under both validation schemes and return the metrics.4849    Two independent models are fit: one on the random 80/20 split and one on50    the geographic (state-holdout) split.51    """52    random_split = _fit_and_score(X, y, idx_train, idx_test)53    geo_split = _fit_and_score(X, y, idx_geo_train, idx_geo_test)54    return {55        "r2_random": random_split["r2"],56        "rmse_random": random_split["rmse"],57        "r2_geo": geo_split["r2"],58        "rmse_geo": geo_split["rmse"],59        "best_iter_random": random_split["best_iter"],60        "best_iter_geo": geo_split["best_iter"],61        "n_features": X.shape[1],62    }63