Files
2026-07-13 13:26:28 +08:00

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