You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

517 lines
21 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

"""Point-in-time snapshot data loader for the backtester.
The public interface is:
store = SnapshotStore.load(snapshot_dir, split_name, oracle_url, db_dsn)
rows = store.get_candidates_for_date(date)
bar = store.get_bar(symbol, date)
macro = store.get_macro_for_date(date)
For tests, inject data directly:
store = SnapshotStore(candidates_by_exec_date=..., bars_by_symbol_date=..., macro_by_date=...)
"""
from __future__ import annotations
import asyncio
import datetime as dt
from pathlib import Path
from typing import Any
import pyarrow.parquet as pq
from libs.common.logging import get_logger
logger = get_logger(__name__)
class SnapshotStore:
"""In-memory cache of point-in-time backtest data."""
def __init__(
self,
candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]],
bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]],
macro_by_date: dict[dt.date, dict[str, Any]] | None = None,
) -> None:
self._candidates = {
date: list(rows)
for date, rows in candidates_by_exec_date.items()
}
self._candidates_by_reaction_date = self._build_reaction_index(self._candidates)
self._bars = bars_by_symbol_date
self._macro = macro_by_date or {}
# ------------------------------------------------------------------
# Public query interface
# ------------------------------------------------------------------
def get_candidates_for_date(self, date: dt.date) -> list[dict[str, Any]]:
"""Return candidates where execution_date == date. No look-ahead."""
return list(self._candidates.get(date, []))
def get_candidates_for_reaction_date(self, date: dt.date) -> list[dict[str, Any]]:
"""Return candidates where reaction_date == date. No look-ahead."""
return list(self._candidates_by_reaction_date.get(date, []))
def get_bar(self, symbol: str, date: dt.date) -> dict[str, Any] | None:
"""Return OHLCV bar for symbol on date, or None if missing."""
sym_bars = self._bars.get(symbol)
if sym_bars is None:
return None
return sym_bars.get(date)
def get_macro_for_date(self, date: dt.date) -> dict[str, Any]:
"""Return macro observations for date (empty dict if none)."""
return dict(self._macro.get(date, {}))
def all_execution_dates(self) -> list[dt.date]:
"""Sorted list of dates that have at least one candidate."""
return sorted(self._candidates.keys())
def all_reaction_dates(self) -> list[dt.date]:
"""Sorted list of dates that have at least one reaction-date candidate."""
return sorted(self._candidates_by_reaction_date.keys())
def all_trading_days(self, include_reaction_dates: bool = False) -> list[dt.date]:
"""All NYSE trading days from first to last execution date (inclusive).
Use this to drive the simulation loop so stop/target/time exits are
checked on every trading day, not just candidate days.
"""
from libs.backtest.calendar import get_trading_days
dates = set(self.all_execution_dates())
if include_reaction_dates:
dates.update(self.all_reaction_dates())
if not dates:
return []
ordered = sorted(dates)
return get_trading_days(ordered[0], ordered[-1])
# ------------------------------------------------------------------
# Factory: load from Parquet + DB + Oracle
# ------------------------------------------------------------------
@classmethod
def load(
cls,
snapshot_dir: str | Path,
split_name: str,
oracle_url: str,
db_dsn: str,
scoring_fn: Any | None = None,
) -> "SnapshotStore":
"""Synchronous factory. Internally uses asyncio.run() to prefetch data.
scoring_fn: optional callable(row_dict) -> float to override default scoring.
Raises RuntimeError if called from within a running event loop.
"""
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None and loop.is_running():
raise RuntimeError(
"SnapshotStore.load() cannot be called from a running event loop. "
"Use SnapshotStore._async_load() directly in async contexts."
)
data = asyncio.run(
cls._async_load(Path(snapshot_dir), split_name, oracle_url, db_dsn, scoring_fn)
)
return cls(**data)
@classmethod
async def _async_load(
cls,
snapshot_dir: Path,
split_name: str,
oracle_url: str,
db_dsn: str,
scoring_fn: Any | None = None,
) -> dict[str, Any]:
"""Async data loading pipeline. Returns kwargs dict for __init__."""
parquet_path = snapshot_dir / f"{split_name}.parquet"
if not parquet_path.exists():
raise FileNotFoundError(f"Parquet file not found: {parquet_path}")
# Step 1: Read Parquet
logger.info("snapshot_store_reading_parquet", path=str(parquet_path))
table = pq.read_table(str(parquet_path))
rows: list[dict[str, Any]] = table.to_pydict()
# Convert column-oriented dict to list of row dicts
num_rows = table.num_rows
col_names = list(rows.keys())
row_list: list[dict[str, Any]] = [
{col: rows[col][i] for col in col_names} for i in range(num_rows)
]
logger.info("snapshot_store_rows_loaded", count=num_rows)
# Collect event_ids for DB lookup
event_ids = [str(r.get("event_id", "")) for r in row_list]
# Step 26: DB + Oracle enrichment
event_meta = await cls._fetch_event_metadata(event_ids, db_dsn)
unique_symbols = list({m.get("ticker", "") for m in event_meta.values() if m.get("ticker")})
date_range = cls._compute_date_range(row_list)
bars_by_symbol, avg_dvol = await cls._fetch_price_data(
unique_symbols, date_range, oracle_url
)
sectors = await cls._fetch_sectors(unique_symbols, oracle_url)
macro_by_date = await cls._fetch_macro(date_range, db_dsn)
# Fetch SPY bars for macro regime filter (SMA computation)
spy_macro = await cls._fetch_spy_macro(date_range, oracle_url)
for d, spy_data in spy_macro.items():
macro_by_date.setdefault(d, {}).update(spy_data)
# Step 7: Build candidates_by_exec_date
candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]] = {}
for row in row_list:
eid = str(row.get("event_id", ""))
meta = event_meta.get(eid, {})
ticker = meta.get("ticker")
if not ticker:
logger.debug("snapshot_store_skip_no_ticker", event_id=eid)
continue
# Map entry_date → execution_date at this boundary
raw_exec = row.get("entry_date") or row.get("execution_date")
if raw_exec is None:
continue
if isinstance(raw_exec, str):
exec_date = dt.date.fromisoformat(raw_exec)
elif isinstance(raw_exec, dt.date):
exec_date = raw_exec
else:
continue
enriched = dict(row)
enriched["execution_date"] = exec_date
enriched["symbol"] = ticker
enriched["issuer_id"] = meta.get("issuer_id")
enriched["event_date"] = meta.get("event_date")
enriched["event_type"] = meta.get("event_type", "")
enriched["event_timestamp"] = meta.get("event_timestamp")
enriched["avg_dollar_volume"] = avg_dvol.get(ticker, 0.0)
enriched["sector"] = sectors.get(ticker, "UNKNOWN")
# Map Parquet-specific columns to canonical backtest names
# event_close (reaction-day close) → entry_price_est baseline
if "entry_price_est" not in enriched and "event_close" in enriched:
enriched["entry_price_est"] = enriched["event_close"]
# score: use existing column or compute from market features
if "score" not in enriched or enriched.get("score") is None:
if scoring_fn is not None:
enriched["score"] = scoring_fn(enriched)
else:
from libs.backtest.scoring import compute_entry_score
enriched["score"] = compute_entry_score(enriched)
candidates_by_exec_date.setdefault(exec_date, []).append(enriched)
# Build bars_by_symbol_date: symbol -> date -> bar dict
bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]] = {}
for sym, date_bars in bars_by_symbol.items():
bars_by_symbol_date[sym] = date_bars
logger.info(
"snapshot_store_built",
exec_dates=len(candidates_by_exec_date),
symbols=len(bars_by_symbol_date),
)
return {
"candidates_by_exec_date": candidates_by_exec_date,
"bars_by_symbol_date": bars_by_symbol_date,
"macro_by_date": macro_by_date,
}
# ------------------------------------------------------------------
# Internal async helpers
# ------------------------------------------------------------------
@staticmethod
async def _fetch_event_metadata(
event_ids: list[str],
db_dsn: str,
) -> dict[str, dict[str, Any]]:
"""Batch-query Event + SymbolMaster for event metadata."""
if not event_ids:
return {}
try:
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from libs.db.models import Event, SymbolMaster
engine = create_async_engine(db_dsn, echo=False)
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
result: dict[str, dict[str, Any]] = {}
async with async_session() as session:
# Single IN clause for all event_ids
stmt = (
select(Event, SymbolMaster)
.outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id)
.where(Event.event_id.in_(event_ids))
)
rows = (await session.execute(stmt)).all()
_UTC = __import__("zoneinfo").ZoneInfo("UTC")
for event, sym in rows:
# Use filed_at_utc if available; fallback to filing_date + 21:00 UTC
# (transparent enrichment in the loader — not silent substitution in selector)
ts = event.filed_at_utc
if ts is None and event.event_date is not None:
ts = dt.datetime.combine(
event.event_date, dt.time(21, 0), tzinfo=_UTC
)
result[event.event_id] = {
"issuer_id": event.issuer_id,
"event_date": event.event_date,
"event_type": event.event_type,
"event_timestamp": ts,
"ticker": sym.ticker if sym else None,
}
await engine.dispose()
return result
except Exception as exc:
logger.warning("snapshot_store_db_fetch_failed", error=str(exc))
return {}
@staticmethod
async def _fetch_price_data(
symbols: list[str],
date_range: tuple[dt.date, dt.date] | None,
oracle_url: str,
concurrency: int = 16,
) -> tuple[dict[str, dict[dt.date, dict[str, Any]]], dict[str, float]]:
"""Fetch daily OHLCV bars and compute avg_dollar_volume per symbol."""
if not symbols or date_range is None:
return {}, {}
try:
from libs.oracle_client import OracleClient, PriceService
start_str = date_range[0].isoformat()
end_str = date_range[1].isoformat()
bars_by_symbol: dict[str, dict[dt.date, dict[str, Any]]] = {}
avg_dvol: dict[str, float] = {}
semaphore = asyncio.Semaphore(concurrency)
async with OracleClient(base_url=oracle_url) as client:
svc = PriceService(client)
async def _fetch_symbol(sym: str) -> tuple[str, dict[dt.date, dict[str, Any]], float]:
async with semaphore:
try:
resp = await svc.get_daily_bars(sym, start=start_str, end=end_str)
date_bars: dict[dt.date, dict[str, Any]] = {}
dollar_vols: list[float] = []
for bar in resp.bars:
d = dt.date.fromisoformat(bar.date)
b = {
"date": d,
"open": bar.open,
"high": bar.high,
"low": bar.low,
"close": bar.close,
"volume": bar.volume,
}
date_bars[d] = b
dollar_vols.append(bar.close * bar.volume)
if dollar_vols:
last_20 = dollar_vols[-20:]
mean_dvol = sum(last_20) / len(last_20)
else:
mean_dvol = 0.0
return sym, date_bars, mean_dvol
except Exception as sym_exc:
logger.warning(
"snapshot_store_price_fetch_failed",
symbol=sym,
error=str(sym_exc),
)
return sym, {}, 0.0
results = await asyncio.gather(*(_fetch_symbol(sym) for sym in symbols))
for sym, date_bars, mean_dvol in results:
bars_by_symbol[sym] = date_bars
avg_dvol[sym] = mean_dvol
return bars_by_symbol, avg_dvol
except Exception as exc:
logger.warning("snapshot_store_oracle_failed", error=str(exc))
return {}, {}
@staticmethod
async def _fetch_sectors(
symbols: list[str],
oracle_url: str,
concurrency: int = 16,
) -> dict[str, str]:
"""Fetch company sector for each symbol. Default 'UNKNOWN' if unavailable."""
if not symbols:
return {}
result: dict[str, str] = {}
semaphore = asyncio.Semaphore(concurrency)
try:
from libs.oracle_client import CompanyService, OracleClient
async with OracleClient(base_url=oracle_url) as client:
company_svc = CompanyService(client)
async def _fetch_sector(sym: str) -> tuple[str, str]:
async with semaphore:
try:
info = await company_svc.get_company(sym)
return sym, info.sector or "UNKNOWN"
except Exception:
return sym, "UNKNOWN"
sector_results = await asyncio.gather(*(_fetch_sector(sym) for sym in symbols))
for sym, sector in sector_results:
result[sym] = sector
except Exception as exc:
logger.warning("snapshot_store_sector_fetch_failed", error=str(exc))
# Default all remaining to UNKNOWN
for sym in symbols:
result.setdefault(sym, "UNKNOWN")
return result
@staticmethod
async def _fetch_macro(
date_range: tuple[dt.date, dt.date] | None,
db_dsn: str,
) -> dict[dt.date, dict[str, Any]]:
"""Load MacroObservation regime data from DB."""
if date_range is None:
return {}
try:
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from libs.db.models import MacroObservation
engine = create_async_engine(db_dsn, echo=False)
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
result: dict[dt.date, dict[str, Any]] = {}
async with async_session() as session:
stmt = select(MacroObservation).where(
MacroObservation.observation_date >= date_range[0],
MacroObservation.observation_date <= date_range[1],
)
rows = (await session.execute(stmt)).scalars().all()
for obs in rows:
d = obs.observation_date
result.setdefault(d, {})[obs.series_id] = (
float(obs.value) if obs.value is not None else None
)
await engine.dispose()
return result
except Exception as exc:
logger.warning("snapshot_store_macro_fetch_failed", error=str(exc))
return {}
@staticmethod
async def _fetch_spy_macro(
date_range: tuple[dt.date, dt.date] | None,
oracle_url: str,
sma_period: int = 20,
) -> dict[dt.date, dict[str, Any]]:
"""Fetch SPY daily bars and compute SMA for macro regime filtering.
Returns dict: date -> {"spy_close": float, "spy_sma_20": float|None}.
SMA is None for the first (sma_period - 1) bars.
"""
if date_range is None:
return {}
try:
from libs.oracle_client import OracleClient, PriceService
# Extend start date back by sma_period trading days for SMA warm-up
warmup_days = sma_period * 2 # calendar days (conservative buffer)
extended_start = date_range[0] - dt.timedelta(days=warmup_days)
async with OracleClient(base_url=oracle_url) as client:
svc = PriceService(client)
resp = await svc.get_daily_bars(
"SPY", start=extended_start.isoformat(), end=date_range[1].isoformat()
)
# Sort bars by date
sorted_bars = sorted(resp.bars, key=lambda b: b.date)
closes: list[tuple[dt.date, float]] = [
(dt.date.fromisoformat(b.date), float(b.close)) for b in sorted_bars
]
result: dict[dt.date, dict[str, Any]] = {}
for i, (d, close) in enumerate(closes):
sma = None
if i >= sma_period - 1:
window = [c for _, c in closes[i - sma_period + 1 : i + 1]]
sma = sum(window) / len(window)
# Only store data within the actual date range
if d >= date_range[0]:
result[d] = {"spy_close": close, "spy_sma_20": sma}
logger.info(
"snapshot_store_spy_macro_loaded",
bars=len(closes),
dates_with_sma=sum(1 for v in result.values() if v.get("spy_sma_20") is not None),
)
return result
except Exception as exc:
logger.warning("snapshot_store_spy_macro_failed", error=str(exc))
return {}
@staticmethod
def _compute_date_range(
rows: list[dict[str, Any]],
) -> tuple[dt.date, dt.date] | None:
"""Compute (min_date, max_date) from execution and reaction-date columns."""
dates: list[dt.date] = []
for r in rows:
for raw in (
r.get("entry_date"),
r.get("execution_date"),
r.get("reaction_date"),
):
if raw is None:
continue
if isinstance(raw, str):
try:
dates.append(dt.date.fromisoformat(raw))
except ValueError:
pass
elif isinstance(raw, dt.date):
dates.append(raw)
if not dates:
return None
return min(dates), max(dates)
@staticmethod
def _build_reaction_index(
candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]],
) -> dict[dt.date, list[dict[str, Any]]]:
reaction_index: dict[dt.date, list[dict[str, Any]]] = {}
for rows in candidates_by_exec_date.values():
for row in rows:
reaction_date = SnapshotStore._normalize_date(row.get("reaction_date"))
if reaction_date is None:
continue
reaction_index.setdefault(reaction_date, []).append(row)
return reaction_index
@staticmethod
def _normalize_date(raw: Any) -> dt.date | None:
if isinstance(raw, dt.datetime):
return raw.date()
if isinstance(raw, dt.date):
return raw
if isinstance(raw, str):
try:
return dt.date.fromisoformat(raw)
except ValueError:
return None
return None