data loader with sanity assertions

578ee58 verified 2 days ago

4.67 kB

"""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())

	"""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())