203 lines
7.8 KiB
Python
203 lines
7.8 KiB
Python
"""TradeAC lake configuration helpers.
|
|
|
|
The lake is a hive-partitioned parquet store (see ``tac-engine/skills/tradeac-lake``):
|
|
|
|
$TAC_LAKE_DIR/
|
|
├── market=US/
|
|
│ └── timeframe=1d/
|
|
│ └── symbol=AAPL.parquet # OHLCV bars: t, date, o, h, l, c, v, n, vw
|
|
├── features/ # indicators, wide format, family tier
|
|
│ └── market=US/
|
|
│ └── timeframe=1d/
|
|
│ ├── family=ta/symbol=AAPL.parquet # t, sma_5, sma_20, rsi_14, ...
|
|
│ └── family=sp/symbol=AAPL.parquet # t, sp_ou_*, sp_hmm_*, ...
|
|
├── calendar.parquet # trading days per market
|
|
├── coverage.parquet # per (market,timeframe,symbol) loaded windows
|
|
└── symbols.parquet # asset master
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
import pandas as pd
|
|
|
|
#: qlib freq string (Freq.__str__) -> lake timeframe partition name
|
|
FREQ_TO_TIMEFRAME: Dict[str, str] = {
|
|
"day": "1d",
|
|
"1d": "1d",
|
|
"min": "1m",
|
|
"1min": "1m",
|
|
"5min": "5m",
|
|
"10min": "10m",
|
|
"15min": "15m",
|
|
"30min": "30m",
|
|
"hour": "1h",
|
|
"1hour": "1h",
|
|
"2hour": "2h",
|
|
"4hour": "4h",
|
|
"week": "1w",
|
|
"1week": "1w",
|
|
"month": "1M",
|
|
"1month": "1M",
|
|
}
|
|
|
|
#: bar-field map: qlib field name (without the leading ``$``) -> lake bar column
|
|
BAR_FIELD_MAP: Dict[str, str] = {
|
|
"open": "o",
|
|
"high": "h",
|
|
"low": "l",
|
|
"close": "c",
|
|
"volume": "v",
|
|
"vwap": "vw",
|
|
"avg_amount": "vw", # amount / volume
|
|
}
|
|
|
|
#: fields that qlib core/backtest queries but the lake does not store -> all-NaN
|
|
UNKNOWN_FIELD_NAMES = ("factor", "change", "trade_unit", "suspend_flag")
|
|
|
|
#: columns in the parquet files that are not features
|
|
NON_FEATURE_COLUMNS = ("t", "date", "market", "timeframe", "symbol")
|
|
|
|
|
|
def timeframe_for_freq(freq: str) -> str:
|
|
"""Map a qlib frequency (e.g. ``day``, ``1min``) to a lake timeframe (e.g. ``1d``)."""
|
|
f = str(freq).lower()
|
|
if f not in FREQ_TO_TIMEFRAME:
|
|
raise ValueError(
|
|
f"unsupported qlib freq {freq!r}; supported freqs: {sorted(set(FREQ_TO_TIMEFRAME))}"
|
|
)
|
|
return FREQ_TO_TIMEFRAME[f]
|
|
|
|
|
|
def resolve_lake_root(lake_root: Optional[str] = None) -> Path:
|
|
"""Resolve the lake root: explicit arg > ``TAC_LAKE_DIR`` (no fallback).
|
|
|
|
``TAC_LAKE_DIR`` is **mandatory** — there is deliberately no default
|
|
A missing/empty value raises so a
|
|
misconfigured environment never silently points at a wrong directory.
|
|
"""
|
|
if lake_root is None:
|
|
lake_root = os.environ.get("TAC_LAKE_DIR")
|
|
if not lake_root:
|
|
raise RuntimeError(
|
|
"TAC_LAKE_DIR is not set. Point it at the TradeAC lake root, e.g. "
|
|
"export TAC_LAKE_DIR=/home/data/lake (docker) or set an absolute "
|
|
"path in your local .env."
|
|
)
|
|
return Path(str(lake_root)).expanduser().resolve()
|
|
|
|
|
|
class LakeConfig:
|
|
"""Path helpers + cached readers for a (lake_root, market) combination."""
|
|
|
|
def __init__(self, lake_root: Optional[str] = None, market: str = "US"):
|
|
self.lake_root: Path = resolve_lake_root(lake_root)
|
|
self.market: str = (market or "US").upper()
|
|
|
|
# ---- paths --------------------------------------------------------------
|
|
def bar_dir(self, timeframe: str) -> Path:
|
|
return self.lake_root / f"market={self.market}" / f"timeframe={timeframe}"
|
|
|
|
def bar_path(self, timeframe: str, symbol: str) -> Path:
|
|
return self.bar_dir(timeframe) / f"symbol={str(symbol).upper()}.parquet"
|
|
|
|
def features_dir(self, timeframe: str) -> Path:
|
|
return self.lake_root / "features" / f"market={self.market}" / f"timeframe={timeframe}"
|
|
|
|
def features_path(self, timeframe: str, symbol: str) -> Path:
|
|
# Legacy flat path (no family tier). Prefer `load_features` which
|
|
# resolves the family=ta|sp partition layout.
|
|
return self.features_dir(timeframe) / f"symbol={str(symbol).upper()}.parquet"
|
|
|
|
def load_features(self, timeframe: str, symbol: str) -> pd.DataFrame:
|
|
"""All feature columns for a symbol, merging the `family=ta` and
|
|
`family=sp` partitions by timestamp. Returns an empty frame when no
|
|
feature files exist (legacy flat layout falls back transparently)."""
|
|
sym = str(symbol).upper()
|
|
frames = []
|
|
for family in ("ta", "sp"):
|
|
p = self.features_dir(timeframe) / f"family={family}" / f"symbol={sym}.parquet"
|
|
if p.exists():
|
|
frames.append(pd.read_parquet(p))
|
|
if not frames:
|
|
flat = self.features_dir(timeframe) / f"symbol={sym}.parquet"
|
|
if flat.exists():
|
|
return pd.read_parquet(flat)
|
|
return pd.DataFrame()
|
|
if len(frames) == 1:
|
|
return frames[0]
|
|
merged = frames[0]
|
|
for extra in frames[1:]:
|
|
merged = merged.merge(extra, on="t", how="outer", suffixes=("", "_dup"))
|
|
for c in [c for c in merged.columns if c.endswith("_dup")]:
|
|
merged = merged.drop(columns=c)
|
|
return merged
|
|
|
|
def calendar_path(self) -> Path:
|
|
return self.lake_root / "calendar.parquet"
|
|
|
|
def symbols_path(self) -> Path:
|
|
return self.lake_root / "symbols.parquet"
|
|
|
|
def coverage_path(self) -> Path:
|
|
return self.lake_root / "coverage.parquet"
|
|
|
|
# ---- metadata readers ----------------------------------------------------
|
|
def load_symbols(self) -> List[str]:
|
|
"""All symbols known to the lake (from ``symbols.parquet``)."""
|
|
p = self.symbols_path()
|
|
if not p.exists():
|
|
return []
|
|
df = pd.read_parquet(p)
|
|
if "symbol" not in df.columns:
|
|
return []
|
|
return sorted(df["symbol"].astype(str).str.upper().tolist())
|
|
|
|
def symbol_spans(self, symbol: str, timeframe: str) -> List[tuple]:
|
|
"""Listing span(s) ``[(start_iso, end_iso)]`` for a symbol from coverage.parquet."""
|
|
p = self.coverage_path()
|
|
if p.exists():
|
|
try:
|
|
df = pd.read_parquet(p)
|
|
except Exception: # pragma: no cover - defensive
|
|
df = pd.DataFrame()
|
|
if len(df):
|
|
df = df[
|
|
(df.get("market") == self.market)
|
|
& (df.get("timeframe") == timeframe)
|
|
& (df.get("symbol") == str(symbol).upper())
|
|
]
|
|
if len(df):
|
|
row = df.iloc[0]
|
|
first = pd.Timestamp(row["first_t"]).date()
|
|
last = pd.Timestamp(row["last_t"]).date()
|
|
return [(first.isoformat(), last.isoformat())]
|
|
# fallback: derive from the bar file itself
|
|
p = self.bar_path(timeframe, symbol)
|
|
if p.exists():
|
|
import pyarrow.parquet as pq
|
|
|
|
tbl = pq.read_table(p, columns=["t"])
|
|
first = pd.Timestamp(tbl.column("t")[0].as_py()).date()
|
|
last = pd.Timestamp(tbl.column("t")[-1].as_py()).date()
|
|
return [(first.isoformat(), last.isoformat())]
|
|
return [("1970-01-01", "2099-12-31")]
|
|
|
|
def load_calendar_dates(self) -> List[pd.Timestamp]:
|
|
"""Trading days (midnight timestamps) for the market, from ``calendar.parquet``."""
|
|
p = self.calendar_path()
|
|
if p.exists():
|
|
df = pd.read_parquet(p)
|
|
if "date" in df.columns:
|
|
if "market" in df.columns:
|
|
df = df[df["market"] == self.market]
|
|
dates = pd.to_datetime(df["date"]).dt.normalize().sort_values().unique()
|
|
return [pd.Timestamp(x) for x in dates]
|
|
return []
|
|
|
|
def __repr__(self) -> str: # pragma: no cover
|
|
return f"LakeConfig(lake_root={self.lake_root}, market={self.market})"
|