| """ |
| XGBoost training script for the surf condition model (Section 3.3.5). |
| |
| Workflow: |
| 1. Load train / val splits from ml/data/processed/. |
| 2. Apply synthetic labels (ml/labels.py). |
| 3. Extract features (app/ml/feature_extractor.py). |
| 4. Fit a SimpleImputer on training data. |
| 5. Grid-search XGBoost with TimeSeriesSplit (5-fold). |
| 6. Evaluate best model on the val set. |
| 7. Persist model + imputer + metadata to ml/models/. |
| |
| Usage: |
| python -m ml.train |
| python -m ml.train --no-search # skip grid search, use default params |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| import joblib |
| import numpy as np |
| import pandas as pd |
| from sklearn.ensemble import HistGradientBoostingRegressor |
| from sklearn.impute import SimpleImputer |
| from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error |
| from sklearn.model_selection import GridSearchCV, TimeSeriesSplit |
|
|
| from app.ml.feature_extractor import FEATURE_NAMES, ForecastPointFeatureExtractor |
| from ml.labels import compute_synthetic_score |
| from ml.splits import split |
|
|
| MODELS_DIR = Path("ml/models") |
| PROCESSED = Path("ml/data/processed") |
|
|
| PARAM_GRID = { |
| "max_iter": [100, 300, 500], |
| "max_depth": [3, 5, 7], |
| "learning_rate": [0.01, 0.05, 0.1], |
| "min_samples_leaf": [10, 20, 50], |
| } |
|
|
| DEFAULT_PARAMS = { |
| "max_iter": 300, |
| "max_depth": 5, |
| "learning_rate": 0.05, |
| "min_samples_leaf": 20, |
| "l2_regularization": 0.1, |
| "random_state": 42, |
| } |
|
|
|
|
| def _prepare(df: pd.DataFrame, extractor: ForecastPointFeatureExtractor) -> np.ndarray: |
| """Extract features and add synthetic labels.""" |
| X = extractor.transform_batch(df, skill_level="intermediate") |
| return X |
|
|
|
|
| def _add_labels(df: pd.DataFrame) -> pd.Series: |
| return df.apply(compute_synthetic_score, axis=1) |
|
|
|
|
| def _dataset_hash(df: pd.DataFrame) -> str: |
| return hashlib.md5(pd.util.hash_pandas_object(df, index=False).values).hexdigest()[:12] |
|
|
|
|
| def train(grid_search: bool = True) -> None: |
| MODELS_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| print("Loading splits …") |
| train_df, val_df, _ = split() |
| print(f" train={len(train_df):,} val={len(val_df):,}") |
|
|
| extractor = ForecastPointFeatureExtractor() |
|
|
| print("Extracting features …") |
| X_train = _prepare(train_df, extractor) |
| X_val = _prepare(val_df, extractor) |
|
|
| print("Computing synthetic labels …") |
| y_train = _add_labels(train_df).values |
| y_val = _add_labels(val_df).values |
|
|
| print("Fitting imputer on training data …") |
| imputer = SimpleImputer(strategy="median") |
| X_train = imputer.fit_transform(X_train) |
| X_val = imputer.transform(X_val) |
|
|
| if grid_search: |
| print("Grid-searching hyperparameters (TimeSeriesSplit 5-fold) …") |
| base = HistGradientBoostingRegressor(random_state=42) |
| tscv = TimeSeriesSplit(n_splits=5) |
| gs = GridSearchCV( |
| base, PARAM_GRID, cv=tscv, |
| scoring="r2", n_jobs=-1, verbose=1, |
| refit=True, |
| ) |
| gs.fit(X_train, y_train) |
| model = gs.best_estimator_ |
| best_params = gs.best_params_ |
| cv_r2 = float(gs.best_score_) |
| print(f" Best CV R²: {cv_r2:.4f}") |
| print(f" Best params: {best_params}") |
| else: |
| print("Training with default parameters …") |
| model = HistGradientBoostingRegressor(**DEFAULT_PARAMS) |
| model.fit(X_train, y_train) |
| best_params = DEFAULT_PARAMS |
| cv_r2 = float("nan") |
|
|
| print("Evaluating on validation set …") |
| y_pred = model.predict(X_val) |
| val_r2 = r2_score(y_val, y_pred) |
| val_mae = mean_absolute_error(y_val, y_pred) |
| val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) |
| print(f" Val R²={val_r2:.4f} MAE={val_mae:.2f} RMSE={val_rmse:.2f}") |
|
|
| model_path = MODELS_DIR / "surf_condition_model.joblib" |
| imputer_path = MODELS_DIR / "imputer.joblib" |
| meta_path = MODELS_DIR / "model_metadata.json" |
|
|
| joblib.dump(model, model_path) |
| joblib.dump(imputer, imputer_path) |
|
|
| metadata = { |
| "training_timestamp": datetime.now(timezone.utc).isoformat(), |
| "dataset_hash": _dataset_hash(train_df), |
| "feature_list": FEATURE_NAMES, |
| "best_params": best_params, |
| "cv_r2_mean": cv_r2, |
| "val_r2": val_r2, |
| "val_mae": val_mae, |
| "val_rmse": val_rmse, |
| "train_rows": len(train_df), |
| "val_rows": len(val_df), |
| } |
| with open(meta_path, "w") as f: |
| json.dump(metadata, f, indent=2) |
|
|
| print(f"\nSaved:") |
| print(f" {model_path} ({model_path.stat().st_size // 1024} KB)") |
| print(f" {imputer_path}") |
| print(f" {meta_path}") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--no-search", action="store_true", |
| help="Skip grid search and use default hyperparameters") |
| args = parser.parse_args() |
| train(grid_search=not args.no_search) |
|
|
