init

tests/test_pipeline.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from cbpred.pipeline import (
+    CSV_COLUMNS,
+    EIGHT_H_MS,
+    aggregate_8h,
+    build_feature_dataset,
+    compute_indicators,
+    fit_asset_scaler,
+    _vectorized_window_features,
+    _window_feature_names,
+    _window_features,
+    run_training,
+)
+
+
+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_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_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()