mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-28 07:57:44 +00:00
feat: add Kronos CLI commands, expand tests, document in README
- predix kronos-factor: generate KronosPredReturn alpha factor via CLI - predix kronos-eval: evaluate Kronos IC/hit-rate vs LightGBM via CLI - 19 tests covering adapter, factor builder, model evaluator, CLI (mock-based) - README: Kronos section in Features + CLI commands table - Total test suite: 153 passed Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
TPTBusiness
@@ -1,102 +1,277 @@
|
||||
"""Tests for KronosAdapter — mock-based, no real model download needed."""
|
||||
"""Tests for KronosAdapter and CLI commands — mock-based, no real model download needed."""
|
||||
|
||||
import json
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
||||
def _make_ohlcv(n: int = 600) -> pd.DataFrame:
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_ohlcv(n: int = 600, freq: str = "1min") -> pd.DataFrame:
|
||||
"""Synthetic 1-min OHLCV DataFrame."""
|
||||
idx = pd.date_range("2024-01-01", periods=n, freq="1min")
|
||||
idx = pd.date_range("2024-01-01", periods=n, freq=freq)
|
||||
close = 1.1000 + np.cumsum(np.random.randn(n) * 0.0001)
|
||||
df = pd.DataFrame({
|
||||
return pd.DataFrame({
|
||||
"open": close + np.random.randn(n) * 0.00005,
|
||||
"high": close + np.abs(np.random.randn(n) * 0.0001),
|
||||
"low": close - np.abs(np.random.randn(n) * 0.0001),
|
||||
"close": close,
|
||||
"volume": np.abs(np.random.randn(n) * 100),
|
||||
}, index=idx)
|
||||
return df
|
||||
|
||||
|
||||
def test_ohlcv_conversion():
|
||||
"""_ohlcv_from_predix renames $ columns correctly."""
|
||||
from rdagent.components.coder.kronos_adapter import _ohlcv_from_predix
|
||||
|
||||
def _make_predix_hdf5(tmp_path: Path, n: int = 300) -> Path:
|
||||
"""Write a minimal Predix-format HDF5 file and return its path."""
|
||||
idx = pd.MultiIndex.from_arrays(
|
||||
[pd.date_range("2024-01-01", periods=3, freq="1min"), ["EURUSD"] * 3],
|
||||
[pd.date_range("2024-01-01", periods=n, freq="1min"), ["EURUSD"] * n],
|
||||
names=["datetime", "instrument"],
|
||||
)
|
||||
predix_df = pd.DataFrame({
|
||||
"$open": [1.1, 1.2, 1.3],
|
||||
"$high": [1.15, 1.25, 1.35],
|
||||
"$low": [1.05, 1.15, 1.25],
|
||||
"$close": [1.12, 1.22, 1.32],
|
||||
"$volume": [100.0, 200.0, 300.0],
|
||||
df = pd.DataFrame({
|
||||
"$open": (np.random.rand(n) + 1.1).astype("float32"),
|
||||
"$close": (np.random.rand(n) + 1.1).astype("float32"),
|
||||
"$high": (np.random.rand(n) + 1.11).astype("float32"),
|
||||
"$low": (np.random.rand(n) + 1.09).astype("float32"),
|
||||
"$volume": (np.random.rand(n) * 100).astype("float32"),
|
||||
}, index=idx)
|
||||
|
||||
ohlcv = _ohlcv_from_predix(predix_df)
|
||||
assert "close" in ohlcv.columns
|
||||
assert "$close" not in ohlcv.columns
|
||||
assert list(ohlcv.columns) == ["open", "high", "low", "close", "volume"]
|
||||
h5 = tmp_path / "intraday_pv.h5"
|
||||
df.to_hdf(h5, key="data", mode="w")
|
||||
return h5
|
||||
|
||||
|
||||
def test_kronos_adapter_load_skipped_without_repo(tmp_path, monkeypatch):
|
||||
"""KronosAdapter gracefully reports unavailable when repo is missing."""
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KRONOS_REPO", tmp_path / "nonexistent")
|
||||
monkeypatch.setattr(mod, "_KRONOS_AVAILABLE", None)
|
||||
|
||||
from rdagent.components.coder.kronos_adapter import KronosAdapter, _ensure_kronos
|
||||
available = _ensure_kronos()
|
||||
assert available is False
|
||||
|
||||
|
||||
def test_build_kronos_factor_mock(tmp_path, monkeypatch):
|
||||
"""build_kronos_factor produces correct MultiIndex output with mocked predictor."""
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
|
||||
# Mock the adapter so no real Kronos load happens
|
||||
def _make_mock_adapter():
|
||||
"""Return a mock KronosAdapter whose predict_next_bars is deterministic."""
|
||||
class MockAdapter:
|
||||
def load(self): return self
|
||||
def predict_next_bars(self, ohlcv_df, context_bars, pred_bars, **kw):
|
||||
idx = pd.date_range(ohlcv_df.index[-1], periods=pred_bars + 1, freq="1min")[1:]
|
||||
last_close = float(ohlcv_df["close"].iloc[-1])
|
||||
return pd.DataFrame({
|
||||
"open": last_close * (1 + np.random.randn(pred_bars) * 0.001),
|
||||
"close": last_close * (1 + np.random.randn(pred_bars) * 0.001),
|
||||
"high": last_close * 1.001,
|
||||
"open": last_close * 1.001,
|
||||
"close": last_close * 1.002,
|
||||
"high": last_close * 1.003,
|
||||
"low": last_close * 0.999,
|
||||
"volume": 100.0,
|
||||
}, index=idx)
|
||||
def predict_return(self, ohlcv_df, context_bars=512, pred_bars=1):
|
||||
return 0.001
|
||||
return MockAdapter()
|
||||
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: MockAdapter())
|
||||
|
||||
# Write minimal HDF5
|
||||
n = 300
|
||||
idx = pd.MultiIndex.from_arrays(
|
||||
[pd.date_range("2024-01-01", periods=n, freq="1min"), ["EURUSD"] * n],
|
||||
names=["datetime", "instrument"],
|
||||
)
|
||||
df = pd.DataFrame({
|
||||
"$open": np.random.rand(n).astype("float32") + 1.1,
|
||||
"$close": np.random.rand(n).astype("float32") + 1.1,
|
||||
"$high": np.random.rand(n).astype("float32") + 1.11,
|
||||
"$low": np.random.rand(n).astype("float32") + 1.09,
|
||||
"$volume": np.random.rand(n).astype("float32") * 100,
|
||||
}, index=idx)
|
||||
h5_path = tmp_path / "intraday_pv.h5"
|
||||
df.to_hdf(h5_path, key="data", mode="w")
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: _ohlcv_from_predix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
result = mod.build_kronos_factor(
|
||||
hdf5_path=h5_path,
|
||||
context_bars=100,
|
||||
pred_bars=20,
|
||||
stride_bars=20,
|
||||
device="cpu",
|
||||
)
|
||||
class TestOhlcvConversion:
|
||||
def test_renames_dollar_columns(self):
|
||||
from rdagent.components.coder.kronos_adapter import _ohlcv_from_predix
|
||||
idx = pd.MultiIndex.from_arrays(
|
||||
[pd.date_range("2024-01-01", periods=3, freq="1min"), ["EURUSD"] * 3],
|
||||
names=["datetime", "instrument"],
|
||||
)
|
||||
df = pd.DataFrame({
|
||||
"$open": [1.1, 1.2, 1.3], "$high": [1.15, 1.25, 1.35],
|
||||
"$low": [1.05, 1.15, 1.25], "$close": [1.12, 1.22, 1.32],
|
||||
"$volume": [100.0, 200.0, 300.0],
|
||||
}, index=idx)
|
||||
result = _ohlcv_from_predix(df)
|
||||
assert list(result.columns) == ["open", "high", "low", "close", "volume"]
|
||||
|
||||
assert isinstance(result, pd.DataFrame)
|
||||
assert result.index.names == ["datetime", "instrument"]
|
||||
assert "KronosPredReturn" in result.columns
|
||||
assert result["KronosPredReturn"].notna().sum() > 0
|
||||
def test_no_dollar_columns_passthrough(self):
|
||||
from rdagent.components.coder.kronos_adapter import _ohlcv_from_predix
|
||||
df = pd.DataFrame({"open": [1.0], "close": [1.1], "high": [1.2], "low": [0.9], "volume": [100.0]})
|
||||
result = _ohlcv_from_predix(df)
|
||||
assert "close" in result.columns
|
||||
|
||||
def test_output_is_float64(self):
|
||||
from rdagent.components.coder.kronos_adapter import _ohlcv_from_predix
|
||||
df = pd.DataFrame({
|
||||
"$open": np.array([1.1], dtype="float32"),
|
||||
"$close": np.array([1.1], dtype="float32"),
|
||||
"$high": np.array([1.1], dtype="float32"),
|
||||
"$low": np.array([1.1], dtype="float32"),
|
||||
"$volume": np.array([100.0], dtype="float32"),
|
||||
})
|
||||
result = _ohlcv_from_predix(df)
|
||||
assert result["close"].dtype == np.float64
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: KronosAdapter availability check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestKronosAvailability:
|
||||
def test_unavailable_without_repo(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KRONOS_REPO", tmp_path / "nonexistent")
|
||||
monkeypatch.setattr(mod, "_KRONOS_AVAILABLE", None)
|
||||
assert mod._ensure_kronos() is False
|
||||
|
||||
def test_load_raises_without_repo(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KRONOS_REPO", tmp_path / "nonexistent")
|
||||
monkeypatch.setattr(mod, "_KRONOS_AVAILABLE", None)
|
||||
adapter = mod.KronosAdapter()
|
||||
with pytest.raises(RuntimeError, match="Kronos not available"):
|
||||
adapter.load()
|
||||
|
||||
def test_predict_without_load_raises(self):
|
||||
from rdagent.components.coder.kronos_adapter import KronosAdapter
|
||||
adapter = KronosAdapter()
|
||||
with pytest.raises(RuntimeError, match="Call .load()"):
|
||||
adapter.predict_next_bars(_make_ohlcv(100), 50, 10)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: build_kronos_factor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestBuildKronosFactor:
|
||||
def test_output_has_correct_multiindex(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path)
|
||||
result = mod.build_kronos_factor(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert result.index.names == ["datetime", "instrument"]
|
||||
assert result.index.nlevels == 2
|
||||
|
||||
def test_output_column_name(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path)
|
||||
result = mod.build_kronos_factor(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert "KronosPredReturn" in result.columns
|
||||
|
||||
def test_output_has_non_nan_values(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path)
|
||||
result = mod.build_kronos_factor(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert result["KronosPredReturn"].notna().sum() > 0
|
||||
|
||||
def test_output_length_matches_input(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
n = 300
|
||||
h5 = _make_predix_hdf5(tmp_path, n=n)
|
||||
result = mod.build_kronos_factor(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert len(result) == n
|
||||
|
||||
def test_forward_fill_propagates_signal(self, tmp_path, monkeypatch):
|
||||
"""Values within a predicted window should be forward-filled, not NaN."""
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path, n=300)
|
||||
result = mod.build_kronos_factor(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
non_nan_ratio = result["KronosPredReturn"].notna().mean()
|
||||
assert non_nan_ratio > 0.5, f"Expected >50% non-NaN, got {non_nan_ratio:.2%}"
|
||||
|
||||
def test_raises_on_missing_hdf5(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
with pytest.raises(Exception):
|
||||
mod.build_kronos_factor(tmp_path / "missing.h5", context_bars=50, pred_bars=10, stride_bars=10)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: evaluate_kronos_model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEvaluateKronosModel:
|
||||
def test_returns_required_keys(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path, n=400)
|
||||
metrics = mod.evaluate_kronos_model(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
for key in ["IC_mean", "IC_std", "IC_IR", "hit_rate", "n_predictions"]:
|
||||
assert key in metrics, f"Missing key: {key}"
|
||||
|
||||
def test_hit_rate_in_valid_range(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path, n=400)
|
||||
metrics = mod.evaluate_kronos_model(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert 0.0 <= metrics["hit_rate"] <= 1.0
|
||||
|
||||
def test_n_predictions_positive(self, tmp_path, monkeypatch):
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
h5 = _make_predix_hdf5(tmp_path, n=400)
|
||||
metrics = mod.evaluate_kronos_model(h5, context_bars=100, pred_bars=20, stride_bars=20, device="cpu")
|
||||
assert metrics["n_predictions"] > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests: CLI commands (via typer test runner)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCLICommands:
|
||||
def test_kronos_factor_missing_data_exits(self, tmp_path, monkeypatch):
|
||||
"""kronos-factor exits with code 1 when HDF5 data is missing."""
|
||||
from typer.testing import CliRunner
|
||||
import predix as predix_mod
|
||||
monkeypatch.chdir(tmp_path)
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(predix_mod.app, ["kronos-factor"])
|
||||
assert result.exit_code == 1
|
||||
|
||||
def test_kronos_eval_missing_data_exits(self, tmp_path, monkeypatch):
|
||||
"""kronos-eval exits with code 1 when HDF5 data is missing."""
|
||||
from typer.testing import CliRunner
|
||||
import predix as predix_mod
|
||||
monkeypatch.chdir(tmp_path)
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(predix_mod.app, ["kronos-eval"])
|
||||
assert result.exit_code == 1
|
||||
|
||||
def test_kronos_factor_runs_with_mock(self, tmp_path, monkeypatch):
|
||||
"""kronos-factor completes and saves parquet + json when adapter is mocked."""
|
||||
from typer.testing import CliRunner
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
import predix as predix_mod
|
||||
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
|
||||
data_dir = tmp_path / "git_ignore_folder" / "factor_implementation_source_data"
|
||||
data_dir.mkdir(parents=True)
|
||||
_make_predix_hdf5(data_dir.parent.parent, n=300)
|
||||
h5_src = tmp_path / "intraday_pv.h5"
|
||||
# Put HDF5 where the CLI expects it
|
||||
import shutil
|
||||
src = _make_predix_hdf5(tmp_path, n=300)
|
||||
shutil.copy(src, data_dir / "intraday_pv.h5")
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(predix_mod.app, [
|
||||
"kronos-factor", "--context", "100", "--pred", "20", "--device", "cpu"
|
||||
])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "saved" in result.output.lower()
|
||||
|
||||
def test_kronos_eval_runs_with_mock(self, tmp_path, monkeypatch):
|
||||
"""kronos-eval completes and prints IC metrics when adapter is mocked."""
|
||||
from typer.testing import CliRunner
|
||||
import rdagent.components.coder.kronos_adapter as mod
|
||||
import predix as predix_mod
|
||||
|
||||
monkeypatch.setattr(mod, "KronosAdapter", lambda **kw: _make_mock_adapter())
|
||||
|
||||
data_dir = tmp_path / "git_ignore_folder" / "factor_implementation_source_data"
|
||||
data_dir.mkdir(parents=True)
|
||||
_make_predix_hdf5(data_dir.parent.parent, n=400)
|
||||
src = _make_predix_hdf5(tmp_path, n=400)
|
||||
import shutil
|
||||
shutil.copy(src, data_dir / "intraday_pv.h5")
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(predix_mod.app, [
|
||||
"kronos-eval", "--context", "100", "--pred", "20", "--device", "cpu"
|
||||
])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "IC" in result.output
|
||||
|
||||
Reference in New Issue
Block a user