925 lines
31 KiB
Python
925 lines
31 KiB
Python
"""Shared GBM pipeline infrastructure for Ch12 notebooks and case study templates.
|
|
|
|
Provides:
|
|
- load_gbm_config(): Load canonical params from YAML config files
|
|
- make_model_params(): Transparent library-specific parameter mapping
|
|
- create_model(): Factory for unfitted sklearn-compatible GBM regressors
|
|
|
|
Usage:
|
|
from case_studies.utils.gbm import load_gbm_config, make_model_params, create_model
|
|
|
|
config = load_gbm_config("medium")
|
|
params = make_model_params(config, "lightgbm", "cpu")
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import time
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
# Import torch before ml4t.diagnostic. ml4t.diagnostic transitively loads the
|
|
# `cuda` Python package, which dlopens the older system `libcudart.so.12`
|
|
# (12.0.146) and wins the symbol resolution; subsequent torch imports then
|
|
# fail with `undefined symbol: cudaGetDriverEntryPointByVersion`. Loading
|
|
# torch first ensures its bundled CUDA runtime wins. Same pattern as in
|
|
# `case_studies/utils/latent_factors/__init__.py` and `model_analysis.py`.
|
|
import torch # noqa: F401
|
|
import yaml
|
|
from ml4t.diagnostic.metrics import cross_sectional_ic
|
|
|
|
from utils.config import REPO_ROOT
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Configuration
|
|
# ---------------------------------------------------------------------------
|
|
from utils.modeling import RANDOM_SEED, seed_everything
|
|
|
|
_CONFIGS_DIR = REPO_ROOT / "case_studies" / "_configs" / "gbm"
|
|
_MODELING_CONFIG = REPO_ROOT / "case_studies" / "_configs" / "modeling.yaml"
|
|
|
|
# Fallback presets (used when YAML files are not available)
|
|
PRESETS: dict[str, dict[str, Any]] = {
|
|
"light": {
|
|
"n_trees": 200,
|
|
"max_depth": 4,
|
|
"lr": 0.10,
|
|
"l2": 1.0,
|
|
"subsample": 0.8,
|
|
"colsample": 0.8,
|
|
"min_leaf": 20,
|
|
},
|
|
"medium": {
|
|
"n_trees": 500,
|
|
"max_depth": 6,
|
|
"lr": 0.05,
|
|
"l2": 1.0,
|
|
"subsample": 0.8,
|
|
"colsample": 0.8,
|
|
"min_leaf": 20,
|
|
},
|
|
"heavy": {
|
|
"n_trees": 1000,
|
|
"max_depth": 8,
|
|
"lr": 0.01,
|
|
"l2": 1.0,
|
|
"subsample": 0.8,
|
|
"colsample": 0.8,
|
|
"min_leaf": 20,
|
|
},
|
|
"default": {
|
|
"n_trees": 500,
|
|
},
|
|
}
|
|
|
|
|
|
def load_gbm_config(preset: str = "medium") -> dict[str, Any]:
|
|
"""Load canonical GBM parameters from YAML config.
|
|
|
|
Falls back to built-in PRESETS if YAML file not found.
|
|
|
|
Parameters
|
|
----------
|
|
preset : str
|
|
Preset name ("light", "medium", "heavy") or path to YAML file.
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
Canonical parameters (n_trees, max_depth, lr, l2, ...).
|
|
"""
|
|
yaml_path = _CONFIGS_DIR / f"{preset}.yaml"
|
|
if yaml_path.exists():
|
|
with open(yaml_path) as f:
|
|
return yaml.safe_load(f)
|
|
|
|
if preset in PRESETS:
|
|
return dict(PRESETS[preset])
|
|
|
|
raise ValueError(f"Unknown preset '{preset}'. Available: {list(PRESETS.keys())}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parameter Translation (transparent name mapping)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
PARAM_NAMES: dict[str, dict[str, str]] = {
|
|
"xgboost": {
|
|
"n_trees": "n_estimators",
|
|
"lr": "learning_rate",
|
|
"l1": "reg_alpha",
|
|
"l2": "reg_lambda",
|
|
"colsample": "colsample_bytree",
|
|
"min_leaf": "min_child_weight",
|
|
},
|
|
"lightgbm": {
|
|
"n_trees": "n_estimators",
|
|
"lr": "learning_rate",
|
|
"l1": "reg_alpha",
|
|
"l2": "reg_lambda",
|
|
"colsample": "colsample_bytree",
|
|
"min_leaf": "min_child_samples",
|
|
},
|
|
"catboost": {
|
|
"n_trees": "iterations",
|
|
"lr": "learning_rate",
|
|
"l1": "model_size_reg",
|
|
"l2": "l2_leaf_reg",
|
|
"max_depth": "depth",
|
|
"min_leaf": "min_data_in_leaf",
|
|
},
|
|
"sklearn_hgb": {
|
|
"n_trees": "max_iter",
|
|
"lr": "learning_rate",
|
|
"l2": "l2_regularization",
|
|
"min_leaf": "min_samples_leaf",
|
|
"colsample": "max_features",
|
|
},
|
|
}
|
|
|
|
# Canonical params that have no equivalent in certain libraries
|
|
_SKIP_PARAMS: dict[str, set[str]] = {
|
|
"sklearn_hgb": {"subsample"},
|
|
"catboost": {"colsample"},
|
|
}
|
|
|
|
# Cached GPU device per library: "cuda" or None (CPU only)
|
|
# OpenCL ("gpu") is NEVER used — it is slower and produces misleading benchmarks.
|
|
_BEST_GPU: dict[str, str | None] = {}
|
|
|
|
|
|
def _best_gpu_device(library: str) -> str | None:
|
|
"""Return "cuda" if library supports CUDA on this system, else None.
|
|
|
|
Only CUDA is accepted. OpenCL (device="gpu") is explicitly excluded —
|
|
it is orders of magnitude slower and produces misleading benchmark results.
|
|
"""
|
|
if library not in _BEST_GPU:
|
|
import numpy as _np
|
|
|
|
_X = _np.random.randn(10, 2).astype(_np.float32)
|
|
_y = _np.random.randn(10).astype(_np.float32)
|
|
|
|
try:
|
|
if library == "lightgbm":
|
|
import lightgbm as lgb
|
|
|
|
lgb.LGBMRegressor(n_estimators=2, device="cuda", verbose=-1).fit(_X, _y)
|
|
elif library == "xgboost":
|
|
import xgboost as xgb
|
|
|
|
xgb.XGBRegressor(
|
|
n_estimators=2, device="cuda", tree_method="hist", verbosity=0
|
|
).fit(_X, _y)
|
|
_BEST_GPU[library] = "cuda"
|
|
except Exception:
|
|
_BEST_GPU[library] = None
|
|
return _BEST_GPU[library]
|
|
|
|
|
|
# Library-specific defaults (not in canonical config)
|
|
_LIB_DEFAULTS: dict[str, dict[str, Any]] = {
|
|
"xgboost": {"tree_method": "hist", "random_state": RANDOM_SEED, "verbosity": 0, "n_jobs": -1},
|
|
"lightgbm": {"random_state": RANDOM_SEED, "verbose": -1, "n_jobs": -1},
|
|
"catboost": {
|
|
"bootstrap_type": "Bernoulli",
|
|
"random_seed": RANDOM_SEED,
|
|
"verbose": 0,
|
|
"allow_writing_files": False,
|
|
"thread_count": -1,
|
|
},
|
|
"sklearn_hgb": {"random_state": RANDOM_SEED},
|
|
}
|
|
|
|
# Canonical objective → library-specific mapping
|
|
_OBJECTIVE_MAP: dict[str, dict[str, str]] = {
|
|
"lightgbm": {
|
|
"mse": "regression",
|
|
"mae": "regression_l1",
|
|
"huber": "huber",
|
|
"binary": "binary",
|
|
"multiclass": "multiclass",
|
|
},
|
|
"xgboost": {
|
|
"mse": "reg:squarederror",
|
|
"mae": "reg:absoluteerror",
|
|
"huber": "reg:pseudohubererror",
|
|
"binary": "binary:logistic",
|
|
"multiclass": "multi:softprob",
|
|
},
|
|
"catboost": {
|
|
"mse": "RMSE",
|
|
"mae": "MAE",
|
|
"huber": "Huber",
|
|
"binary": "Logloss",
|
|
"multiclass": "MultiClass",
|
|
},
|
|
}
|
|
|
|
|
|
def make_model_params(
|
|
canonical: dict[str, Any],
|
|
library: str,
|
|
device: str = "cpu",
|
|
) -> dict[str, Any]:
|
|
"""Map canonical params to library-specific kwargs.
|
|
|
|
GPU overrides come from the config's ``gpu`` section (visible in YAML),
|
|
not from hidden internal logic.
|
|
|
|
Parameters
|
|
----------
|
|
canonical : dict
|
|
Canonical params (n_trees, max_depth, lr, ...).
|
|
library : str
|
|
One of "xgboost", "lightgbm", "catboost", "sklearn_hgb".
|
|
device : str
|
|
"cpu" or "gpu".
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
Library-specific kwargs ready for model constructor.
|
|
"""
|
|
if library not in PARAM_NAMES:
|
|
raise ValueError(f"Unknown library: {library}. Use xgboost/lightgbm/catboost/sklearn_hgb.")
|
|
|
|
name_map = PARAM_NAMES[library]
|
|
lib_params: dict[str, Any] = {}
|
|
|
|
# Map canonical names to library names
|
|
skip = _SKIP_PARAMS.get(library, set())
|
|
for k, v in canonical.items():
|
|
if k in ("gpu", "objective"):
|
|
continue # Handled separately
|
|
if k in skip:
|
|
continue # No equivalent in this library
|
|
lib_name = name_map.get(k, k)
|
|
lib_params[lib_name] = v
|
|
|
|
# LightGBM: num_leaves — explicit value wins over max_depth derivation
|
|
if library == "lightgbm":
|
|
if "num_leaves" in canonical:
|
|
lib_params["num_leaves"] = canonical["num_leaves"]
|
|
elif "max_depth" in canonical:
|
|
lib_params["num_leaves"] = 2 ** canonical["max_depth"] - 1
|
|
|
|
# Objective mapping (canonical → library-specific)
|
|
if "objective" in canonical and library in _OBJECTIVE_MAP:
|
|
lib_params["objective"] = _OBJECTIVE_MAP[library].get(
|
|
canonical["objective"], canonical["objective"]
|
|
)
|
|
|
|
# Library defaults
|
|
lib_params.update(_LIB_DEFAULTS.get(library, {}))
|
|
|
|
# GPU: device params + config overrides (visible in YAML)
|
|
# Accept both "gpu" and "cuda" — both mean "use CUDA" (OpenCL is never used)
|
|
if device in ("gpu", "cuda"):
|
|
if library in ("xgboost", "lightgbm"):
|
|
gpu_dev = _best_gpu_device(library)
|
|
if gpu_dev:
|
|
lib_params["device"] = gpu_dev
|
|
else:
|
|
raise RuntimeError(
|
|
f"{library} GPU requested but CUDA is not available. "
|
|
f"Run with device='cpu' or install {library} with CUDA support."
|
|
)
|
|
elif library == "catboost":
|
|
lib_params["task_type"] = "GPU"
|
|
lib_params["devices"] = "0"
|
|
|
|
gpu_overrides = canonical.get("gpu", {}).get(library, {})
|
|
lib_params.update(gpu_overrides)
|
|
|
|
return lib_params
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model Factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def create_model(
|
|
library: str,
|
|
params: dict[str, Any] | None = None,
|
|
device: str = "cpu",
|
|
gpu_adjustments: bool = True,
|
|
task_type: str = "regression",
|
|
):
|
|
"""Create an unfitted sklearn-compatible GBM model.
|
|
|
|
Parameters
|
|
----------
|
|
library : str
|
|
One of "xgboost", "lightgbm", "catboost", "sklearn_hgb".
|
|
params : dict, optional
|
|
Canonical params. Defaults to "medium" preset.
|
|
device : str
|
|
"cpu" or "gpu".
|
|
gpu_adjustments : bool
|
|
If True and device="gpu", applies GPU-specific params from config.
|
|
task_type : str
|
|
"regression" or "classification".
|
|
|
|
Returns
|
|
-------
|
|
Unfitted sklearn-compatible model (regressor or classifier).
|
|
"""
|
|
import catboost as cb
|
|
import lightgbm as lgb
|
|
import xgboost as xgb
|
|
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
|
|
|
|
if params is None:
|
|
params = load_gbm_config("medium")
|
|
|
|
effective_device = device if gpu_adjustments else "cpu"
|
|
lib_params = make_model_params(params, library, effective_device)
|
|
|
|
if task_type == "classification":
|
|
if library == "sklearn_hgb":
|
|
return HistGradientBoostingClassifier(**lib_params)
|
|
if library == "xgboost":
|
|
return xgb.XGBClassifier(**lib_params)
|
|
if library == "lightgbm":
|
|
return lgb.LGBMClassifier(**lib_params)
|
|
if library == "catboost":
|
|
return cb.CatBoostClassifier(**lib_params)
|
|
else:
|
|
if library == "sklearn_hgb":
|
|
return HistGradientBoostingRegressor(**lib_params)
|
|
if library == "xgboost":
|
|
return xgb.XGBRegressor(**lib_params)
|
|
if library == "lightgbm":
|
|
return lgb.LGBMRegressor(**lib_params)
|
|
if library == "catboost":
|
|
return cb.CatBoostRegressor(**lib_params)
|
|
|
|
raise ValueError(f"Unknown library: {library}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ---------------------------------------------------------------------------
|
|
# Checkpoint Prediction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _predict_at_checkpoint(model, X: np.ndarray, n_trees: int, library: str) -> np.ndarray:
|
|
"""Predict using only the first `n_trees` trees from a trained model.
|
|
|
|
Supports partial-iteration prediction for LightGBM, XGBoost, and CatBoost.
|
|
For sklearn HistGradientBoosting, returns full prediction (no checkpoint support).
|
|
"""
|
|
if library == "lightgbm":
|
|
return model.predict(X, num_iteration=n_trees)
|
|
elif library == "xgboost":
|
|
return model.predict(X, iteration_range=(0, n_trees))
|
|
elif library == "catboost":
|
|
return model.predict(X, ntree_end=n_trees)
|
|
else: # sklearn_hgb — no checkpoint support
|
|
return model.predict(X)
|
|
|
|
|
|
def _extract_feature_importance(
|
|
model, feature_names: list[str], library: str, top_n: int = 10
|
|
) -> list[tuple[str, float]]:
|
|
"""Extract top-N feature importances from a fitted model."""
|
|
try:
|
|
importances = model.feature_importances_
|
|
if importances is None or len(importances) == 0:
|
|
return []
|
|
pairs = sorted(
|
|
zip(feature_names, importances, strict=False), key=lambda x: abs(x[1]), reverse=True
|
|
)
|
|
return pairs[:top_n]
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config-driven GBM training (public API for notebooks)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def prepare_gbm_folds(
|
|
dataset_pd,
|
|
splits: list[dict[str, Any]],
|
|
feature_names: list[str],
|
|
label_col: str,
|
|
date_col: str,
|
|
entity_col: str = "symbol",
|
|
task_type: str = "regression",
|
|
class_values: list | None = None,
|
|
temporal_by_fold=None,
|
|
temporal_keys: list[str] | None = None,
|
|
temporal_feature_names: list[str] | None = None,
|
|
train_sample_frac: float = 1.0,
|
|
) -> list[dict[str, Any]]:
|
|
"""Prepare CV fold data for GBM training.
|
|
|
|
Unlike linear folds, GBM folds:
|
|
- Use float32 (LightGBM native precision)
|
|
- No imputation or scaling (GBM handles NaN natively)
|
|
- Include remapped labels for classification (0-indexed for LightGBM)
|
|
|
|
Parameters
|
|
----------
|
|
dataset_pd : pandas DataFrame
|
|
Full dataset.
|
|
splits : list[dict]
|
|
Walk-forward splits.
|
|
feature_names : list[str]
|
|
Feature column names.
|
|
label_col, date_col, entity_col : str
|
|
Column names.
|
|
task_type : str
|
|
"regression" or "classification".
|
|
class_values : list, optional
|
|
Sorted unique class values for classification.
|
|
temporal_by_fold : pd.DataFrame, optional
|
|
Per-fold temporal features with a 'fold' column.
|
|
temporal_keys : list[str], optional
|
|
Join keys for temporal features.
|
|
temporal_feature_names : list[str], optional
|
|
Temporal feature column names to replace per fold.
|
|
train_sample_frac : float, optional
|
|
Fraction of training rows to keep per fold (1.0 = keep all).
|
|
Walk-forward CV structure is preserved (date ranges unchanged);
|
|
only the within-fold row density is reduced. Validation set is
|
|
NEVER sampled — OOS IC is always computed on the full val slice.
|
|
Seed is tied to fold_id for reproducibility. Use < 1.0 for
|
|
memory/compute-constrained runs on large datasets (e.g.,
|
|
nasdaq100 minute bars). Default 1.0.
|
|
|
|
Returns
|
|
-------
|
|
list[dict]
|
|
Each dict has: fold, X_train, y_train, y_train_lgb, X_val, y_val,
|
|
y_val_lgb, dates, entities, n_train, n_val.
|
|
"""
|
|
from utils.modeling import _replace_temporal_columns
|
|
|
|
dates_series = dataset_pd[date_col]
|
|
entity_series = dataset_pd.get(entity_col)
|
|
is_classification = task_type == "classification" and class_values
|
|
has_fold_temporal = temporal_by_fold is not None and temporal_keys and temporal_feature_names
|
|
|
|
folds = []
|
|
for split in splits:
|
|
fold_id = split["fold"]
|
|
train_mask = (dates_series >= split["train_start"]) & (dates_series <= split["train_end"])
|
|
val_start = split.get("val_start", split.get("test_start"))
|
|
val_end = split.get("val_end", split.get("test_end"))
|
|
val_mask = (dates_series >= val_start) & (dates_series <= val_end)
|
|
|
|
if has_fold_temporal:
|
|
train_rows = _replace_temporal_columns(
|
|
dataset_pd,
|
|
train_mask,
|
|
temporal_by_fold,
|
|
temporal_keys,
|
|
temporal_feature_names,
|
|
fold_id,
|
|
)
|
|
val_rows = _replace_temporal_columns(
|
|
dataset_pd,
|
|
val_mask,
|
|
temporal_by_fold,
|
|
temporal_keys,
|
|
temporal_feature_names,
|
|
fold_id,
|
|
)
|
|
X_train = train_rows[feature_names].values.astype(np.float32)
|
|
y_train = train_rows[label_col].values.astype(np.float32)
|
|
X_val = val_rows[feature_names].values.astype(np.float32)
|
|
y_val = val_rows[label_col].values.astype(np.float32)
|
|
val_dates = val_rows[date_col].values
|
|
del train_rows, val_rows
|
|
else:
|
|
X_train = dataset_pd.loc[train_mask, feature_names].values.astype(np.float32)
|
|
y_train = dataset_pd.loc[train_mask, label_col].values.astype(np.float32)
|
|
X_val = dataset_pd.loc[val_mask, feature_names].values.astype(np.float32)
|
|
y_val = dataset_pd.loc[val_mask, label_col].values.astype(np.float32)
|
|
val_dates = dataset_pd.loc[val_mask, date_col].values
|
|
|
|
# Drop NaN labels
|
|
tv = ~np.isnan(y_train)
|
|
vv = ~np.isnan(y_val)
|
|
X_train, y_train = X_train[tv], y_train[tv]
|
|
X_val, y_val = X_val[vv], y_val[vv]
|
|
val_dates = val_dates[vv]
|
|
val_entities = (
|
|
dataset_pd.loc[val_mask, entity_col].values[vv] if entity_series is not None else None
|
|
)
|
|
|
|
# Optional train subsample (never touch val — OOS IC uses full val slice).
|
|
# Seed is tied to fold_id for reproducibility.
|
|
if 0.0 < train_sample_frac < 1.0 and len(X_train) > 0:
|
|
n_keep = max(1, int(len(X_train) * train_sample_frac))
|
|
rng = np.random.default_rng(RANDOM_SEED + fold_id)
|
|
keep_idx = rng.choice(len(X_train), size=n_keep, replace=False)
|
|
keep_idx.sort() # preserve row order
|
|
X_train = X_train[keep_idx]
|
|
y_train = y_train[keep_idx]
|
|
|
|
# Classification: remap labels to 0-indexed for LightGBM
|
|
if is_classification:
|
|
y_train_lgb, _ = _remap_labels_for_lgb(y_train.astype(int), class_values)
|
|
y_val_lgb, _ = _remap_labels_for_lgb(y_val.astype(int), class_values)
|
|
else:
|
|
y_train_lgb = y_train
|
|
y_val_lgb = y_val
|
|
|
|
folds.append(
|
|
{
|
|
"fold": split["fold"],
|
|
"X_train": X_train,
|
|
"y_train": y_train,
|
|
"y_train_lgb": y_train_lgb,
|
|
"X_val": X_val,
|
|
"y_val": y_val,
|
|
"y_val_lgb": y_val_lgb,
|
|
"dates": val_dates,
|
|
"entities": val_entities,
|
|
"n_train": len(X_train),
|
|
"n_val": len(X_val),
|
|
}
|
|
)
|
|
|
|
return folds
|
|
|
|
|
|
def train_gbm_config(
|
|
config: dict[str, Any],
|
|
fold_data: list[dict[str, Any]],
|
|
*,
|
|
feature_names: list[str],
|
|
device: str = "cuda",
|
|
max_bin: int | None = None,
|
|
entity_col: str = "symbol",
|
|
date_col: str = "timestamp",
|
|
task_type: str = "regression",
|
|
class_values: list | None = None,
|
|
save_dir: Path | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Train a single GBM config across all CV folds.
|
|
|
|
Trains to max_iterations, evaluates cross-sectional IC at checkpoints,
|
|
and returns the best checkpoint along with predictions and learning curves.
|
|
|
|
Parameters
|
|
----------
|
|
config : dict
|
|
Preset dict with config_name, params, max_iterations, checkpoint_interval.
|
|
fold_data : list[dict]
|
|
From prepare_gbm_folds().
|
|
feature_names : list[str]
|
|
For feature importance extraction.
|
|
device : str
|
|
"cpu" or "cuda"/"gpu".
|
|
max_bin : int, optional
|
|
Override max_bin (GPU typically needs 63).
|
|
entity_col, date_col : str
|
|
For IC computation.
|
|
task_type : str
|
|
"regression" or "classification".
|
|
class_values : list, optional
|
|
For classification score extraction.
|
|
save_dir : Path, optional
|
|
Save booster files here.
|
|
|
|
Returns
|
|
-------
|
|
dict with keys:
|
|
config_name, best_iter, best_ic, elapsed_s, fold_ics (dict[int, list]),
|
|
learning_curves (list[dict]), predictions (list[dict]), top_features.
|
|
"""
|
|
import lightgbm as lgb
|
|
|
|
config_name = config["config_name"]
|
|
num_boost_round = config.get("max_iterations", 500)
|
|
checkpoint_interval = config.get("checkpoint_interval", 50)
|
|
is_classification = task_type == "classification" and class_values
|
|
|
|
# Build LightGBM params from preset
|
|
params = dict(config["params"])
|
|
params["metric"] = "None"
|
|
params["verbosity"] = params.get("verbosity", -1)
|
|
|
|
# Device setup
|
|
if device in ("cuda", "gpu"):
|
|
gpu_dev = _best_gpu_device("lightgbm")
|
|
if gpu_dev:
|
|
params["device"] = gpu_dev
|
|
if max_bin is not None:
|
|
params["max_bin"] = max_bin
|
|
|
|
# Classification: ensure num_class for multiclass
|
|
if is_classification and class_values and len(class_values) > 2:
|
|
params["num_class"] = len(class_values)
|
|
|
|
checkpoints = list(range(checkpoint_interval, num_boost_round + 1, checkpoint_interval))
|
|
if not checkpoints or checkpoints[-1] != num_boost_round:
|
|
checkpoints.append(num_boost_round)
|
|
|
|
t0 = time.perf_counter()
|
|
checkpoint_ics: dict[int, list[float]] = {cp: [] for cp in checkpoints}
|
|
all_preds: list[dict] = []
|
|
top_features: list[tuple[str, float]] = []
|
|
booster_dir = save_dir / "boosters" if save_dir else None
|
|
if booster_dir:
|
|
booster_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
for fd in fold_data:
|
|
if fd["n_train"] == 0 or fd["n_val"] == 0:
|
|
continue
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
dtrain = lgb.Dataset(
|
|
fd["X_train"],
|
|
label=fd["y_train_lgb"],
|
|
feature_name=feature_names,
|
|
free_raw_data=False,
|
|
)
|
|
# Print progress every 50 iterations so long runs aren't silent.
|
|
# Also print per-fold heartbeat so we see which fold is active
|
|
# on large datasets.
|
|
print(
|
|
f" fold {fd['fold']}: training "
|
|
f"n_train={fd['n_train']:,} n_val={fd['n_val']:,} "
|
|
f"trees={num_boost_round} num_leaves={params.get('num_leaves', '?')} "
|
|
f"obj={params.get('objective', '?')}",
|
|
flush=True,
|
|
)
|
|
_fold_t0 = time.perf_counter()
|
|
model = lgb.train(
|
|
params,
|
|
dtrain,
|
|
num_boost_round=num_boost_round,
|
|
callbacks=[lgb.log_evaluation(period=50)],
|
|
)
|
|
print(
|
|
f" fold {fd['fold']}: done in {time.perf_counter() - _fold_t0:.0f}s",
|
|
flush=True,
|
|
)
|
|
|
|
if booster_dir:
|
|
model.save_model(str(booster_dir / f"fold_{fd['fold']}.txt"))
|
|
|
|
# Feature importance (first fold only)
|
|
if not top_features:
|
|
imp = model.feature_importance(importance_type="gain")
|
|
pairs = sorted(zip(feature_names, imp, strict=False), key=lambda x: x[1], reverse=True)
|
|
top_features = pairs[:10]
|
|
|
|
# Predict at all checkpoints
|
|
for cp in checkpoints:
|
|
raw_preds = model.predict(fd["X_val"], num_iteration=cp)
|
|
if is_classification:
|
|
preds = _extract_gbm_score(raw_preds, class_values, len(fd["X_val"]))
|
|
else:
|
|
preds = raw_preds
|
|
ic_frame = pl.DataFrame(
|
|
{
|
|
"date": fd["dates"],
|
|
"symbol": fd["entities"],
|
|
"y_true": fd["y_val"],
|
|
"y_pred": preds,
|
|
}
|
|
)
|
|
ic = cross_sectional_ic(
|
|
ic_frame,
|
|
ic_frame,
|
|
pred_col="y_pred",
|
|
ret_col="y_true",
|
|
date_col="date",
|
|
entity_col="symbol",
|
|
min_obs=5,
|
|
)["ic_mean"]
|
|
checkpoint_ics[cp].append(ic)
|
|
all_preds.append(
|
|
{
|
|
"dates": fd["dates"],
|
|
"entities": fd["entities"],
|
|
"y_true": fd["y_val"],
|
|
"y_pred": preds,
|
|
"fold": fd["fold"],
|
|
"n_trees": cp,
|
|
}
|
|
)
|
|
|
|
del dtrain, model
|
|
|
|
# Best checkpoint by mean IC
|
|
best_cp = max(
|
|
checkpoints, key=lambda cp: np.mean(checkpoint_ics[cp]) if checkpoint_ics[cp] else -1
|
|
)
|
|
best_ic = float(np.mean(checkpoint_ics[best_cp])) if checkpoint_ics[best_cp] else 0.0
|
|
best_ic_std = float(np.std(checkpoint_ics[best_cp])) if checkpoint_ics[best_cp] else 0.0
|
|
elapsed = time.perf_counter() - t0
|
|
|
|
# Learning curves
|
|
curves = [
|
|
{
|
|
"config": config_name,
|
|
"iteration": cp,
|
|
"ic_mean": float(np.mean(checkpoint_ics[cp])) if checkpoint_ics[cp] else 0.0,
|
|
"ic_std": float(np.std(checkpoint_ics[cp])) if checkpoint_ics[cp] else 0.0,
|
|
}
|
|
for cp in checkpoints
|
|
]
|
|
|
|
# Per-fold metrics at best checkpoint
|
|
def _fold_ic(e: dict[str, Any]) -> float:
|
|
frame = pl.DataFrame(
|
|
{
|
|
"date": e["dates"],
|
|
"symbol": e["entities"],
|
|
"y_true": e["y_true"],
|
|
"y_pred": e["y_pred"],
|
|
}
|
|
)
|
|
return cross_sectional_ic(
|
|
frame,
|
|
frame,
|
|
pred_col="y_pred",
|
|
ret_col="y_true",
|
|
date_col="date",
|
|
entity_col="symbol",
|
|
min_obs=5,
|
|
)["ic_mean"]
|
|
|
|
fold_metrics = [
|
|
{
|
|
"fold_id": e["fold"],
|
|
"ic_mean": _fold_ic(e),
|
|
"n_train": [fd for fd in fold_data if fd["fold"] == e["fold"]][0]["n_train"],
|
|
"n_test": len(e["y_true"]),
|
|
}
|
|
for e in all_preds
|
|
if e["n_trees"] == best_cp
|
|
]
|
|
|
|
gc.collect()
|
|
|
|
return {
|
|
"config_name": config_name,
|
|
"best_iter": best_cp,
|
|
"best_ic": best_ic,
|
|
"best_ic_std": best_ic_std,
|
|
"elapsed_s": elapsed,
|
|
"checkpoint_ics": checkpoint_ics,
|
|
"learning_curves": curves,
|
|
"predictions": all_preds,
|
|
"fold_metrics": fold_metrics,
|
|
"top_features": top_features,
|
|
}
|
|
|
|
|
|
def _make_lgb_native_params(canonical: dict[str, Any], device: str) -> dict[str, Any]:
|
|
"""Convert canonical config to native lgb.train() params dict.
|
|
|
|
Strips sklearn-only keys (n_estimators) and disables built-in metrics.
|
|
"""
|
|
params = make_model_params(canonical, "lightgbm", device)
|
|
params.pop("n_estimators", None)
|
|
params.pop("n_jobs", None)
|
|
params["metric"] = "None"
|
|
params["seed"] = params.pop("random_state", RANDOM_SEED)
|
|
# Subsampling requires bagging_freq in native API
|
|
if params.get("subsample", 1.0) < 1.0:
|
|
params["bagging_freq"] = 1
|
|
return params
|
|
|
|
|
|
def _remap_labels_for_lgb(y: np.ndarray, class_values: list) -> tuple[np.ndarray, dict]:
|
|
"""Remap class labels to 0-indexed for LightGBM native API.
|
|
|
|
E.g., {-1, 0, 1} -> {0, 1, 2}. Returns (remapped_y, mapping_dict).
|
|
"""
|
|
sorted_vals = sorted(class_values)
|
|
mapping = {v: i for i, v in enumerate(sorted_vals)}
|
|
remapped = np.array([mapping[v] for v in y], dtype=np.int32)
|
|
return remapped, mapping
|
|
|
|
|
|
def _extract_gbm_score(raw_preds: np.ndarray, class_values: list, n_samples: int) -> np.ndarray:
|
|
"""Extract continuous score from GBM classification output for IC computation.
|
|
|
|
Binary: raw_preds is P(class=1) directly.
|
|
Multiclass: raw_preds shape = (n_samples, n_classes) -> expected value.
|
|
"""
|
|
sorted_vals = sorted(class_values)
|
|
if len(sorted_vals) == 2:
|
|
# Binary: LightGBM native returns P(class=1) directly
|
|
return raw_preds.ravel()
|
|
# Multiclass: raw_preds shape = (n_samples, n_classes)
|
|
proba = raw_preds.reshape(n_samples, len(sorted_vals))
|
|
return proba @ np.array(sorted_vals, dtype=np.float64)
|
|
|
|
|
|
def register_gbm_result(
|
|
case_study_id: str,
|
|
result: dict,
|
|
cfg: dict,
|
|
label_col: str,
|
|
n_folds: int,
|
|
*,
|
|
max_bin: int | None = None,
|
|
entry_point: str = "07_gbm",
|
|
date_col: str = "timestamp",
|
|
entity_col: str = "symbol",
|
|
train_sample_frac: float = 1.0,
|
|
prediction_split: str = "validation",
|
|
) -> str:
|
|
"""Register a single GBM config's result to the registry.
|
|
|
|
Called INSIDE the training loop (per-config) so each config is persisted
|
|
immediately after it trains. This protects against interruption losing
|
|
all completed configs — a failure rule enforced by the memory file
|
|
``feedback_incremental_save_violation.md``.
|
|
|
|
Writes training_run, prediction_set (best-iter predictions),
|
|
learning_curves.parquet, and fold_metrics.parquet.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
The training_hash for the registered run.
|
|
"""
|
|
import polars as pl
|
|
|
|
from case_studies.utils.registry import (
|
|
build_training_spec,
|
|
get_training_dir,
|
|
register_prediction_set,
|
|
register_training_run,
|
|
)
|
|
|
|
spec = build_training_spec(
|
|
cfg["family"],
|
|
cfg["config_name"],
|
|
label_col,
|
|
n_folds=n_folds,
|
|
max_bin=max_bin,
|
|
checkpoint_interval=cfg.get("checkpoint_interval", 50),
|
|
train_sample_frac=train_sample_frac,
|
|
)
|
|
t_hash = register_training_run(
|
|
case_study_id,
|
|
spec=spec,
|
|
entry_point=entry_point,
|
|
elapsed_s=result.get("elapsed_s"),
|
|
)
|
|
|
|
# Best-checkpoint predictions as a DataFrame
|
|
best_preds = [e for e in result["predictions"] if e["n_trees"] == result["best_iter"]]
|
|
if best_preds:
|
|
pred_rows = []
|
|
for e in best_preds:
|
|
n = len(e["y_pred"])
|
|
pred_rows.append(
|
|
pl.DataFrame(
|
|
{
|
|
date_col: e["dates"],
|
|
entity_col: e["entities"] if e["entities"] is not None else ["unknown"] * n,
|
|
"fold": [e["fold"]] * n,
|
|
"prediction": e["y_pred"],
|
|
"actual": e["y_true"],
|
|
}
|
|
)
|
|
)
|
|
pred_df = pl.concat(pred_rows).to_pandas()
|
|
register_prediction_set(
|
|
case_study_id,
|
|
t_hash,
|
|
split=prediction_split,
|
|
predictions=pred_df,
|
|
metrics={
|
|
"ic_mean": result["best_ic"],
|
|
"ic_std": result["best_ic_std"],
|
|
},
|
|
)
|
|
|
|
# Save learning curves and fold metrics to registry training dir
|
|
reg_dir = get_training_dir(case_study_id, spec)
|
|
cfg_curves = list(result.get("learning_curves", []))
|
|
if cfg_curves:
|
|
pl.DataFrame(cfg_curves).write_parquet(reg_dir / "learning_curves.parquet")
|
|
|
|
cfg_fold_metrics = result.get("fold_metrics", [])
|
|
if cfg_fold_metrics:
|
|
fm_df = pl.DataFrame(cfg_fold_metrics)
|
|
if "config_name" not in fm_df.columns:
|
|
fm_df = fm_df.with_columns(pl.lit(result["config_name"]).alias("config_name"))
|
|
fm_df.write_parquet(reg_dir / "fold_metrics.parquet")
|
|
|
|
return t_hash
|