Files
NexQuant/test/local/test_ml_trainer.py
T
TPTBusiness cbe1c52e00 refactor: rename project from Predix to NexQuant
Rename all source files, scripts, tests, documentation, and configuration
from Predix/predix to NexQuant/nexquant across the entire codebase.
2026-05-09 17:48:22 +02:00

786 lines
28 KiB
Python

"""
Tests for MLTrainer - ML Training Pipeline for NexQuant quant trading system.
Tests cover:
- Feature matrix building
- Train/validation split (time-series)
- LightGBM training (mocked)
- Feature importance extraction
- Model saving/loading
- Feedback generation
- Edge cases and error handling
Run: pytest test/local/test_ml_trainer.py -v
"""
import json
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch
import numpy as np
import pandas as pd
import pytest
from rdagent.scenarios.qlib.local.ml_trainer import MLTrainer, get_ml_trainer
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def tmp_models_dir(tmp_path):
"""Temporary directory for model output."""
return str(tmp_path / "models")
@pytest.fixture
def trainer(tmp_models_dir):
"""MLTrainer instance with temp directory."""
return MLTrainer(models_dir=tmp_models_dir, random_state=42)
@pytest.fixture
def sample_factors():
"""List of sample factor info dicts."""
return [
{
"factor_name": "momentum_5min",
"ic": 0.12,
"sharpe_ratio": 1.5,
"status": "success",
"workspace_hash": "abc123",
},
{
"factor_name": "mean_reversion_zscore",
"ic": 0.08,
"sharpe_ratio": 1.2,
"status": "success",
"workspace_hash": "def456",
},
{
"factor_name": "volatility_atr",
"ic": -0.05,
"sharpe_ratio": 0.8,
"status": "success",
"workspace_hash": "ghi789",
},
{
"factor_name": "low_quality_factor",
"ic": 0.005,
"status": "success",
"workspace_hash": "jkl012",
},
{
"factor_name": "failed_factor",
"status": "failed",
},
]
@pytest.fixture
def sample_feature_matrix():
"""Sample X, y data for training tests."""
np.random.seed(42)
n_samples = 1000
n_features = 5
X = pd.DataFrame(
np.random.randn(n_samples, n_features),
columns=[
"momentum_5min",
"momentum_15min",
"zscore_20",
"atr_ratio",
"volume_spike",
],
)
# Create y with some signal from X
y = pd.Series(
0.5 * X["momentum_5min"]
+ 0.3 * X["zscore_20"]
+ 0.2 * X["atr_ratio"]
+ np.random.randn(n_samples) * 0.1
)
return X, y
@pytest.fixture
def mock_model_info(sample_feature_matrix):
"""Mock model info for testing without actual training."""
X, y = sample_feature_matrix
return {
"model": MagicMock(),
"feature_names": list(X.columns),
"feature_importance": {
"momentum_5min": {"gain": 500.0, "split": 30, "gain_normalized": 0.50},
"momentum_15min": {"gain": 200.0, "split": 15, "gain_normalized": 0.20},
"zscore_20": {"gain": 150.0, "split": 12, "gain_normalized": 0.15},
"atr_ratio": {"gain": 100.0, "split": 8, "gain_normalized": 0.10},
"volume_spike": {"gain": 50.0, "split": 5, "gain_normalized": 0.05},
},
"feature_ranking": [
"momentum_5min",
"momentum_15min",
"zscore_20",
"atr_ratio",
"volume_spike",
],
"ic_train": 0.15,
"ic_valid": 0.10,
"rank_ic_train": 0.12,
"rank_ic_valid": 0.08,
"sharpe_valid": 0.05,
"mse_train": 0.01,
"mse_valid": 0.015,
"n_features": 5,
"n_samples_train": 800,
"n_samples_valid": 200,
"trained_at": "2026-04-09T12:00:00",
"params": {
"n_estimators": 500,
"max_depth": 8,
"learning_rate": 0.03,
"subsample": 0.8,
"colsample_bytree": 0.8,
"random_state": 42,
},
}
# ---------------------------------------------------------------------------
# Test: MLTrainer Initialization
# ---------------------------------------------------------------------------
class TestMLTrainerInit:
"""Test MLTrainer initialization."""
def test_default_init(self):
"""Test default initialization."""
trainer = MLTrainer()
assert trainer.random_state == 42
assert "results" in str(trainer.models_dir)
assert "models" in str(trainer.models_dir)
def test_custom_models_dir(self, tmp_path):
"""Test custom models directory."""
custom_dir = str(tmp_path / "custom_models")
trainer = MLTrainer(models_dir=custom_dir, random_state=123)
assert str(trainer.models_dir) == custom_dir
assert trainer.random_state == 123
def test_models_dir_created(self, tmp_path):
"""Test models directory is created if missing."""
new_dir = str(tmp_path / "nested" / "models")
trainer = MLTrainer(models_dir=new_dir)
assert Path(new_dir).exists()
# ---------------------------------------------------------------------------
# Test: load_top_factors
# ---------------------------------------------------------------------------
class TestLoadTopFactors:
"""Test loading top factors from JSON files."""
def test_load_factors_from_dir(self, trainer, tmp_path, sample_factors):
"""Test loading factors from directory."""
# Create factor JSON files
factors_dir = tmp_path / "factors"
factors_dir.mkdir()
for i, factor in enumerate(sample_factors[:3]):
(factors_dir / f"factor_{i}.json").write_text(
json.dumps(factor), encoding="utf-8"
)
result = trainer.load_top_factors(
top_n=2, min_ic=0.01, factors_dir=str(factors_dir)
)
assert len(result) == 2
# Sorted by |IC| descending
assert abs(result[0]["ic"]) >= abs(result[1]["ic"])
assert result[0]["factor_name"] == "momentum_5min"
def test_filters_by_status(self, trainer, tmp_path, sample_factors):
"""Test filtering out failed factors."""
factors_dir = tmp_path / "factors"
factors_dir.mkdir()
# Write one success, one failed
(factors_dir / "success.json").write_text(
json.dumps(sample_factors[0]), encoding="utf-8"
)
(factors_dir / "failed.json").write_text(
json.dumps(sample_factors[4]), encoding="utf-8"
)
result = trainer.load_top_factors(factors_dir=str(factors_dir))
assert len(result) == 1
assert result[0]["status"] == "success"
def test_filters_by_min_ic(self, trainer, tmp_path, sample_factors):
"""Test filtering by minimum IC threshold."""
factors_dir = tmp_path / "factors"
factors_dir.mkdir()
# Write factors with different IC
(factors_dir / "high_ic.json").write_text(
json.dumps(sample_factors[0]), encoding="utf-8"
)
(factors_dir / "low_ic.json").write_text(
json.dumps(sample_factors[3]), encoding="utf-8"
)
result = trainer.load_top_factors(min_ic=0.05, factors_dir=str(factors_dir))
assert len(result) == 1
assert result[0]["factor_name"] == "momentum_5min"
def test_empty_directory(self, trainer, tmp_path):
"""Test loading from empty directory."""
factors_dir = tmp_path / "empty"
factors_dir.mkdir()
result = trainer.load_top_factors(factors_dir=str(factors_dir))
assert result == []
def test_nonexistent_directory(self, trainer):
"""Test loading from non-existent directory."""
result = trainer.load_top_factors(factors_dir="/nonexistent/path")
assert result == []
def test_top_n_limit(self, trainer, tmp_path, sample_factors):
"""Test top_n limit is respected."""
factors_dir = tmp_path / "factors"
factors_dir.mkdir()
for i in range(10):
factor = {
"factor_name": f"factor_{i}",
"ic": 0.1 - i * 0.005,
"status": "success",
"workspace_hash": f"hash_{i}",
}
(factors_dir / f"factor_{i}.json").write_text(
json.dumps(factor), encoding="utf-8"
)
result = trainer.load_top_factors(top_n=3, min_ic=0.01, factors_dir=str(factors_dir))
assert len(result) == 3
# ---------------------------------------------------------------------------
# Test: build_feature_matrix
# ---------------------------------------------------------------------------
class TestBuildFeatureMatrix:
"""Test feature matrix building."""
def test_data_file_not_found(self, trainer, sample_factors):
"""Test when data file does not exist."""
X, y = trainer.build_feature_matrix(
sample_factors, data_file="/nonexistent/data.h5"
)
assert X is None
assert y is None
def test_no_workspace_hash(self, trainer, tmp_path, sample_factors):
"""Test handling of factors without workspace_hash."""
factors_no_hash = [{"factor_name": "test", "ic": 0.1, "status": "success"}]
X, y = trainer.build_feature_matrix(
factors_no_hash, data_file=str(tmp_path / "data.h5")
)
assert X is None
assert y is None
# ---------------------------------------------------------------------------
# Test: train_lightgbm
# ---------------------------------------------------------------------------
class TestTrainLightGBM:
"""Test LightGBM training."""
def _mock_lgbm_regressor(self):
"""Helper to create a properly configured mock LGBMRegressor."""
mock_model = MagicMock()
mock_model.predict.side_effect = lambda x: np.random.randn(len(x)) * 0.1
mock_booster = MagicMock()
mock_booster.feature_importance.return_value = np.array([100, 80, 60, 40, 20])
mock_model.booster_ = mock_booster
return mock_model
def test_train_with_mock(self, trainer, sample_feature_matrix):
"""Test training with mocked LightGBM."""
X, y = sample_feature_matrix
mock_model = self._mock_lgbm_regressor()
with patch("lightgbm.LGBMRegressor", return_value=mock_model):
result = trainer.train_lightgbm(X, y)
assert result is not None
assert "model" in result
assert "feature_importance" in result
assert "ic_train" in result
assert "ic_valid" in result
assert result["n_features"] == 5
def test_time_series_split(self, trainer, sample_feature_matrix):
"""Test chronological train/val split."""
X, y = sample_feature_matrix
mock_model = self._mock_lgbm_regressor()
with patch("lightgbm.LGBMRegressor", return_value=mock_model):
result = trainer.train_lightgbm(X, y, time_series_split=True)
assert result is not None
# 80/20 split
assert result["n_samples_train"] == 800
assert result["n_samples_valid"] == 200
def test_custom_params(self, trainer, sample_feature_matrix):
"""Test training with custom hyperparameters."""
X, y = sample_feature_matrix
mock_model = self._mock_lgbm_regressor()
with patch("lightgbm.LGBMRegressor", return_value=mock_model):
custom_params = {
"n_estimators": 100,
"max_depth": 4,
"learning_rate": 0.1,
}
result = trainer.train_lightgbm(X, y, params=custom_params)
assert result is not None
assert result["params"]["n_estimators"] == 100
assert result["params"]["max_depth"] == 4
def test_lightgbm_not_installed(self, trainer, sample_feature_matrix):
"""Test graceful degradation when LightGBM is missing."""
X, y = sample_feature_matrix
import sys
orig = sys.modules.get("lightgbm")
sys.modules["lightgbm"] = None # type: ignore[assignment]
try:
# Force fresh import check in train_lightgbm
import importlib
import rdagent.scenarios.qlib.local.ml_trainer as mt
importlib.reload(mt)
trainer2 = mt.MLTrainer(models_dir=trainer.models_dir)
result = trainer2.train_lightgbm(X, y)
assert result is None
finally:
if orig is not None:
sys.modules["lightgbm"] = orig
elif "lightgbm" in sys.modules:
del sys.modules["lightgbm"]
def test_feature_importance_extraction(self, trainer, sample_feature_matrix):
"""Test feature importance is correctly extracted."""
X, y = sample_feature_matrix
mock_model = self._mock_lgbm_regressor()
with patch("lightgbm.LGBMRegressor", return_value=mock_model):
result = trainer.train_lightgbm(X, y)
assert "momentum_5min" in result["feature_importance"]
assert result["feature_importance"]["momentum_5min"]["gain"] == 100.0
assert "gain_normalized" in result["feature_importance"]["momentum_5min"]
assert len(result["feature_ranking"]) == 5
assert result["feature_ranking"][0] == "momentum_5min"
def test_metrics_calculated(self, trainer, sample_feature_matrix):
"""Test that all metrics are calculated."""
X, y = sample_feature_matrix
mock_model = MagicMock()
# Predictions correlated with y for valid metrics
mock_model.predict.side_effect = lambda x: x.iloc[:, 0].values * 0.5
mock_booster = MagicMock()
mock_booster.feature_importance.return_value = np.array([100, 100, 100, 100, 100])
mock_model.booster_ = mock_booster
with patch("lightgbm.LGBMRegressor", return_value=mock_model):
result = trainer.train_lightgbm(X, y)
assert "ic_train" in result
assert "ic_valid" in result
assert "rank_ic_train" in result
assert "rank_ic_valid" in result
assert "sharpe_valid" in result
assert "mse_train" in result
assert "mse_valid" in result
# ---------------------------------------------------------------------------
# Test: extract_feature_importance
# ---------------------------------------------------------------------------
class TestExtractFeatureImportance:
"""Test feature importance extraction."""
def test_extract_all(self, trainer, mock_model_info):
"""Test extracting all feature importances."""
df = trainer.extract_feature_importance(mock_model_info)
assert len(df) == 5
assert df.iloc[0]["feature"] == "momentum_5min"
assert "gain" in df.columns
assert "split" in df.columns
def test_extract_top_n(self, trainer, mock_model_info):
"""Test extracting only top N features."""
df = trainer.extract_feature_importance(mock_model_info, top_n=3)
assert len(df) == 3
assert df.iloc[0]["feature"] == "momentum_5min"
def test_empty_importance(self, trainer):
"""Test handling empty importance dict."""
df = trainer.extract_feature_importance({})
assert df.empty
def test_sorted_by_gain(self, trainer, mock_model_info):
"""Test that result is sorted by gain descending."""
df = trainer.extract_feature_importance(mock_model_info)
gains = df["gain"].values
assert all(gains[i] >= gains[i + 1] for i in range(len(gains) - 1))
# ---------------------------------------------------------------------------
# Test: save_model
# ---------------------------------------------------------------------------
class TestSaveModel:
"""Test model persistence."""
def test_save_model(self, trainer, mock_model_info):
"""Test saving model to directory."""
mock_model_info["model"].save_model = MagicMock()
path = trainer.save_model(mock_model_info, model_name="test_model")
assert path is not None
assert path.exists()
assert (path / "model.txt").exists() or mock_model_info["model"].save_model.called
assert (path / "metadata.json").exists()
assert (path / "feature_importance.json").exists()
def test_save_model_metadata(self, trainer, mock_model_info, tmp_path):
"""Test metadata is correctly saved."""
mock_model_info["model"].save_model = MagicMock()
path = trainer.save_model(mock_model_info, model_name="test_meta")
metadata_file = path / "metadata.json"
with open(metadata_file, encoding="utf-8") as fh:
metadata = json.load(fh)
assert metadata["model_type"] == "LightGBM"
assert metadata["ic_valid"] == 0.10
assert metadata["n_features"] == 5
def test_save_model_feature_importance(self, trainer, mock_model_info):
"""Test feature importance JSON is saved."""
mock_model_info["model"].save_model = MagicMock()
path = trainer.save_model(mock_model_info, model_name="test_importance")
imp_file = path / "feature_importance.json"
with open(imp_file, encoding="utf-8") as fh:
importance = json.load(fh)
assert "momentum_5min" in importance
assert importance["momentum_5min"]["gain"] == 500.0
def test_save_model_csv(self, trainer, mock_model_info):
"""Test feature importance CSV is saved."""
mock_model_info["model"].save_model = MagicMock()
path = trainer.save_model(mock_model_info, model_name="test_csv")
csv_file = path / "feature_importance.csv"
assert csv_file.exists()
df = pd.read_csv(csv_file)
assert "feature" in df.columns
assert "gain" in df.columns
def test_save_model_none_info(self, trainer):
"""Test saving with None model_info."""
path = trainer.save_model(None)
assert path is None
def test_save_model_missing_model(self, trainer):
"""Test saving without model object."""
path = trainer.save_model({"feature_names": []})
assert path is None
def test_save_model_default_name(self, trainer, mock_model_info):
"""Test auto-generated model name."""
mock_model_info["model"].save_model = MagicMock()
path = trainer.save_model(mock_model_info)
assert path is not None
assert "lgbm_" in path.name
# ---------------------------------------------------------------------------
# Test: load_model
# ---------------------------------------------------------------------------
class TestLoadModel:
"""Test model loading."""
def test_load_model_by_name(self, trainer, mock_model_info, tmp_path):
"""Test loading model by name."""
# Train a real model to have a valid model.txt file
import lightgbm as lgb
X = pd.DataFrame(np.random.randn(200, 3), columns=["f1", "f2", "f3"])
y = pd.Series(np.random.randn(200))
real_model = lgb.LGBMRegressor(n_estimators=10, max_depth=3, verbose=-1)
real_model.fit(X, y)
model_dir = trainer.models_dir / "load_test"
model_dir.mkdir(parents=True, exist_ok=True)
# Save via booster
real_model.booster_.save_model(str(model_dir / "model.txt"))
metadata = {
"model_type": "LightGBM",
"ic_valid": 0.1,
"n_features": 3,
"feature_names": ["f1", "f2", "f3"],
"feature_importance": {"f1": {"gain": 100, "split": 5}},
}
with open(model_dir / "metadata.json", "w", encoding="utf-8") as fh:
json.dump(metadata, fh)
loaded = trainer.load_model(model_name="load_test")
assert loaded is not None
assert "model" in loaded
assert loaded["model_type"] == "LightGBM"
def test_load_model_nonexistent(self, trainer):
"""Test loading non-existent model."""
loaded = trainer.load_model(model_name="nonexistent")
assert loaded is None
def test_load_model_no_models_dir(self, trainer, tmp_path):
"""Test loading when no models exist."""
empty_dir = tmp_path / "empty"
empty_dir.mkdir()
trainer.models_dir = empty_dir
loaded = trainer.load_model()
assert loaded is None
# ---------------------------------------------------------------------------
# Test: generate_feedback
# ---------------------------------------------------------------------------
class TestGenerateFeedback:
"""Test feedback generation for factor generation loop."""
def test_generate_feedback_success(self, trainer, mock_model_info):
"""Test feedback generation with valid model info."""
feedback = trainer.generate_feedback(mock_model_info)
assert feedback["status"] == "success"
assert "suggestions" in feedback
assert len(feedback["suggestions"]) > 0
assert "factor_type_analysis" in feedback
assert feedback["n_high_importance"] > 0
def test_feedback_contains_suggestions(self, trainer, mock_model_info):
"""Test that feedback contains actionable suggestions."""
feedback = trainer.generate_feedback(mock_model_info)
assert len(feedback["suggestions"]) >= 1
assert all(isinstance(s, str) for s in feedback["suggestions"])
def test_feedback_top_5_features(self, trainer, mock_model_info):
"""Test top 5 features are included."""
feedback = trainer.generate_feedback(mock_model_info)
assert "top_5_features" in feedback
assert len(feedback["top_5_features"]) <= 5
def test_feedback_factor_type_analysis(self, trainer, mock_model_info):
"""Test factor type analysis is included."""
feedback = trainer.generate_feedback(mock_model_info)
analysis = feedback["factor_type_analysis"]
assert "momentum" in analysis
assert "mean_reversion" in analysis
assert "volatility" in analysis
assert "volume" in analysis
def test_feedback_low_ic_warning(self, trainer, mock_model_info):
"""Test feedback warns about low IC."""
mock_model_info["ic_valid"] = 0.02 # Below threshold
feedback = trainer.generate_feedback(mock_model_info)
assert any("IC is low" in s for s in feedback["suggestions"])
def test_feedback_empty_importance(self, trainer):
"""Test feedback with no importance data."""
feedback = trainer.generate_feedback({"feature_importance": {}})
assert feedback["status"] == "no_importance_data"
assert feedback["suggestions"] == []
def test_feedback_min_importance_threshold(self, trainer, mock_model_info):
"""Test min_importance_threshold affects classification."""
feedback_low = trainer.generate_feedback(
mock_model_info, min_importance_threshold=0.01
)
feedback_high = trainer.generate_feedback(
mock_model_info, min_importance_threshold=0.50
)
assert feedback_low["n_high_importance"] >= feedback_high["n_high_importance"]
# ---------------------------------------------------------------------------
# Test: save_feedback
# ---------------------------------------------------------------------------
class TestSaveFeedback:
"""Test feedback persistence."""
def test_save_feedback(self, trainer, mock_model_info):
"""Test saving feedback to JSON."""
feedback = trainer.generate_feedback(mock_model_info)
path = trainer.save_feedback(feedback)
assert path.exists()
assert path.suffix == ".json"
with open(path, encoding="utf-8") as fh:
loaded = json.load(fh)
assert loaded["status"] == "success"
def test_save_feedback_custom_path(self, trainer, mock_model_info, tmp_path):
"""Test saving feedback to custom path."""
feedback = trainer.generate_feedback(mock_model_info)
custom_path = str(tmp_path / "custom_feedback.json")
path = trainer.save_feedback(feedback, output_path=custom_path)
assert str(path) == custom_path
assert path.exists()
# ---------------------------------------------------------------------------
# Test: train_top_factors (full pipeline)
# ---------------------------------------------------------------------------
class TestTrainTopFactors:
"""Test complete training pipeline."""
@patch.object(MLTrainer, "load_top_factors")
@patch.object(MLTrainer, "build_feature_matrix")
@patch.object(MLTrainer, "train_lightgbm")
@patch.object(MLTrainer, "save_model")
@patch.object(MLTrainer, "generate_feedback")
@patch.object(MLTrainer, "save_feedback")
def test_full_pipeline(
self,
mock_save_feedback,
mock_gen_feedback,
mock_save_model,
mock_train,
mock_build,
mock_load,
trainer,
sample_feature_matrix,
mock_model_info,
):
"""Test complete pipeline with all steps mocked."""
X, y = sample_feature_matrix
mock_load.return_value = [{"factor_name": "f1", "ic": 0.1, "status": "success"}]
mock_build.return_value = (X, y)
mock_train.return_value = mock_model_info
mock_save_model.return_value = Path("/fake/path")
mock_gen_feedback.return_value = {
"status": "success",
"suggestions": ["test"],
}
mock_save_feedback.return_value = Path("/fake/feedback.json")
result = trainer.train_top_factors(top_n=10, min_ic=0.05)
assert result is not None
assert "feedback" in result
assert "model_path" in result
assert "feedback_path" in result
def test_pipeline_no_factors(self, trainer):
"""Test pipeline when no factors found."""
with patch.object(MLTrainer, "load_top_factors", return_value=[]):
result = trainer.train_top_factors()
assert result is None
def test_pipeline_no_feature_matrix(self, trainer):
"""Test pipeline when feature matrix build fails."""
with patch.object(MLTrainer, "load_top_factors", return_value=[{"ic": 0.1}]):
with patch.object(MLTrainer, "build_feature_matrix", return_value=(None, None)):
result = trainer.train_top_factors()
assert result is None
def test_pipeline_training_fails(self, trainer):
"""Test pipeline when training fails."""
with patch.object(MLTrainer, "load_top_factors", return_value=[{"ic": 0.1}]):
with patch.object(
MLTrainer, "build_feature_matrix", return_value=(pd.DataFrame(), pd.Series())
):
with patch.object(MLTrainer, "train_lightgbm", return_value=None):
result = trainer.train_top_factors()
assert result is None
# ---------------------------------------------------------------------------
# Test: get_ml_trainer factory
# ---------------------------------------------------------------------------
class TestGetMLTrainer:
"""Test factory function."""
def test_factory_returns_trainer_when_lgb_available(self):
"""Test factory returns MLTrainer when LightGBM available."""
try:
import lightgbm # noqa: F401
has_lgb = True
except ImportError:
has_lgb = False
if has_lgb:
trainer = get_ml_trainer()
assert trainer is not None
# Check it has the expected methods
assert hasattr(trainer, "train_lightgbm")
assert hasattr(trainer, "load_top_factors")
assert hasattr(trainer, "generate_feedback")
else:
# Skip test if LightGBM not installed
pytest.skip("LightGBM not installed")
def test_factory_returns_none_without_lgb(self):
"""Test factory returns None when LightGBM missing."""
import sys
orig = sys.modules.get("lightgbm")
# Temporarily remove lightgbm from modules
if "lightgbm" in sys.modules:
del sys.modules["lightgbm"]
sys.modules["lightgbm"] = None # type: ignore[assignment]
try:
# Force reimport in get_ml_trainer
result = get_ml_trainer()
assert result is None
finally:
if orig is not None:
sys.modules["lightgbm"] = orig
elif "lightgbm" in sys.modules:
del sys.modules["lightgbm"]