data loader with sanity assertions

crypto_eval/data.py

@@ -0,0 +1,113 @@
+"""Data loading for the crypto-quant-eval dataset.
+
+Loads curated parquet files from the dataset repo or a local copy,
+asserts schema and no-lookahead invariants, and exposes simple accessors.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+REQUIRED_BARS = {
+    "datetime", "open", "high", "low", "close", "volume",
+}
+REQUIRED_FUNDING = {"timestamp", "funding_rate"}
+
+
+def sha256_file(path: str | Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+@dataclass
+class EvalData:
+    """Curated bar data plus aligned funding ticks."""
+
+    bars: pd.DataFrame        # index: UTC datetime, OHLCV columns
+    funding: pd.DataFrame     # index: UTC datetime (8h ticks), funding_rate
+    symbol: str
+    freq: str
+
+    def assert_sane(self) -> None:
+        # bar start time may be carried either as the DatetimeIndex or as
+        # a 'datetime' column; the index form is what load_local produces.
+        bar_cols = set(self.bars.columns)
+        if isinstance(self.bars.index, pd.DatetimeIndex):
+            bar_cols = bar_cols | {"datetime"}
+        missing = REQUIRED_BARS - bar_cols
+        if missing:
+            raise ValueError(f"bars missing columns: {sorted(missing)}")
+        fund_cols = set(self.funding.columns)
+        if isinstance(self.funding.index, pd.DatetimeIndex):
+            fund_cols = fund_cols | {"timestamp"}
+        missing = REQUIRED_FUNDING - fund_cols
+        if missing:
+            raise ValueError(f"funding missing columns: {sorted(missing)}")
+        if not self.bars.index.is_monotonic_increasing:
+            raise ValueError("bars index is not sorted")
+        if self.bars.index.has_duplicates:
+            raise ValueError("bars index has duplicate timestamps")
+        # no-lookahead invariant: bars are closed bars only.
+        # The curated dataset stores bar START time; a bar starting at t is
+        # only usable at t+freq. load() shifts nothing here — consumers must
+        # use shift-based accessors below, never bars.loc[t] inside a
+        # signal computed at t.
+        nulls = self.bars[["open", "high", "low", "close"]].isna().sum().sum()
+        if nulls:
+            raise ValueError(f"bars contain {nulls} null OHLC values")
+        bad_ohlc = (
+            (self.bars["high"] < self.bars["low"])
+            | (self.bars["close"] > self.bars["high"])
+            | (self.bars["open"] > self.bars["high"])
+            | (self.bars["close"] < self.bars["low"])
+            | (self.bars["open"] < self.bars["low"])
+        ).sum()
+        if bad_ohlc:
+            raise ValueError(f"{bad_ohlc} bars violate OHLC ordering")
+
+    def check_funding_coverage(self, freq: str) -> int:
+        """Return count of funding ticks falling inside the bar range."""
+        mask = (self.funding.index >= self.bars.index[0]) & (
+            self.funding.index <= self.bars.index[-1]
+        )
+        return int(mask.sum())
+
+
+def load_local(root: str | Path, symbol: str = "BTCUSDT", freq: str = "1h") -> EvalData:
+    """Load bars_<symbol>_<freq>.parquet and funding_<symbol>.parquet from root."""
+    root = Path(root)
+    bars = pd.read_parquet(root / f"bars_{symbol}_{freq}.parquet")
+    funding = pd.read_parquet(root / f"funding_{symbol}.parquet")
+    # curated parquet stores bar start time as a 'datetime' column; the
+    # synthetic/local layout keeps a DatetimeIndex. Accept both.
+    for df in (bars, funding):
+        if "datetime" in df.columns and not isinstance(df.index, pd.DatetimeIndex):
+            ts = pd.to_datetime(df["datetime"], utc=True)
+            df.index = ts
+        elif "timestamp" in df.columns and not isinstance(df.index, pd.DatetimeIndex):
+            df.index = pd.to_datetime(df["timestamp"], utc=True)
+    bars = bars.drop(columns=["datetime"], errors="ignore")
+    funding = funding.drop(columns=["datetime", "timestamp"], errors="ignore")
+    for df in (bars, funding):
+        if not isinstance(df.index, pd.DatetimeIndex):
+            raise ValueError("parquet files must carry a datetime index")
+        if df.index.tz is None:
+            df.index = df.index.tz_localize("UTC")
+        else:
+            df.index = df.index.tz_convert("UTC")
+    data = EvalData(bars=bars, funding=funding, symbol=symbol, freq=freq)
+    data.assert_sane()
+    return data
+
+
+def load_manifest(root: str | Path) -> dict:
+    """Read the source-hash manifest; every eval run must record this."""
+    return json.loads((Path(root) / "manifest.json").read_text())