refactor: class-based alpha factory + month-partitioned data pipeline
Replace the old signal/strategy/backtest modules with a decoupled
data → alpha → combo pipeline (parquet between phases, .pq extension).
Alphas:
- BaseAlpha + @register_alpha factory/plugin registry; one file per
built-in (reversal, reversal_vol, momentum); external alphas via
--alpha-module. Alphas are z-scored position weights, not predictors.
Data:
- baostock primary / akshare fallback, treated consistently.
- New --universe all (~5000 A-shares via query_all_stock, filtered).
- login-once batch downloader; empty-string OHLCV coerced to NaN.
- Month-partitioned dataset {output_dir}/{universe}/month=YYYY-MM/*.pq
with chunked durability flushes; --data-path is the dataset dir.
CLI logs at INFO by default (--log-level) so progress is visible.
Docs (README, CLAUDE.md) updated incl. pipeline diagram and roadmap
TODOs for portfolio construction / backtest / paper trading.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Yuxuan Yan
@@ -0,0 +1,44 @@
|
||||
"""CLI for data download phase."""
|
||||
|
||||
import click
|
||||
from datetime import date
|
||||
|
||||
from pipeline.data.downloader import download_universe
|
||||
|
||||
|
||||
@click.group(name="data")
|
||||
def data():
|
||||
"""Download and manage market data."""
|
||||
|
||||
|
||||
@data.command("download")
|
||||
@click.option(
|
||||
"--universe", default="csi500",
|
||||
help="Which universe: hs300, csi500, all (~5000 A-shares), file path, or comma-separated symbols",
|
||||
)
|
||||
@click.option("--start-date", default="2017-01-01", help="Start date YYYY-MM-DD")
|
||||
@click.option("--end-date", default=str(date.today()), help="End date YYYY-MM-DD")
|
||||
@click.option("--output-dir", default="data/daily_bars", help="Root for the partitioned dataset")
|
||||
@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
|
||||
@click.option("--chunk-size", default=300, type=int, help="Symbols per durability flush")
|
||||
@click.option("--adjust", default="qfq", help="Price adjust: qfq, hfq, or none")
|
||||
def download(universe, start_date, end_date, output_dir, symbols, chunk_size, adjust):
|
||||
"""Download daily bars into a month-partitioned parquet dataset.
|
||||
|
||||
Writes ``{output_dir}/{universe}/month=YYYY-MM/*.pq``. Point ``alpha
|
||||
compute --data-path`` at that dataset directory.
|
||||
"""
|
||||
stats = download_universe(
|
||||
universe=universe,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
output_dir=output_dir,
|
||||
max_symbols=symbols,
|
||||
chunk_size=chunk_size,
|
||||
adjust=adjust,
|
||||
)
|
||||
click.echo(
|
||||
f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
|
||||
f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}"
|
||||
)
|
||||
click.echo(f"Dataset: {stats['dataset_path']}")
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Download daily bar data for a universe and save as a partitioned parquet dataset."""
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.dataset as pads
|
||||
|
||||
# Reuse existing downloader and universe modules
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from data.downloader import download_daily_batch
|
||||
from data.universe import get_all_stocks, get_hs300_stocks, get_zz500_stocks
|
||||
from pipeline.common.schema import DATA_COLUMNS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _fix_baostock_columns(df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""baostock constituent queries return (update_date, code, name) —
|
||||
detect columns by value patterns rather than assuming column order."""
|
||||
cols = df.columns.tolist()
|
||||
result = {}
|
||||
for col in cols:
|
||||
vals = df[col].astype(str)
|
||||
# Stock code: matches sh.NNNNNN or sz.NNNNNN (possibly with dot)
|
||||
if vals.str.match(r"^(sh|sz)\.?\d{6}$").all():
|
||||
result["symbol_id"] = df[col].str.replace(".", "", regex=False)
|
||||
# Stock name: Chinese characters (detected by byte length > str length)
|
||||
elif vals.apply(lambda x: len(x.encode("utf-8")) > len(x)).any() and vals.str.len().max() < 10:
|
||||
result["symbol_name"] = df[col]
|
||||
# Skip date column
|
||||
return pd.DataFrame(result)
|
||||
|
||||
|
||||
def _resolve_universe(universe: str, max_symbols: int = 0) -> pd.DataFrame:
|
||||
"""Resolve a universe name or file path to symbol list with names.
|
||||
|
||||
Returns DataFrame with columns: code (symbol_id), name (symbol_name).
|
||||
"""
|
||||
name = universe.lower()
|
||||
if name == "hs300":
|
||||
df = get_hs300_stocks()
|
||||
# baostock returns (date, code, name) — detect columns by value patterns
|
||||
df = _fix_baostock_columns(df)
|
||||
elif name == "csi500":
|
||||
df = get_zz500_stocks()
|
||||
df = _fix_baostock_columns(df)
|
||||
elif name in ("all", "full"):
|
||||
# Every listed A-share (~5000); already (code, name) with prefixed codes.
|
||||
all_df = get_all_stocks()
|
||||
df = all_df.rename(columns={"code": "symbol_id", "name": "symbol_name"})
|
||||
elif Path(universe).exists():
|
||||
# File with one symbol_id per line
|
||||
with open(universe) as f:
|
||||
symbols = [line.strip() for line in f if line.strip()]
|
||||
df = pd.DataFrame({"symbol_id": symbols, "symbol_name": symbols})
|
||||
else:
|
||||
# Assume comma-separated list
|
||||
symbols = [s.strip() for s in universe.split(",") if s.strip()]
|
||||
df = pd.DataFrame({"symbol_id": symbols, "symbol_name": symbols})
|
||||
|
||||
if max_symbols and max_symbols > 0 and len(df) > max_symbols:
|
||||
df = df.head(max_symbols).copy()
|
||||
|
||||
return df
|
||||
|
||||
|
||||
def _write_month_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: str) -> None:
|
||||
"""Append rows to a Hive-partitioned (month=YYYY-MM) parquet dataset.
|
||||
|
||||
``existing_data_behavior='overwrite_or_ignore'`` plus a per-chunk
|
||||
``basename_prefix`` means each flush adds new ``.pq`` files into the month
|
||||
directories without deleting earlier chunks' files.
|
||||
"""
|
||||
out = df.copy()
|
||||
out["month"] = pd.to_datetime(out["date"]).dt.strftime("%Y-%m")
|
||||
table = pa.Table.from_pandas(out, preserve_index=False)
|
||||
pads.write_dataset(
|
||||
table,
|
||||
str(base_dir),
|
||||
format="parquet",
|
||||
partitioning=["month"],
|
||||
partitioning_flavor="hive",
|
||||
basename_template=f"{basename_prefix}-{{i}}.pq",
|
||||
existing_data_behavior="overwrite_or_ignore",
|
||||
)
|
||||
|
||||
|
||||
def download_universe(
|
||||
universe: str = "csi500",
|
||||
start_date: str = "2017-01-01",
|
||||
end_date: str = "2026-12-31",
|
||||
output_dir: str = "data/daily_bars",
|
||||
max_symbols: int = 0,
|
||||
chunk_size: int = 300,
|
||||
adjust: str = "qfq",
|
||||
) -> dict:
|
||||
"""Download a universe's daily bars into a month-partitioned parquet dataset.
|
||||
|
||||
Streams downloads under a single baostock session and flushes every
|
||||
``chunk_size`` symbols, so memory stays bounded and a crash keeps the
|
||||
partitions already written. The dataset is rebuilt from scratch: any
|
||||
existing ``output_dir/{universe}`` directory is removed first.
|
||||
|
||||
Args:
|
||||
universe: ``hs300``, ``csi500``, ``all``/``full``, a file path, or a
|
||||
comma-separated symbol list.
|
||||
start_date, end_date: ``YYYY-MM-DD`` bounds.
|
||||
output_dir: Root under which ``{universe}/month=YYYY-MM/*.pq`` is written.
|
||||
max_symbols: Cap on symbols (0 = all).
|
||||
chunk_size: Symbols per durability flush.
|
||||
adjust: ``qfq`` / ``hfq`` / ``''``.
|
||||
|
||||
Returns:
|
||||
Stats dict: ``dataset_path``, ``n_symbols`` (succeeded), ``n_requested``,
|
||||
``n_rows``, ``date_min``, ``date_max``.
|
||||
"""
|
||||
constituents = _resolve_universe(universe, max_symbols)
|
||||
symbols = constituents["symbol_id"].tolist()
|
||||
names = dict(zip(constituents["symbol_id"], constituents["symbol_name"]))
|
||||
n_requested = len(symbols)
|
||||
logger.info("Universe %s: %d symbols, %s → %s", universe, n_requested, start_date, end_date)
|
||||
|
||||
base_dir = Path(output_dir) / universe
|
||||
if base_dir.exists():
|
||||
shutil.rmtree(base_dir)
|
||||
base_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
buffer: list[pd.DataFrame] = []
|
||||
chunk_idx = 0
|
||||
succeeded = 0
|
||||
n_rows = 0
|
||||
date_min = None
|
||||
date_max = None
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal buffer, chunk_idx, n_rows, date_min, date_max
|
||||
if not buffer:
|
||||
return
|
||||
chunk = pd.concat(buffer, ignore_index=True)
|
||||
_write_month_partitions(chunk, base_dir, basename_prefix=f"chunk{chunk_idx:04d}")
|
||||
n_rows += len(chunk)
|
||||
cmin, cmax = chunk["date"].min(), chunk["date"].max()
|
||||
date_min = cmin if date_min is None else min(date_min, cmin)
|
||||
date_max = cmax if date_max is None else max(date_max, cmax)
|
||||
logger.info("Flushed chunk %d: %d rows (%d symbols done)", chunk_idx, len(chunk), succeeded)
|
||||
buffer = []
|
||||
chunk_idx += 1
|
||||
|
||||
for i, (symbol, df) in enumerate(
|
||||
download_daily_batch(symbols, start_date, end_date, adjust=adjust), start=1
|
||||
):
|
||||
if df is None:
|
||||
logger.warning(" %s: no data", symbol)
|
||||
else:
|
||||
df["symbol_id"] = symbol
|
||||
df["symbol_name"] = names.get(symbol, symbol)
|
||||
buffer.append(df[DATA_COLUMNS])
|
||||
succeeded += 1
|
||||
if len(buffer) >= chunk_size:
|
||||
flush()
|
||||
if i % 100 == 0:
|
||||
logger.info("Progress: %d/%d symbols", i, n_requested)
|
||||
flush()
|
||||
|
||||
if succeeded == 0:
|
||||
raise RuntimeError("No data downloaded for any symbol")
|
||||
|
||||
return {
|
||||
"dataset_path": str(base_dir),
|
||||
"n_symbols": succeeded,
|
||||
"n_requested": n_requested,
|
||||
"n_rows": n_rows,
|
||||
"date_min": None if date_min is None else str(pd.Timestamp(date_min).date()),
|
||||
"date_max": None if date_max is None else str(pd.Timestamp(date_max).date()),
|
||||
}
|
||||
Reference in New Issue
Block a user