224 lines
8.4 KiB
Python
224 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import subprocess
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from cbpred.pipeline import (
|
|
CSV_COLUMNS,
|
|
DAY_MS,
|
|
EIGHT_H_MS,
|
|
aggregate_ohlcv,
|
|
aggregate_8h,
|
|
build_feature_dataset,
|
|
compute_indicators,
|
|
fit_asset_scaler,
|
|
_vectorized_window_features,
|
|
_window_feature_names,
|
|
_window_features,
|
|
run_training,
|
|
Symbol,
|
|
run_banbot_download,
|
|
)
|
|
|
|
|
|
def make_hour_rows(start_ms: int, hours: int, phase: float = 0.0):
|
|
rows = []
|
|
price = 100.0 + phase
|
|
for index in range(hours):
|
|
timestamp = start_ms + index * 60 * 60 * 1000
|
|
drift = 0.0015 * np.sin(index / 17 + phase)
|
|
open_ = price
|
|
close = price * (1 + drift)
|
|
high = max(open_, close) * 1.003
|
|
low = min(open_, close) * 0.997
|
|
rows.append((timestamp, open_, high, low, close, 1000 + index, 100000 + index * 10, 500, 10))
|
|
price = close
|
|
return rows
|
|
|
|
|
|
class PipelineTests(unittest.TestCase):
|
|
def test_download_skips_invalid_current_symbol_and_retries_chunk(self):
|
|
valid = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0)
|
|
stale = Symbol(2, "binance", "linear", "SXP/USDT:USDT", 1, 0)
|
|
calls = []
|
|
|
|
def run(command, **kwargs):
|
|
calls.append(command)
|
|
if len(calls) == 1:
|
|
raise subprocess.CalledProcessError(
|
|
1, command, stderr="Error: [InvalidSymbol] symbol SXP/USDT:USDT not found\n"
|
|
)
|
|
|
|
with patch("cbpred.pipeline.subprocess.run", side_effect=run):
|
|
downloaded, unavailable = run_banbot_download(
|
|
[valid, stale], Path("/banbot"), Path("/banbot/config.yml"), 1, 2
|
|
)
|
|
|
|
self.assertEqual(downloaded, [valid])
|
|
self.assertEqual(unavailable, [stale])
|
|
self.assertEqual(calls[0][calls[0].index("--pairs") + 1], "BTC/USDT:USDT,SXP/USDT:USDT")
|
|
self.assertEqual(calls[1][calls[1].index("--pairs") + 1], "BTC/USDT:USDT")
|
|
self.assertIn("market_type: linear", calls[0][calls[0].index("--config-data") + 1])
|
|
|
|
def test_aggregate_8h_drops_incomplete_buckets(self):
|
|
rows = make_hour_rows(0, 17)
|
|
frame = aggregate_8h(rows)
|
|
self.assertEqual(len(frame), 2)
|
|
self.assertEqual(frame.iloc[0]["timestamp"], 0)
|
|
self.assertAlmostEqual(frame.iloc[0]["quote"], sum(row[6] for row in rows[:8]))
|
|
|
|
def test_aggregate_1d_drops_incomplete_buckets(self):
|
|
rows = make_hour_rows(0, 49)
|
|
frame = aggregate_ohlcv(rows, "1d")
|
|
self.assertEqual(len(frame), 2)
|
|
self.assertEqual(frame.iloc[0]["timestamp"], 0)
|
|
self.assertEqual(frame.iloc[1]["timestamp"], DAY_MS)
|
|
|
|
def test_daily_feature_dataset_uses_daily_continuity(self):
|
|
raw = aggregate_ohlcv(make_hour_rows(0, 24 * 80), "1d")
|
|
profiles = {
|
|
"x": {
|
|
"asset_atr_price_top20_median": 0.01,
|
|
"asset_quote_volume_top20_median_log": 12.0,
|
|
}
|
|
}
|
|
scaler = fit_asset_scaler(profiles)
|
|
dataset, _, _ = build_feature_dataset(
|
|
{"x": raw},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + DAY_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
step_ms=DAY_MS,
|
|
)
|
|
self.assertFalse(dataset.empty)
|
|
|
|
def test_future_target_changes_but_past_features_do_not(self):
|
|
start = 0
|
|
raw = aggregate_8h(make_hour_rows(start, 8 * 70))
|
|
indicators = compute_indicators(raw)
|
|
profiles = {"x": {"asset_atr_price_top20_median": 0.01, "asset_quote_volume_top20_median_log": 12.0}}
|
|
scaler = fit_asset_scaler(profiles)
|
|
one, names, _ = build_feature_dataset(
|
|
{"x": raw},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
changed = raw.copy()
|
|
changed.loc[62:, "low"] *= 0.5
|
|
two, _, _ = build_feature_dataset(
|
|
{"x": changed},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
anchor = int(raw.iloc[60]["timestamp"])
|
|
first = one[one["timestamp"] == anchor].iloc[0]
|
|
second = two[two["timestamp"] == anchor].iloc[0]
|
|
self.assertEqual(len(names), 18 * 10 + 2)
|
|
self.assertGreater(first["target_scale"], 0.0)
|
|
for name in names:
|
|
if name != "target":
|
|
self.assertEqual(first[name], second[name])
|
|
self.assertNotEqual(first["target"], second["target"])
|
|
|
|
rising = raw.copy()
|
|
anchor_close = float(rising.iloc[60]["close"])
|
|
for index in range(61, 67):
|
|
open_ = anchor_close * (1.01 + 0.002 * (index - 61))
|
|
close = open_ * 1.002
|
|
rising.loc[index, "open"] = open_
|
|
rising.loc[index, "low"] = open_ * 0.999
|
|
rising.loc[index, "high"] = close * 1.001
|
|
rising.loc[index, "close"] = close
|
|
rising_dataset, _, _ = build_feature_dataset(
|
|
{"x": rising},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
rising_row = rising_dataset[rising_dataset["timestamp"] == anchor].iloc[0]
|
|
self.assertGreater(rising_row["target"], 0.0)
|
|
expected_target = (
|
|
np.quantile(rising.loc[61:66, "low"].to_numpy(dtype=float), 0.20)
|
|
/ anchor_close
|
|
- 1.0
|
|
) / rising_row["target_scale"]
|
|
self.assertAlmostEqual(rising_row["target"], expected_target)
|
|
|
|
def test_vectorized_features_match_window_features(self):
|
|
raw = aggregate_8h(make_hour_rows(0, 8 * 80))
|
|
indicators = compute_indicators(raw)
|
|
asset = {
|
|
"asset_atr_price_top20_median": 0.01,
|
|
"asset_quote_volume_top20_median_log": 12.0,
|
|
}
|
|
for stage in (1, 2, 3):
|
|
names = _window_feature_names(stage)
|
|
for anchor in (18, 30, 55):
|
|
window = indicators.iloc[anchor - 17 : anchor + 1]
|
|
expected = _window_features(window, stage, asset)
|
|
actual = _vectorized_window_features(
|
|
indicators, stage, np.array([anchor]), asset
|
|
)[0]
|
|
np.testing.assert_allclose(
|
|
actual,
|
|
np.array([expected[name] for name in names]),
|
|
rtol=1e-12,
|
|
atol=1e-12,
|
|
equal_nan=True,
|
|
)
|
|
|
|
def test_three_stage_training_writes_report_and_model(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
data_dir = root / "data"
|
|
raw_dir = data_dir / "raw"
|
|
raw_dir.mkdir(parents=True)
|
|
symbols = []
|
|
end_ms = 0
|
|
for idx in range(4):
|
|
symbol = f"COIN{idx}/USDT:USDT"
|
|
frame = aggregate_8h(make_hour_rows(0, 8 * 400, idx * 0.1))
|
|
path = raw_dir / f"{idx}.csv"
|
|
frame.to_csv(path, index=False)
|
|
end_ms = max(end_ms, int(frame.iloc[-1]["timestamp"] + EIGHT_H_MS))
|
|
symbols.append({"sid": idx, "symbol": symbol, "path": str(path.relative_to(data_dir)), "bars": len(frame)})
|
|
(data_dir / "manifest.json").write_text(
|
|
json.dumps({"end_ms": end_ms, "symbols": symbols})
|
|
)
|
|
report = run_training(
|
|
data_dir,
|
|
root / "artifacts",
|
|
end_ms=end_ms,
|
|
test_months=1,
|
|
n_estimators=4,
|
|
n_jobs=1,
|
|
)
|
|
self.assertIn(report["selected"]["stage"], {1, 2, 3})
|
|
self.assertEqual(len(report["stages"]), 3)
|
|
self.assertTrue((root / "artifacts" / "model.joblib").exists())
|
|
self.assertTrue((root / "artifacts" / "report.md").exists())
|
|
self.assertGreaterEqual(report["test"]["calibrated_q80_coverage"], 0.0)
|
|
self.assertEqual(report["protocol"]["volatility_filter"]["atr_rank_window"], 200)
|
|
self.assertAlmostEqual(report["protocol"]["volatility_filter"]["min_rank"], 0.0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|