File size: 1,902 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Model training on GP-mined or qlib features."""
from __future__ import annotations
from pathlib import Path
import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
def prepare_ml_matrix(df: pd.DataFrame, feature_prefix: str = "factor_", target_col: str = "target_return"):
feature_cols = [c for c in df.columns if c.startswith(feature_prefix)]
work = df.copy()
for col in feature_cols:
work[col] = work.groupby("date")[col].transform(lambda x: x.rank(pct=True))
work["target_rank"] = work.groupby("date")[target_col].transform(lambda x: x.rank(pct=True))
work = work.replace([np.inf, -np.inf], np.nan).dropna()
return work, feature_cols
def train_lightgbm(
train_df: pd.DataFrame,
valid_df: pd.DataFrame,
feature_cols: list[str],
target_col: str = "target_rank",
model_path: Path | None = None,
) -> lgb.Booster:
train_set = lgb.Dataset(train_df[feature_cols], label=train_df[target_col])
valid_set = lgb.Dataset(valid_df[feature_cols], label=valid_df[target_col], reference=train_set)
params = {
"objective": "regression",
"metric": "mse",
"learning_rate": 0.05,
"num_leaves": 31,
"feature_fraction": 0.8,
"bagging_fraction": 0.8,
"bagging_freq": 5,
"verbose": -1,
"seed": 42,
}
model = lgb.train(
params,
train_set,
num_boost_round=500,
valid_sets=[valid_set],
callbacks=[lgb.early_stopping(50, verbose=False)],
)
if model_path is not None:
model_path.parent.mkdir(parents=True, exist_ok=True)
model.save_model(str(model_path))
valid_pred = model.predict(valid_df[feature_cols])
rmse = mean_squared_error(valid_df[target_col], valid_pred, squared=False)
print(f"LightGBM valid RMSE: {rmse:.6f}")
return model
|