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:
@@ -0,0 +1,57 @@
|
||||
"""Base class for alphas.
|
||||
|
||||
An alpha maps a wide close matrix (date index × symbol_id columns) to signed
|
||||
position weights. Subclasses implement :meth:`signal` — the raw, unnormalized
|
||||
score. The base class turns a signal into cross-sectionally z-scored weights
|
||||
via :meth:`to_weights` (override it for a different normalization).
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class BaseAlpha(ABC):
|
||||
"""A position-weight alpha over a cross-section of stocks.
|
||||
|
||||
Concrete subclasses must set a unique class-level :attr:`name` (the registry
|
||||
key) and implement :meth:`signal`. Construct subclasses with their own typed
|
||||
parameters (e.g. ``lookback``); the factory passes only the parameters a
|
||||
given ``__init__`` accepts.
|
||||
"""
|
||||
|
||||
#: Unique registry key. Every concrete alpha must set this to a non-empty str.
|
||||
name: str = ""
|
||||
|
||||
@abstractmethod
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Compute the raw signal.
|
||||
|
||||
Args:
|
||||
close: Wide close prices, date index × ``symbol_id`` columns.
|
||||
|
||||
Returns:
|
||||
A wide DataFrame aligned to ``close`` where higher values indicate a
|
||||
stronger long. Use NaN where the signal is undefined.
|
||||
"""
|
||||
|
||||
def to_weights(self, signal: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Cross-sectionally z-score a signal into signed position weights.
|
||||
|
||||
Each date is demeaned and scaled by its cross-sectional std; undefined
|
||||
cells become a 0 weight. Override for a custom scheme (rank, neutralized,
|
||||
capped, etc.).
|
||||
"""
|
||||
signal = signal.dropna(how="all")
|
||||
demeaned = signal.subtract(signal.mean(axis=1), axis=0)
|
||||
std = signal.std(axis=1).replace(0, np.nan)
|
||||
weights = demeaned.divide(std, axis=0)
|
||||
return weights.fillna(0.0)
|
||||
|
||||
def weights(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Full pipeline for one alpha: raw signal → normalized weights."""
|
||||
return self.to_weights(self.signal(close))
|
||||
|
||||
def __repr__(self) -> str:
|
||||
params = ", ".join(f"{k}={v!r}" for k, v in vars(self).items())
|
||||
return f"{type(self).__name__}({params})"
|
||||
@@ -0,0 +1,174 @@
|
||||
"""CLI for alpha computation and evaluation."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import click
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.compute import compute_alpha, evaluate_alpha
|
||||
from pipeline.alpha.registry import available_alphas, load_alpha_module
|
||||
|
||||
|
||||
@click.group(name="alpha")
|
||||
def alpha():
|
||||
"""Compute and evaluate alpha weights."""
|
||||
|
||||
|
||||
def _coerce(value: str):
|
||||
"""Best-effort coercion of a CLI string to int, then float, else str."""
|
||||
for cast in (int, float):
|
||||
try:
|
||||
return cast(value)
|
||||
except ValueError:
|
||||
continue
|
||||
return value
|
||||
|
||||
|
||||
def _parse_params(pairs: tuple[str, ...]) -> dict:
|
||||
"""Parse repeated ``name=value`` options into a params dict."""
|
||||
params: dict = {}
|
||||
for pair in pairs:
|
||||
if "=" not in pair:
|
||||
raise click.BadParameter(f"--param must be name=value, got '{pair}'")
|
||||
key, value = pair.split("=", 1)
|
||||
params[key.strip()] = _coerce(value.strip())
|
||||
return params
|
||||
|
||||
|
||||
@alpha.command("list")
|
||||
@click.option(
|
||||
"--alpha-module", "alpha_modules", multiple=True,
|
||||
help="External module(s) to import first (dotted path or .py file)",
|
||||
)
|
||||
def list_(alpha_modules):
|
||||
"""List the registered alpha types."""
|
||||
for spec in alpha_modules:
|
||||
load_alpha_module(spec)
|
||||
for name in available_alphas():
|
||||
click.echo(name)
|
||||
|
||||
|
||||
@alpha.command("compute")
|
||||
@click.option("--data-path", required=True, help="Path to data parquet file")
|
||||
@click.option("--alpha-name", required=True, help="Name for this alpha")
|
||||
@click.option("--alpha-type", required=True, help="Registry key of the alpha class")
|
||||
@click.option("--output-dir", default="alphas", help="Directory to save alpha parquet")
|
||||
@click.option("--lookback", default=5, type=int, help="Lookback days")
|
||||
@click.option("--vol-window", default=20, type=int, help="Volatility window (reversal_vol only)")
|
||||
@click.option(
|
||||
"--alpha-module", "alpha_modules", multiple=True,
|
||||
help="External module(s) to import so their alphas register (dotted path or .py file)",
|
||||
)
|
||||
@click.option(
|
||||
"--param", "extra_params", multiple=True,
|
||||
help="Extra alpha constructor param as name=value (repeatable)",
|
||||
)
|
||||
def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window,
|
||||
alpha_modules, extra_params):
|
||||
"""Compute one alpha from raw data and save as parquet."""
|
||||
for spec in alpha_modules:
|
||||
load_alpha_module(spec)
|
||||
|
||||
options = available_alphas()
|
||||
if alpha_type not in options:
|
||||
raise click.BadParameter(
|
||||
f"Unknown alpha-type '{alpha_type}'. Available: {options}. "
|
||||
f"Use --alpha-module to register an external alpha.",
|
||||
param_hint="--alpha-type",
|
||||
)
|
||||
|
||||
params = {"lookback": lookback, "vol_window": vol_window}
|
||||
params.update(_parse_params(extra_params))
|
||||
|
||||
data = pd.read_parquet(data_path)
|
||||
click.echo(f"Loaded data: {len(data):,} rows from {data_path}")
|
||||
|
||||
result = compute_alpha(
|
||||
data=data,
|
||||
alpha_name=alpha_name,
|
||||
alpha_type=alpha_type,
|
||||
**params,
|
||||
)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
out_path = f"{output_dir}/{alpha_name}.pq"
|
||||
result.to_parquet(out_path, index=False)
|
||||
click.echo(f"Saved alpha: {out_path} ({len(result):,} rows)")
|
||||
click.echo(
|
||||
f"Weight stats — min: {result['weight'].min():.4f}, "
|
||||
f"max: {result['weight'].max():.4f}, "
|
||||
f"mean: {result['weight'].mean():.4f}"
|
||||
)
|
||||
|
||||
|
||||
@alpha.command("reversal")
|
||||
@click.option("--data-path", required=True, help="Path to data parquet file")
|
||||
@click.option("--output-dir", default="alphas", help="Directory to save alpha parquet")
|
||||
@click.option("--lookback", default=5, type=int, help="Lookback days")
|
||||
def reversal(data_path, output_dir, lookback):
|
||||
"""Shortcut: compute a reversal alpha."""
|
||||
alpha_name = f"reversal_{lookback}d"
|
||||
ctx = click.get_current_context()
|
||||
ctx.invoke(
|
||||
compute,
|
||||
data_path=data_path,
|
||||
alpha_name=alpha_name,
|
||||
alpha_type="reversal",
|
||||
output_dir=output_dir,
|
||||
lookback=lookback,
|
||||
)
|
||||
|
||||
|
||||
@alpha.command("reversal-vol")
|
||||
@click.option("--data-path", required=True, help="Path to data parquet file")
|
||||
@click.option("--output-dir", default="alphas", help="Directory to save alpha parquet")
|
||||
@click.option("--lookback", default=5, type=int, help="Lookback days")
|
||||
@click.option("--vol-window", default=20, type=int, help="Volatility window")
|
||||
def reversal_vol(data_path, output_dir, lookback, vol_window):
|
||||
"""Shortcut: compute a volatility-scaled reversal alpha."""
|
||||
alpha_name = f"reversal_vol_{lookback}d_{vol_window}d"
|
||||
ctx = click.get_current_context()
|
||||
ctx.invoke(
|
||||
compute,
|
||||
data_path=data_path,
|
||||
alpha_name=alpha_name,
|
||||
alpha_type="reversal_vol",
|
||||
output_dir=output_dir,
|
||||
lookback=lookback,
|
||||
vol_window=vol_window,
|
||||
)
|
||||
|
||||
|
||||
@alpha.command("eval")
|
||||
@click.option("--alpha-path", required=True, help="Path to alpha parquet file")
|
||||
@click.option("--data-path", required=True, help="Path to data parquet (for price data)")
|
||||
def eval_(alpha_path, data_path):
|
||||
"""Evaluate an alpha's performance (return, Sharpe, turnover).
|
||||
|
||||
Alphas are interpreted as position WEIGHTS, not return predictors.
|
||||
No IC/IR metrics — these are not predictors of future returns.
|
||||
"""
|
||||
alpha_df = pd.read_parquet(alpha_path)
|
||||
data_df = pd.read_parquet(data_path)
|
||||
|
||||
metrics = evaluate_alpha(alpha_df, data_df)
|
||||
|
||||
click.echo("\n" + "=" * 50)
|
||||
click.echo("ALPHA EVALUATION")
|
||||
click.echo("=" * 50)
|
||||
click.echo(f"Cumulative Return: {metrics['cumulative_return']:>10.4%}")
|
||||
click.echo(f"Annual Sharpe: {metrics['sharpe_annual']:>10.4f}")
|
||||
click.echo(f"Annual Turnover: {metrics['turnover_annual']:>10.2%}")
|
||||
click.echo(f"Max Drawdown: {metrics['max_drawdown']:>10.4%}")
|
||||
click.echo(f"Hit Rate: {metrics['hit_rate']:>10.2%}")
|
||||
click.echo(f"Trading Days: {metrics['n_dates']:>10d}")
|
||||
click.echo("=" * 50)
|
||||
|
||||
# Also dump JSON
|
||||
os.makedirs("reports", exist_ok=True)
|
||||
alpha_name = alpha_df["alpha_name"].iloc[0]
|
||||
json_path = f"reports/{alpha_name}_eval.json"
|
||||
with open(json_path, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
click.echo(f"\nReport saved: {json_path}")
|
||||
@@ -0,0 +1,153 @@
|
||||
"""Alpha computation and evaluation.
|
||||
|
||||
Alphas are position WEIGHTS — positive=long, negative=short. They are NOT
|
||||
predictors of future returns. Concrete alphas are classes that live in
|
||||
``pipeline/alpha/library/`` (or any external module) and are resolved by name
|
||||
through :mod:`pipeline.alpha.registry`.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.registry import get_alpha
|
||||
from pipeline.common.schema import ALPHA_COLUMNS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _pivot_close(df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Pivot data to wide format: date index, columns = symbol_id, values = close."""
|
||||
pivot = df.pivot_table(
|
||||
index="date", columns="symbol_id", values="close", aggfunc="first"
|
||||
)
|
||||
return pivot.sort_index()
|
||||
|
||||
|
||||
def _daily_returns(close: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Compute daily returns from wide close DataFrame."""
|
||||
return close.pct_change()
|
||||
|
||||
|
||||
def compute_alpha(
|
||||
data: pd.DataFrame,
|
||||
alpha_name: str,
|
||||
alpha_type: str,
|
||||
**params,
|
||||
) -> pd.DataFrame:
|
||||
"""Compute alpha weights from raw data.
|
||||
|
||||
Args:
|
||||
data: DataFrame with DATA_COLUMNS.
|
||||
alpha_name: Label stored in the ``alpha_name`` output column.
|
||||
alpha_type: Registry key of the alpha class (e.g. ``reversal``).
|
||||
**params: Constructor parameters for the alpha (e.g. ``lookback``,
|
||||
``vol_window``). Only the params the alpha's ``__init__`` accepts are
|
||||
used; extras are ignored.
|
||||
|
||||
Returns:
|
||||
DataFrame with ALPHA_COLUMNS.
|
||||
|
||||
Raises:
|
||||
KeyError: If ``alpha_type`` is not registered.
|
||||
"""
|
||||
alpha = get_alpha(alpha_type, **params)
|
||||
close = _pivot_close(data)
|
||||
weights = alpha.weights(close)
|
||||
|
||||
# Melt to long format
|
||||
weights_melted = weights.reset_index().melt(
|
||||
id_vars="date", var_name="symbol_id", value_name="weight"
|
||||
)
|
||||
weights_melted["alpha_name"] = alpha_name
|
||||
weights_melted = weights_melted[ALPHA_COLUMNS]
|
||||
weights_melted = weights_melted.dropna(subset=["weight"])
|
||||
weights_melted = weights_melted.sort_values(["symbol_id", "date"]).reset_index(drop=True)
|
||||
|
||||
logger.info(
|
||||
"Alpha '%s' (%r): %d symbols × %d dates, weight range [%.4f, %.4f]",
|
||||
alpha_name,
|
||||
alpha,
|
||||
weights_melted["symbol_id"].nunique(),
|
||||
weights_melted["date"].nunique(),
|
||||
weights_melted["weight"].min(),
|
||||
weights_melted["weight"].max(),
|
||||
)
|
||||
return weights_melted
|
||||
|
||||
|
||||
def evaluate_alpha(alpha_df: pd.DataFrame, data_df: pd.DataFrame) -> dict:
|
||||
"""Evaluate an alpha's performance as position weights.
|
||||
|
||||
Computes return, annualized Sharpe, annualized turnover, max drawdown.
|
||||
|
||||
Alpha is interpreted as POSITION WEIGHTS, not predictions.
|
||||
Return on date t = sum(weight[s,t] * realized_return[s,t]) / sum(abs(weight[s,t]))
|
||||
|
||||
Args:
|
||||
alpha_df: DataFrame with ALPHA_COLUMNS.
|
||||
data_df: DataFrame with DATA_COLUMNS (for price data).
|
||||
|
||||
Returns:
|
||||
Dict with metrics: cumulative_return, sharpe_annual, turnover_annual,
|
||||
max_drawdown, hit_rate, n_dates.
|
||||
"""
|
||||
close = _pivot_close(data_df)
|
||||
returns = _daily_returns(close)
|
||||
|
||||
# Pivot alpha weights to wide format
|
||||
weights = alpha_df.pivot_table(
|
||||
index="date", columns="symbol_id", values="weight", aggfunc="first"
|
||||
).sort_index()
|
||||
|
||||
# Align dates
|
||||
common_dates = weights.index.intersection(returns.index)
|
||||
weights = weights.loc[common_dates]
|
||||
returns = returns.loc[common_dates]
|
||||
|
||||
if len(common_dates) < 2:
|
||||
return {
|
||||
"cumulative_return": 0.0,
|
||||
"sharpe_annual": 0.0,
|
||||
"turnover_annual": 0.0,
|
||||
"max_drawdown": 0.0,
|
||||
"hit_rate": 0.0,
|
||||
"n_dates": len(common_dates),
|
||||
}
|
||||
|
||||
# Daily portfolio return = sum(w * r) / sum(|w|) — normalized by gross exposure
|
||||
daily_returns = (weights * returns).sum(axis=1) / weights.abs().sum(axis=1)
|
||||
|
||||
# Cumulative return
|
||||
cumulative_return = float((1.0 + daily_returns).prod() - 1.0)
|
||||
|
||||
# Annualized Sharpe (sqrt(252) * mean / std)
|
||||
mu = daily_returns.mean()
|
||||
sigma = daily_returns.std()
|
||||
sharpe_annual = float(np.sqrt(252) * mu / sigma) if sigma > 0 else 0.0
|
||||
|
||||
# Annualized turnover: avg daily turnover * 252
|
||||
# Daily turnover = sum(|w_t - w_{t-1}|) / sum(|w_{t-1}|)
|
||||
weight_change = weights.diff().abs().sum(axis=1)
|
||||
gross_exposure = weights.abs().sum(axis=1).shift(1)
|
||||
daily_turnover = weight_change / gross_exposure
|
||||
turnover_annual = float(daily_turnover.mean() * 252)
|
||||
|
||||
# Max drawdown
|
||||
equity = (1.0 + daily_returns).cumprod()
|
||||
peak = equity.cummax()
|
||||
drawdown = (equity - peak) / peak
|
||||
max_drawdown = float(drawdown.min())
|
||||
|
||||
# Hit rate
|
||||
hit_rate = float((daily_returns > 0).mean())
|
||||
|
||||
return {
|
||||
"cumulative_return": cumulative_return,
|
||||
"sharpe_annual": sharpe_annual,
|
||||
"turnover_annual": turnover_annual,
|
||||
"max_drawdown": max_drawdown,
|
||||
"hit_rate": hit_rate,
|
||||
"n_dates": len(common_dates),
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
"""Built-in alpha library.
|
||||
|
||||
Importing this package imports each alpha module, which registers the alpha via
|
||||
the ``@register_alpha`` decorator. Add a new built-in by dropping a module here
|
||||
and importing it below.
|
||||
"""
|
||||
from pipeline.alpha.library import momentum, reversal, reversal_vol # noqa: F401
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Short-horizon momentum alpha."""
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
|
||||
@register_alpha
|
||||
class MomentumAlpha(BaseAlpha):
|
||||
"""Positive trailing return: stocks that rose score high."""
|
||||
|
||||
name = "momentum"
|
||||
|
||||
def __init__(self, lookback: int = 5):
|
||||
self.lookback = lookback
|
||||
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
return close.pct_change(self.lookback)
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Short-horizon reversal alpha."""
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
|
||||
@register_alpha
|
||||
class ReversalAlpha(BaseAlpha):
|
||||
"""Negative trailing return: oversold stocks score high."""
|
||||
|
||||
name = "reversal"
|
||||
|
||||
def __init__(self, lookback: int = 5):
|
||||
self.lookback = lookback
|
||||
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
return -close.pct_change(self.lookback)
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Volatility-scaled short-horizon reversal alpha."""
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
|
||||
@register_alpha
|
||||
class ReversalVolAlpha(BaseAlpha):
|
||||
"""Reversal scaled by trailing volatility.
|
||||
|
||||
The raw reversal ``-close.pct_change(lookback)`` is divided by the rolling
|
||||
standard deviation of daily returns over ``vol_window``, so the score favors
|
||||
oversold names whose move is large *relative* to their own volatility.
|
||||
"""
|
||||
|
||||
name = "reversal_vol"
|
||||
|
||||
def __init__(self, lookback: int = 5, vol_window: int = 20):
|
||||
self.lookback = lookback
|
||||
self.vol_window = vol_window
|
||||
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
reversal = -close.pct_change(self.lookback)
|
||||
vol = close.pct_change().rolling(self.vol_window).std()
|
||||
return reversal / vol
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Registry and factory for alphas.
|
||||
|
||||
Built-in alphas live in :mod:`pipeline.alpha.library` and self-register via the
|
||||
:func:`register_alpha` decorator. External alphas authored anywhere can be made
|
||||
available with :func:`load_alpha_module` (a dotted module path or a ``.py`` file),
|
||||
which is how you test an alpha written outside this repo.
|
||||
"""
|
||||
import importlib
|
||||
import importlib.util
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import Optional, Type
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
|
||||
_REGISTRY: dict[str, Type[BaseAlpha]] = {}
|
||||
_builtins_loaded = False
|
||||
|
||||
|
||||
def register_alpha(cls: Type[BaseAlpha]) -> Type[BaseAlpha]:
|
||||
"""Class decorator that registers an alpha under its :attr:`~BaseAlpha.name`.
|
||||
|
||||
Raises:
|
||||
TypeError: If ``cls`` is not a ``BaseAlpha`` subclass.
|
||||
ValueError: If ``name`` is empty or already used by a different class.
|
||||
"""
|
||||
if not (isinstance(cls, type) and issubclass(cls, BaseAlpha)):
|
||||
raise TypeError(f"{cls!r} is not a BaseAlpha subclass")
|
||||
key = getattr(cls, "name", "")
|
||||
if not key:
|
||||
raise ValueError(f"{cls.__name__} must set a non-empty class attribute `name`")
|
||||
existing = _REGISTRY.get(key)
|
||||
if existing is not None and existing is not cls:
|
||||
raise ValueError(
|
||||
f"Alpha name '{key}' already registered by {existing.__name__}"
|
||||
)
|
||||
_REGISTRY[key] = cls
|
||||
return cls
|
||||
|
||||
|
||||
def available_alphas() -> list[str]:
|
||||
"""Sorted names of all registered alphas (built-ins are loaded lazily)."""
|
||||
_ensure_builtins()
|
||||
return sorted(_REGISTRY)
|
||||
|
||||
|
||||
def get_alpha(name: str, **params) -> BaseAlpha:
|
||||
"""Instantiate a registered alpha by name.
|
||||
|
||||
Only the parameters accepted by the alpha's ``__init__`` are forwarded, so a
|
||||
caller may pass a superset (e.g. both ``lookback`` and ``vol_window``) and
|
||||
each alpha picks what it needs.
|
||||
|
||||
Raises:
|
||||
KeyError: If ``name`` is not registered.
|
||||
"""
|
||||
_ensure_builtins()
|
||||
if name not in _REGISTRY:
|
||||
raise KeyError(f"Unknown alpha '{name}'. Available: {sorted(_REGISTRY)}")
|
||||
cls = _REGISTRY[name]
|
||||
accepted = _accepted_params(cls)
|
||||
kwargs = params if accepted is None else {k: v for k, v in params.items() if k in accepted}
|
||||
return cls(**kwargs)
|
||||
|
||||
|
||||
def load_alpha_module(spec: str) -> None:
|
||||
"""Import an external module so its ``@register_alpha`` classes register.
|
||||
|
||||
Args:
|
||||
spec: A dotted module path (``my_pkg.my_alpha``) on ``sys.path``, or a
|
||||
filesystem path to a ``.py`` file (``/path/to/my_alpha.py``).
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If a ``.py`` path is given but does not exist.
|
||||
"""
|
||||
looks_like_file = spec.endswith(".py") or Path(spec).expanduser().exists()
|
||||
if looks_like_file:
|
||||
path = Path(spec).expanduser().resolve()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Alpha module not found: {path}")
|
||||
module_spec = importlib.util.spec_from_file_location(path.stem, path)
|
||||
if module_spec is None or module_spec.loader is None:
|
||||
raise ImportError(f"Cannot load alpha module from {path}")
|
||||
module = importlib.util.module_from_spec(module_spec)
|
||||
module_spec.loader.exec_module(module)
|
||||
else:
|
||||
importlib.import_module(spec)
|
||||
|
||||
|
||||
def _accepted_params(cls: Type[BaseAlpha]) -> Optional[set[str]]:
|
||||
"""Param names ``cls.__init__`` accepts, or None if it takes ``**kwargs``."""
|
||||
sig = inspect.signature(cls.__init__)
|
||||
if any(p.kind is p.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
return None
|
||||
return {name for name in sig.parameters if name != "self"}
|
||||
|
||||
|
||||
def _ensure_builtins() -> None:
|
||||
global _builtins_loaded
|
||||
if not _builtins_loaded:
|
||||
import pipeline.alpha.library # noqa: F401 (importing registers built-ins)
|
||||
_builtins_loaded = True
|
||||
@@ -0,0 +1,47 @@
|
||||
"""CLI for alpha combination."""
|
||||
|
||||
import os
|
||||
import click
|
||||
|
||||
from pipeline.combo.combine import combine_alphas, COMBO_METHODS
|
||||
|
||||
|
||||
@click.group(name="combo")
|
||||
def combo():
|
||||
"""Combine multiple alphas into a single combined weight."""
|
||||
|
||||
|
||||
@combo.command("combine")
|
||||
@click.option(
|
||||
"--alpha-paths", required=True,
|
||||
help="Comma-separated paths to alpha parquet files",
|
||||
)
|
||||
@click.option("--combo-name", required=True, help="Name for this combo")
|
||||
@click.option(
|
||||
"--method", default="equal_weight",
|
||||
type=click.Choice(list(COMBO_METHODS.keys())),
|
||||
help="Combination method",
|
||||
)
|
||||
@click.option("--output-dir", default="combos", help="Directory to save combo parquet")
|
||||
def combine(alpha_paths, combo_name, method, output_dir):
|
||||
"""Combine multiple alphas and save as parquet."""
|
||||
paths = [p.strip() for p in alpha_paths.split(",") if p.strip()]
|
||||
if len(paths) < 2:
|
||||
click.echo("Error: --alpha-paths requires at least 2 comma-separated paths", err=True)
|
||||
return
|
||||
|
||||
result = combine_alphas(
|
||||
alpha_paths=paths,
|
||||
combo_name=combo_name,
|
||||
method=method,
|
||||
)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
out_path = f"{output_dir}/{combo_name}.pq"
|
||||
result.to_parquet(out_path, index=False)
|
||||
click.echo(f"Saved combo: {out_path} ({len(result):,} rows)")
|
||||
click.echo(
|
||||
f"Weight stats — min: {result['weight'].min():.4f}, "
|
||||
f"max: {result['weight'].max():.4f}, "
|
||||
f"mean: {result['weight'].mean():.4f}"
|
||||
)
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Combine multiple alphas into a single combined weight.
|
||||
|
||||
Future combination methods can be registered below.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Callable
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.common.schema import COMBO_COLUMNS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _equal_weight(alpha_dfs: list[pd.DataFrame]) -> pd.DataFrame:
|
||||
"""Equal-weight combination: mean of all alpha weights per (symbol_id, date).
|
||||
|
||||
If any alpha has NaN for a symbol/date, that alpha is skipped for that row.
|
||||
"""
|
||||
# Stack all alphas with (symbol_id, date, alpha_name) as key
|
||||
combined = pd.concat(alpha_dfs, ignore_index=True)
|
||||
# Group by symbol_id + date, take mean of weights
|
||||
result = combined.groupby(["symbol_id", "date"])["weight"].mean().reset_index()
|
||||
return result
|
||||
|
||||
|
||||
# Registry of combo methods — add new functions + register them here
|
||||
COMBO_METHODS: dict[str, Callable] = {
|
||||
"equal_weight": _equal_weight,
|
||||
}
|
||||
|
||||
|
||||
def combine_alphas(
|
||||
alpha_paths: list[str],
|
||||
combo_name: str,
|
||||
method: str = "equal_weight",
|
||||
) -> pd.DataFrame:
|
||||
"""Load alphas from parquet, combine, and return combo weights.
|
||||
|
||||
Args:
|
||||
alpha_paths: List of paths to alpha parquet files.
|
||||
combo_name: Name identifier for this combo.
|
||||
method: Combination method ('equal_weight').
|
||||
|
||||
Returns:
|
||||
DataFrame with COMBO_COLUMNS.
|
||||
|
||||
Raises:
|
||||
ValueError: If method is unknown or alpha grids don't align.
|
||||
"""
|
||||
if method not in COMBO_METHODS:
|
||||
raise ValueError(
|
||||
f"Unknown combo method: {method}. Options: {list(COMBO_METHODS)}"
|
||||
)
|
||||
|
||||
alpha_dfs = []
|
||||
for path in alpha_paths:
|
||||
df = pd.read_parquet(path)
|
||||
alpha_dfs.append(df)
|
||||
logger.info("Loaded alpha: %s (%d rows)", path, len(df))
|
||||
|
||||
# Verify alignment: all alphas must share the same (symbol_id, date) pairs
|
||||
keys = [set(zip(df["symbol_id"], pd.to_datetime(df["date"]).astype(str))) for df in alpha_dfs]
|
||||
common = keys[0]
|
||||
for i, k in enumerate(keys[1:], 1):
|
||||
if k != common:
|
||||
logger.warning("Alpha %d has different (symbol_id, date) grid — intersection used", i)
|
||||
common = common.intersection(k)
|
||||
|
||||
combine_fn = COMBO_METHODS[method]
|
||||
result = combine_fn(alpha_dfs)
|
||||
result["combo_name"] = combo_name
|
||||
result = result[COMBO_COLUMNS]
|
||||
result = result.sort_values(["symbol_id", "date"]).reset_index(drop=True)
|
||||
|
||||
logger.info(
|
||||
"Combo '%s': %d symbols × %d dates, weight range [%.4f, %.4f]",
|
||||
combo_name,
|
||||
result["symbol_id"].nunique(),
|
||||
result["date"].nunique(),
|
||||
result["weight"].min(),
|
||||
result["weight"].max(),
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Column contracts for pipeline parquet files."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
# Required columns for data parquet files (daily bars, alternative data, etc.)
|
||||
DATA_COLUMNS: Final[list[str]] = [
|
||||
"symbol_id", # str: internal code like 'sh600000'
|
||||
"symbol_name", # str: stock name like '浦发银行'
|
||||
"date", # date
|
||||
"open", # float64
|
||||
"high", # float64
|
||||
"low", # float64
|
||||
"close", # float64
|
||||
"volume", # float64 (shares)
|
||||
"amount", # float64 (turnover in yuan)
|
||||
]
|
||||
|
||||
# Required columns for alpha parquet files.
|
||||
# Alphas are position WEIGHTS: positive=long, negative=short.
|
||||
ALPHA_COLUMNS: Final[list[str]] = [
|
||||
"symbol_id", # str: matches DATA_COLUMNS symbol_id
|
||||
"date", # date: aligned with data dates
|
||||
"alpha_name", # str: identifies which alpha (e.g. 'reversal_5d')
|
||||
"weight", # float64: position weight, signed
|
||||
]
|
||||
|
||||
# Required columns for combo parquet files.
|
||||
COMBO_COLUMNS: Final[list[str]] = [
|
||||
"symbol_id", # str
|
||||
"date", # date
|
||||
"combo_name", # str: identifies which combo (e.g. 'equal_weight')
|
||||
"weight", # float64: combined weight, signed
|
||||
]
|
||||
@@ -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