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
+80
-28
@@ -1,14 +1,16 @@
|
||||
"""Unified data downloader: akshare primary, baostock fallback."""
|
||||
"""Unified data downloader: baostock primary, akshare fallback."""
|
||||
import logging
|
||||
from datetime import date, datetime
|
||||
from typing import Optional
|
||||
from typing import Iterable, Iterator, Optional, Tuple
|
||||
import pandas as pd
|
||||
import akshare as ak
|
||||
import baostock as bs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
BAOSTOCK_FREQ_MAP = {"d": "d", "w": "w", "m": "m"} # baostock only supports daily
|
||||
# Map the adjust argument to baostock's adjustflag codes.
|
||||
_BAOSTOCK_ADJUST = {"qfq": "2", "hfq": "1", "": "3", "none": "3"}
|
||||
_BAOSTOCK_FIELDS = "date,open,high,low,close,volume,amount"
|
||||
_OHLCV = ["open", "high", "low", "close", "volume", "amount"]
|
||||
|
||||
|
||||
def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
|
||||
@@ -44,8 +46,8 @@ def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") ->
|
||||
return None
|
||||
|
||||
|
||||
def _download_baostock(symbol: str, start: str, end: str, frequency: str = "d") -> Optional[pd.DataFrame]:
|
||||
"""Download daily bars from baostock as fallback."""
|
||||
def _download_baostock(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
|
||||
"""Download daily bars from baostock (primary source)."""
|
||||
try:
|
||||
bs.login()
|
||||
# baostock format: sh.600000
|
||||
@@ -55,8 +57,8 @@ def _download_baostock(symbol: str, start: str, end: str, frequency: str = "d")
|
||||
fields="date,open,high,low,close,volume,amount",
|
||||
start_date=start,
|
||||
end_date=end,
|
||||
frequency=frequency,
|
||||
adjustflag="2", # qfq
|
||||
frequency="d",
|
||||
adjustflag=_BAOSTOCK_ADJUST.get(adjust, "2"),
|
||||
)
|
||||
if rs.error_code != "0":
|
||||
logger.warning(f"baostock error for {symbol}: {rs.error_msg}")
|
||||
@@ -64,22 +66,22 @@ def _download_baostock(symbol: str, start: str, end: str, frequency: str = "d")
|
||||
data_list = []
|
||||
while rs.next():
|
||||
data_list.append(rs.get_row_data())
|
||||
bs.logout()
|
||||
if not data_list:
|
||||
return None
|
||||
df = pd.DataFrame(data_list, columns=["date", "open", "high", "low", "close", "volume", "amount"])
|
||||
df[["open", "high", "low", "close", "volume", "amount"]] = df[
|
||||
["open", "high", "low", "close", "volume", "amount"]
|
||||
].astype(float)
|
||||
].apply(pd.to_numeric, errors="coerce")
|
||||
df["symbol"] = symbol
|
||||
return df[["symbol", "date", "open", "high", "low", "close", "volume", "amount"]]
|
||||
except Exception as e:
|
||||
logger.warning(f"baostock download failed for {symbol}: {e}")
|
||||
return None
|
||||
finally:
|
||||
try:
|
||||
bs.logout()
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def download_daily(
|
||||
@@ -90,24 +92,24 @@ def download_daily(
|
||||
source: str = "auto",
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Download daily OHLCV data. Tries akshare first, falls back to baostock.
|
||||
Download daily OHLCV data. Tries baostock first, falls back to akshare.
|
||||
|
||||
Args:
|
||||
symbol: Stock symbol like 'sh600000' or 'sz000001'
|
||||
start: Start date 'YYYY-MM-DD'
|
||||
end: End date 'YYYY-MM-DD'
|
||||
adjust: 'qfq' (forward-adjusted), 'hfq' (backward), '' (none)
|
||||
source: 'auto' (akshare then baostock fallback), 'akshare' only,
|
||||
or 'baostock' only
|
||||
source: 'auto' (baostock then akshare fallback), 'baostock' only,
|
||||
or 'akshare' only
|
||||
|
||||
Returns:
|
||||
DataFrame with columns: symbol, date, open, high, low, close, volume, amount
|
||||
"""
|
||||
df = None
|
||||
if source in ("akshare", "auto"):
|
||||
if source in ("baostock", "auto"):
|
||||
df = _download_baostock(symbol, start, end, adjust)
|
||||
if df is None and source in ("akshare", "auto"):
|
||||
df = _download_akshare(symbol, start, end, adjust)
|
||||
if df is None and source in ("baostock", "auto"):
|
||||
df = _download_baostock(symbol, start, end)
|
||||
|
||||
if df is None or df.empty:
|
||||
raise RuntimeError(f"Failed to download data for {symbol} from {start} to {end}")
|
||||
@@ -117,18 +119,68 @@ def download_daily(
|
||||
return df
|
||||
|
||||
|
||||
def download_batch(
|
||||
symbols: list[str],
|
||||
def download_daily_batch(
|
||||
symbols: Iterable[str],
|
||||
start: str,
|
||||
end: str,
|
||||
adjust: str = "qfq",
|
||||
) -> dict[str, pd.DataFrame]:
|
||||
"""Download daily data for multiple symbols. Returns {symbol: DataFrame}."""
|
||||
results = {}
|
||||
for sym in symbols:
|
||||
akshare_fallback: bool = True,
|
||||
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
|
||||
"""Download many symbols under a single baostock session.
|
||||
|
||||
Logging into baostock once per call (instead of per symbol) is the dominant
|
||||
speed-up when fetching thousands of symbols. Yields ``(symbol, df)`` as each
|
||||
symbol completes so callers can stream results to disk; ``df`` is ``None``
|
||||
when both sources fail. Each ``df`` has the same 8 columns as
|
||||
:func:`download_daily`.
|
||||
|
||||
Args:
|
||||
symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
|
||||
start, end: ``YYYY-MM-DD`` bounds.
|
||||
adjust: ``qfq`` / ``hfq`` / ``''``.
|
||||
akshare_fallback: Retry a failed symbol through akshare before yielding
|
||||
``None``.
|
||||
"""
|
||||
flag = _BAOSTOCK_ADJUST.get(adjust, "2")
|
||||
bs.login()
|
||||
try:
|
||||
for symbol in symbols:
|
||||
df: Optional[pd.DataFrame] = None
|
||||
try:
|
||||
code = f"{symbol[:2]}.{symbol[2:]}"
|
||||
rs = bs.query_history_k_data_plus(
|
||||
code=code, fields=_BAOSTOCK_FIELDS,
|
||||
start_date=start, end_date=end,
|
||||
frequency="d", adjustflag=flag,
|
||||
)
|
||||
if rs.error_code == "0":
|
||||
rows = []
|
||||
while rs.next():
|
||||
rows.append(rs.get_row_data())
|
||||
if rows:
|
||||
df = pd.DataFrame(rows, columns=["date", *_OHLCV])
|
||||
# Suspended-trading days come back as empty strings;
|
||||
# coerce to NaN rather than crashing the whole symbol.
|
||||
df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce")
|
||||
df["symbol"] = symbol
|
||||
df = df[["symbol", "date", *_OHLCV]]
|
||||
else:
|
||||
logger.warning("baostock error for %s: %s", symbol, rs.error_msg)
|
||||
except Exception as e:
|
||||
logger.warning("baostock download failed for %s: %s", symbol, e)
|
||||
|
||||
if (df is None or df.empty) and akshare_fallback:
|
||||
df = _download_akshare(symbol, start, end, adjust)
|
||||
|
||||
if df is not None and not df.empty:
|
||||
df["date"] = pd.to_datetime(df["date"])
|
||||
df = df.sort_values("date").reset_index(drop=True)
|
||||
yield symbol, df
|
||||
else:
|
||||
yield symbol, None
|
||||
finally:
|
||||
try:
|
||||
results[sym] = download_daily(sym, start, end, adjust)
|
||||
logger.info(f"Downloaded {sym}: {len(results[sym])} bars")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed {sym}: {e}")
|
||||
return results
|
||||
bs.logout()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@@ -1,44 +0,0 @@
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import date
|
||||
from typing import Optional
|
||||
import pandas as pd
|
||||
|
||||
|
||||
@dataclass
|
||||
class DailyBar:
|
||||
"""Single daily bar for one stock."""
|
||||
symbol: str
|
||||
date: date
|
||||
open: float
|
||||
high: float
|
||||
low: float
|
||||
close: float
|
||||
volume: float
|
||||
amount: float # turnover in yuan
|
||||
|
||||
@classmethod
|
||||
def from_dataframe(cls, df: pd.DataFrame, symbol_col: str = "symbol") -> list["DailyBar"]:
|
||||
"""Convert akshare/baostock DataFrame to list of DailyBar."""
|
||||
bars = []
|
||||
for _, row in df.iterrows():
|
||||
bars.append(cls(
|
||||
symbol=row.get(symbol_col, ""),
|
||||
date=pd.Timestamp(row["date"]).date(),
|
||||
open=float(row["open"]),
|
||||
high=float(row["high"]),
|
||||
low=float(row["low"]),
|
||||
close=float(row["close"]),
|
||||
volume=float(row["volume"]),
|
||||
amount=float(row.get("amount", 0)),
|
||||
))
|
||||
return bars
|
||||
|
||||
def to_series(self) -> dict:
|
||||
return {
|
||||
"date": self.date,
|
||||
"open": self.open,
|
||||
"high": self.high,
|
||||
"low": self.low,
|
||||
"close": self.close,
|
||||
"volume": self.volume,
|
||||
}
|
||||
+47
-26
@@ -1,36 +1,16 @@
|
||||
"""CSI 300 (HS300) and CSI 500 (ZZ500) universe helpers."""
|
||||
"""CSI 300 (HS300), CSI 500 (ZZ500), and full A-share universe helpers."""
|
||||
import logging
|
||||
from datetime import date, timedelta
|
||||
|
||||
import baostock as bs
|
||||
import pandas as pd
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# First 30 HS300 constituents (large caps) in 'shXXXXXX' / 'szXXXXXX' format.
|
||||
# Hardcoded for fast, deterministic smoke tests. Use get_hs300_stocks() for the
|
||||
# live, full list — downloading daily bars for all ~300 takes roughly 10 minutes.
|
||||
SYMBOLS = [
|
||||
"sh600000", "sh600009", "sh600010", "sh600028", "sh600030",
|
||||
"sh600036", "sh600048", "sh600050", "sh600104", "sh600276",
|
||||
"sh600309", "sh600519", "sh600585", "sh600887", "sh600900",
|
||||
"sh601012", "sh601166", "sh601288", "sh601318", "sh601398",
|
||||
"sh601628", "sh601668", "sh601857", "sh601888", "sh601988",
|
||||
"sz000001", "sz000002", "sz000333", "sz000651", "sz000858",
|
||||
]
|
||||
|
||||
|
||||
# First 30 CSI 500 (ZZ500) constituents (mid/small caps) in 'shXXXXXX' /
|
||||
# 'szXXXXXX' format. Hardcoded for fast, deterministic smoke tests. Use
|
||||
# get_zz500_stocks() for the live, full list. Mean reversion tends to be
|
||||
# stronger in these smaller caps than in the HS300 large caps.
|
||||
CSI500_SYMBOLS = [
|
||||
"sh600006", "sh600008", "sh600017", "sh600020", "sh600021",
|
||||
"sh600026", "sh600037", "sh600039", "sh600053", "sh600056",
|
||||
"sh600060", "sh600061", "sh600062", "sh600073", "sh600089",
|
||||
"sh600095", "sh600118", "sh600125", "sh600126", "sh600143",
|
||||
"sh600153", "sh600160", "sh600169", "sh600176", "sh600183",
|
||||
"sz000009", "sz000012", "sz000021", "sz000025", "sz000027",
|
||||
]
|
||||
# A-share code patterns (baostock dotted form): SH main/STAR (sh.6xxxxx),
|
||||
# SZ main/SME (sz.0xxxxx), ChiNext (sz.3xxxxx). Excludes indices and B-shares.
|
||||
_ASHARE_RE = r"^sh\.6\d{5}$|^sz\.[03]\d{5}$"
|
||||
_SZ_INDEX_RE = r"^sz\.399"
|
||||
|
||||
|
||||
def get_hs300_stocks() -> pd.DataFrame:
|
||||
@@ -69,3 +49,44 @@ def get_zz500_stocks() -> pd.DataFrame:
|
||||
df = pd.DataFrame(stocks, columns=["code", "name", "date"])
|
||||
df["code"] = df["code"].str.replace(".", "", regex=False)
|
||||
return df
|
||||
|
||||
|
||||
def get_all_stocks(day: str = "") -> pd.DataFrame:
|
||||
"""Fetch every listed A-share from baostock's all-stock snapshot.
|
||||
|
||||
Queries ``query_all_stock`` for a single trading day and keeps only A-shares
|
||||
(SH main/STAR, SZ main/SME/ChiNext), dropping indices and B-shares. If the
|
||||
given day is a non-trading day baostock returns nothing, so we walk back up
|
||||
to 10 days to land on the most recent trading day.
|
||||
|
||||
Args:
|
||||
day: ``YYYY-MM-DD`` snapshot day; defaults to today (walks back to the
|
||||
last trading day).
|
||||
|
||||
Returns:
|
||||
DataFrame with columns ``code`` (e.g. ``sh600000``), ``name``.
|
||||
"""
|
||||
start = date.fromisoformat(day) if day else date.today()
|
||||
bs.login()
|
||||
try:
|
||||
rows: list = []
|
||||
fields: list = []
|
||||
for back in range(11):
|
||||
probe = (start - timedelta(days=back)).isoformat()
|
||||
rs = bs.query_all_stock(day=probe)
|
||||
fields = rs.fields
|
||||
while rs.next():
|
||||
rows.append(rs.get_row_data())
|
||||
if rows:
|
||||
logger.info("query_all_stock: %d rows on %s", len(rows), probe)
|
||||
break
|
||||
finally:
|
||||
bs.logout()
|
||||
|
||||
df = pd.DataFrame(rows, columns=fields)
|
||||
code = df["code"]
|
||||
keep = code.str.match(_ASHARE_RE) & ~code.str.match(_SZ_INDEX_RE)
|
||||
df = df[keep].copy()
|
||||
df["code"] = df["code"].str.replace(".", "", regex=False)
|
||||
df = df.rename(columns={"code_name": "name"})
|
||||
return df[["code", "name"]].reset_index(drop=True)
|
||||
|
||||
Reference in New Issue
Block a user