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
@@ -4,3 +4,9 @@ __pycache__/
|
||||
*.egg-info/
|
||||
.venv/
|
||||
venv/
|
||||
|
||||
# Pipeline outputs (regenerated by the CLI; can be large)
|
||||
data/daily_bars/
|
||||
alphas/
|
||||
combos/
|
||||
reports/
|
||||
|
||||
@@ -1,16 +1,79 @@
|
||||
# Chinese Equity Quant Research Framework
|
||||
# CLAUDE.md
|
||||
|
||||
## Architecture
|
||||
- backtrader is the backtesting engine — never reimplement backtest logic
|
||||
- akshare primary data source, baostock secondary fallback
|
||||
- Daily frequency only (Phase 1)
|
||||
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||||
|
||||
## Key Commands
|
||||
- `python3 run_example.py` — smoke test
|
||||
- `python3 -m pytest tests/ -v` — run tests
|
||||
- `pip install -r requirements.txt` — install deps
|
||||
A modular Chinese A-share quant research framework. Daily frequency only (Phase 1).
|
||||
|
||||
## Code Standards
|
||||
- Type hints on public functions
|
||||
- Google-style docstrings
|
||||
- 4-space indentation for Python
|
||||
## Commands
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt # install deps
|
||||
python3 -m pytest tests/ -v # all tests
|
||||
python3 -m pytest tests/test_alpha.py -v # single file (test_alpha is network-free)
|
||||
python3 -m pytest tests/test_alpha.py::test_evaluate_alpha_keys -v # single test
|
||||
|
||||
# Pipeline — each phase is independent: reads parquet, writes parquet.
|
||||
python3 cli.py data download --universe csi500 --start-date 2017-01-01 # → data/daily_bars/csi500/ (month-partitioned)
|
||||
python3 cli.py alpha reversal --data-path data/daily_bars/<universe> # --data-path is the dataset DIR
|
||||
python3 cli.py alpha eval --alpha-path alphas/<file>.pq --data-path data/daily_bars/<universe>
|
||||
python3 cli.py combo combine --alpha-paths a.pq,b.pq --combo-name eq --method equal_weight
|
||||
```
|
||||
|
||||
Note: `tests/test_downloader.py` hits the network (live baostock/akshare); `tests/test_alpha.py` is pure and fast.
|
||||
|
||||
## Architecture: one decoupled pipeline
|
||||
|
||||
The system is a phase-based CLI (`cli.py` + `pipeline/`). Each phase communicates **only** through parquet files on disk, so phases can be run, cached, and inspected independently:
|
||||
|
||||
```
|
||||
data → alpha → combo
|
||||
```
|
||||
|
||||
- `pipeline/data/` — download daily bars for a universe → `data/daily_bars/{universe}/month=YYYY-MM/*.pq` (Hive-partitioned dataset; pass the `{universe}` dir as `--data-path`)
|
||||
- `pipeline/alpha/` — compute one alpha's position weights from a data parquet → `alphas/*.pq`, and `alpha eval` to score it
|
||||
- `pipeline/combo/` — combine several alpha parquets into one → `combos/*.pq`
|
||||
|
||||
The pipeline reuses two top-level helper modules via a `sys.path.insert` at the top of `pipeline/data/downloader.py`: `data/downloader.py` (network download) and `data/universe.py` (constituent lists). This path hack is load-bearing — keep it.
|
||||
|
||||
## Alphas are weights, not predictors
|
||||
|
||||
An **alpha** is a signed cross-sectional **position weight**: positive = long, negative = short. It is produced by applying a formula to the wide close matrix, then **cross-sectional z-scoring** per date (`compute_alpha` in `pipeline/alpha/compute.py`). Alphas are evaluated by **return / Sharpe / turnover / max-drawdown** in `evaluate_alpha` — interpreting the weight series as a portfolio. There is deliberately **no IC/IR** anywhere: those frame a signal as a return *predictor*, which this codebase does not do. Do not reintroduce IC-style evaluation.
|
||||
|
||||
## Parquet schema contracts
|
||||
|
||||
`pipeline/common/schema.py` defines the column contracts that are the *only* interface between phases. Any new phase or alpha must conform:
|
||||
- `DATA_COLUMNS` (data output): `symbol_id, symbol_name, date, open, high, low, close, volume, amount`
|
||||
- `ALPHA_COLUMNS` (alpha output): `symbol_id, date, alpha_name, weight`
|
||||
- `COMBO_COLUMNS` (combo output): `symbol_id, date, combo_name, weight`
|
||||
|
||||
Data is stored **long/tidy**, not wide, as a Hive-partitioned dataset keyed by `month=YYYY-MM` (so reads of the dataset directory carry an extra `month` partition column, which `_pivot_close` ignores). Compute code pivots to wide (date index × symbol_id columns) internally via `_pivot_close`, where all formulas are vectorized column-wise.
|
||||
|
||||
## Alphas: factory + plugin pattern
|
||||
|
||||
Each alpha is a class subclassing `BaseAlpha` (`pipeline/alpha/base.py`), living in its own module. It implements `signal(close) -> wide DataFrame` (the raw score); the base class's `to_weights` cross-sectionally z-scores that into position weights (override for custom normalization). Subclasses declare their own typed `__init__` params (e.g. `lookback`, `vol_window`, or anything custom).
|
||||
|
||||
- Built-in alphas: one file each under `pipeline/alpha/library/` (`reversal.py`, `reversal_vol.py`, `momentum.py`). Each uses the `@register_alpha` decorator and is imported by `library/__init__.py` so it self-registers. **Add a built-in** by dropping a module there + adding it to that `__init__`.
|
||||
- Factory: `pipeline/alpha/registry.py` — `register_alpha`, `get_alpha(name, **params)` (forwards only the params the alpha's `__init__` accepts, via signature inspection), `available_alphas()`, and `load_alpha_module(spec)`.
|
||||
- **External alphas** (authored outside this repo) are the point of the design: write a `@register_alpha class MyAlpha(BaseAlpha)` in any `.py` file, then register it at runtime with `--alpha-module path/to/file.py` (or a dotted module path). See `examples/alphas/mean_reversion.py` for a working example, and `tests/test_alpha.py::test_load_external_alpha_module`.
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha list # registered alpha types
|
||||
python3 cli.py alpha list --alpha-module my_alpha.py # incl. an external one
|
||||
python3 cli.py alpha compute --alpha-module my_alpha.py \
|
||||
--alpha-type my_alpha --alpha-name run1 --param decay=0.9 --data-path <data>.pq
|
||||
```
|
||||
|
||||
`compute_alpha(data, alpha_name, alpha_type, **params)` in `pipeline/alpha/compute.py` resolves the class via `get_alpha`, applies `.weights()`, and melts to `ALPHA_COLUMNS`. `--lookback`/`--vol-window` are passed as params for convenience; arbitrary params go through repeatable `--param name=value`. `evaluate_alpha` (return/Sharpe/turnover, no IC) is unchanged.
|
||||
|
||||
Combo methods are still a plain dict registry: `COMBO_METHODS` in `pipeline/combo/combine.py`.
|
||||
|
||||
## Data sources & symbol conventions
|
||||
|
||||
- **baostock is the primary source, akshare the fallback** (`data/downloader.py`, `download_daily(source="auto")` tries baostock first). This ordering is intentional — akshare is less reliable on the deployment network.
|
||||
- Internal symbol format is `sh600000` / `sz000001` (exchange prefix + code). baostock uses the dotted form `sh.600000`; akshare's `stock_zh_a_hist` wants the bare code (prefix stripped). Both download paths return the identical 8-column schema and map the `adjust` argument consistently (`qfq`/`hfq`/none → baostock `adjustflag` via `_BAOSTOCK_ADJUST`).
|
||||
- baostock constituent queries (`get_hs300_stocks`, `get_zz500_stocks` in `data/universe.py`) return columns in an unreliable order, so `pipeline/data/downloader.py:_fix_baostock_columns` detects them by value pattern, not position.
|
||||
- Universes accepted by `data download --universe`: `hs300`, `csi500`, `all`/`full` (every listed A-share, ~5000, via `get_all_stocks` → baostock `query_all_stock` filtered to SH 6xxxxx/68xxxx + SZ 0xxxxx/3xxxxx, excluding indices & B-shares), a file path (one symbol per line), or a comma-separated symbol list. Bulk downloads use `download_daily_batch` (one baostock login for the whole run) rather than per-symbol `download_daily`.
|
||||
|
||||
## Code standards
|
||||
|
||||
- Type hints on public functions; Google-style docstrings; 4-space indentation.
|
||||
|
||||
@@ -1,8 +1,41 @@
|
||||
# Chinese Equity Quant Research Framework
|
||||
|
||||
A modular Chinese A-share quant research framework built on
|
||||
[backtrader](https://www.backtrader.com/) for backtesting, with
|
||||
akshare (primary) and baostock (fallback) for daily bar data.
|
||||
A modular Chinese A-share quant research framework. Daily frequency only (Phase 1).
|
||||
|
||||
It is a **decoupled, file-based pipeline**: each phase reads parquet and writes
|
||||
parquet, so phases run, cache, and inspect independently.
|
||||
|
||||
```
|
||||
baostock (primary) one weight series
|
||||
akshare (fallback) interpreted as a
|
||||
│ portfolio
|
||||
▼ ▲
|
||||
┌──────────────┐ ┌───────────────┐ ┌───────────────┐ │
|
||||
│ DATA │ │ ALPHA │ │ COMBO │ ┌────┴─────┐
|
||||
│ download │─────▶│ compute │─────▶│ combine │ │ EVAL │
|
||||
│ daily bars │ │ signal→weights│ │ merge alphas │ │ score it │
|
||||
└──────┬───────┘ └───────┬───────┘ └───────┬───────┘ └────┬─────┘
|
||||
│ │ │ │
|
||||
▼ ▼ ▼ │
|
||||
data/daily_bars/ alphas/*.pq combos/*.pq │
|
||||
{universe}/ (ALPHA_COLUMNS) (COMBO_COLUMNS) │
|
||||
month=YYYY-MM/*.pq │ │
|
||||
(DATA_COLUMNS) │ │
|
||||
└──────── price ───────┴───────────────────────────────────────┘
|
||||
│
|
||||
▼ (planned — not yet implemented)
|
||||
┌ ─ ─ ─ ─ ─ ─ ┐ ┌ ─ ─ ─ ─ ─ ─ ┐ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ┐
|
||||
PORTFOLIO BACKTEST PAPER TRADING
|
||||
│ construct │ │ simulate │ │ forward / live │ TODO
|
||||
positions fills + costs execution
|
||||
└ ─ ─ ─ ─ ─ ─ ┘ └ ─ ─ ─ ─ ─ ─ ┘ └ ─ ─ ─ ─ ─ ─ ─ ─ ┘
|
||||
|
||||
Each phase reads parquet and writes parquet — run, cache, and inspect
|
||||
independently. The only interface between phases is the parquet schema.
|
||||
Solid boxes are implemented; dashed boxes are on the roadmap (see TODO below).
|
||||
```
|
||||
|
||||
Data comes from **baostock (primary)** with **akshare (fallback)**.
|
||||
|
||||
## Install
|
||||
|
||||
@@ -13,13 +46,218 @@ pip install -r requirements.txt
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
python3 run_example.py # end-to-end smoke test (SMA crossover on 浦发银行)
|
||||
python3 -m pytest tests/ -v # run tests
|
||||
# 1. Download daily bars for a few symbols (writes a month-partitioned dataset).
|
||||
python3 cli.py data download \
|
||||
--universe sh600000,sz000001,sh600519 \
|
||||
--start-date 2024-01-01 --end-date 2024-03-31 \
|
||||
--output-dir data/daily_bars
|
||||
|
||||
# 2. Compute an alpha (position weights) from that data.
|
||||
# --data-path is the dataset DIRECTORY ({output-dir}/{universe}).
|
||||
python3 cli.py alpha reversal \
|
||||
--data-path "data/daily_bars/sh600000,sz000001,sh600519"
|
||||
|
||||
# 3. Evaluate it (return / Sharpe / turnover / drawdown).
|
||||
python3 cli.py alpha eval \
|
||||
--alpha-path alphas/reversal_5d.pq \
|
||||
--data-path "data/daily_bars/sh600000,sz000001,sh600519"
|
||||
|
||||
# Tests
|
||||
python3 -m pytest tests/ -v # tests/test_alpha.py is network-free; test_downloader.py hits the network
|
||||
```
|
||||
|
||||
## CLI reference
|
||||
|
||||
All commands are subcommands of `python3 cli.py`. Add `--help` to any of them.
|
||||
|
||||
### `data download` — fetch daily bars → partitioned parquet dataset
|
||||
|
||||
| Option | Default | Description |
|
||||
| --- | --- | --- |
|
||||
| `--universe` | `csi500` | `hs300`, `csi500`, `all` (~5000 A-shares), a file path (one symbol per line), or comma-separated symbols (`sh600000,sz000001`) |
|
||||
| `--start-date` | `2017-01-01` | `YYYY-MM-DD` |
|
||||
| `--end-date` | today | `YYYY-MM-DD` |
|
||||
| `--output-dir` | `data/daily_bars` | Root for the dataset directory |
|
||||
| `--symbols` | `0` | Max symbols to download (`0` = all) |
|
||||
| `--chunk-size` | `300` | Symbols per durability flush (each flush appends `.pq` files) |
|
||||
| `--adjust` | `qfq` | Price adjustment: `qfq` (forward), `hfq` (backward), `none` |
|
||||
|
||||
Writes a **Hive-partitioned dataset** at `{output_dir}/{universe}/month=YYYY-MM/*.pq`
|
||||
(one partition per calendar month). The `{universe}` directory is rebuilt from
|
||||
scratch on each run. Downloads stream under a single baostock session and flush
|
||||
every `--chunk-size` symbols, so memory stays bounded and a crash keeps the
|
||||
partitions already written. Pass the **dataset directory** (`{output_dir}/{universe}`)
|
||||
as `--data-path` to later phases — `pd.read_parquet` reads the whole partitioned
|
||||
set. Symbols use the internal `sh600000` / `sz000001` form (exchange prefix + code).
|
||||
|
||||
### `alpha list` — show registered alpha types
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha list
|
||||
python3 cli.py alpha list --alpha-module path/to/my_alpha.py # include an external alpha
|
||||
```
|
||||
|
||||
### `alpha compute` — alpha class → weights parquet
|
||||
|
||||
| Option | Default | Description |
|
||||
| --- | --- | --- |
|
||||
| `--data-path` | (required) | Data parquet from `data download` |
|
||||
| `--alpha-name` | (required) | Label stored in the `alpha_name` column / output filename |
|
||||
| `--alpha-type` | (required) | Registry key of the alpha class (see `alpha list`) |
|
||||
| `--output-dir` | `alphas` | Output directory |
|
||||
| `--lookback` | `5` | Lookback days (passed to alphas that accept it) |
|
||||
| `--vol-window` | `20` | Volatility window (passed to alphas that accept it) |
|
||||
| `--alpha-module` | — | External module(s) to import first; repeatable. Dotted path or `.py` file |
|
||||
| `--param` | — | Extra constructor param as `name=value`; repeatable |
|
||||
|
||||
Only the params an alpha's `__init__` accepts are forwarded, so passing extras
|
||||
(e.g. `--vol-window` to a reversal alpha) is harmless.
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha compute \
|
||||
--data-path <data>.pq \
|
||||
--alpha-type reversal_vol --alpha-name rv_5_20 \
|
||||
--lookback 5 --vol-window 20
|
||||
```
|
||||
|
||||
Shortcuts for the two most common built-ins:
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha reversal --data-path <data>.pq --lookback 5
|
||||
python3 cli.py alpha reversal-vol --data-path <data>.pq --lookback 5 --vol-window 20
|
||||
```
|
||||
|
||||
### `alpha eval` — score an alpha as a portfolio
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha eval --alpha-path alphas/<name>.pq --data-path <data>.pq
|
||||
```
|
||||
|
||||
Interprets the weights as a portfolio and reports cumulative return, annual
|
||||
Sharpe, annual turnover, max drawdown, and hit rate; also dumps
|
||||
`reports/<alpha_name>_eval.json`. There is deliberately **no IC/IR** — alphas are
|
||||
position weights, not return predictors.
|
||||
|
||||
### `combo combine` — merge several alphas into one weight
|
||||
|
||||
| Option | Default | Description |
|
||||
| --- | --- | --- |
|
||||
| `--alpha-paths` | (required) | Comma-separated alpha parquet paths (≥ 2) |
|
||||
| `--combo-name` | (required) | Label stored in the `combo_name` column / output filename |
|
||||
| `--method` | `equal_weight` | Combination method (see `COMBO_METHODS`) |
|
||||
| `--output-dir` | `combos` | Output directory |
|
||||
|
||||
```bash
|
||||
python3 cli.py combo combine \
|
||||
--alpha-paths alphas/reversal_5d.pq,alphas/reversal_vol_5d_20d.pq \
|
||||
--combo-name eq --method equal_weight
|
||||
```
|
||||
|
||||
## Alphas: the factory / plugin interface
|
||||
|
||||
An **alpha** is a class that maps a wide close matrix (date index × `symbol_id`
|
||||
columns) to **signed position weights** (positive = long, negative = short).
|
||||
Every alpha subclasses `BaseAlpha` (`pipeline/alpha/base.py`) and is resolved by
|
||||
name through the registry (`pipeline/alpha/registry.py`).
|
||||
|
||||
### Minimal alpha
|
||||
|
||||
```python
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
|
||||
@register_alpha
|
||||
class MyAlpha(BaseAlpha):
|
||||
name = "my_alpha" # unique registry key (required)
|
||||
|
||||
def __init__(self, lookback: int = 5):
|
||||
self.lookback = lookback # declare whatever params you need
|
||||
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
# Raw score: wide (date × symbol_id), higher = stronger long, NaN where undefined.
|
||||
return -close.pct_change(self.lookback)
|
||||
```
|
||||
|
||||
That is the whole contract:
|
||||
|
||||
- `name` — the `--alpha-type` key; must be unique.
|
||||
- `signal(close)` — the only required method; return a wide DataFrame.
|
||||
- `to_weights(signal)` — provided by the base class: cross-sectionally z-scores
|
||||
each date into weights (NaN → 0). **Override** it for a different scheme (rank,
|
||||
dollar-neutral caps, etc.).
|
||||
|
||||
### Built-in alphas
|
||||
|
||||
One file per alpha under `pipeline/alpha/library/`:
|
||||
|
||||
| `--alpha-type` | Params | Description |
|
||||
| --- | --- | --- |
|
||||
| `reversal` | `lookback` | Negative trailing return (oversold scores high) |
|
||||
| `reversal_vol` | `lookback`, `vol_window` | Reversal scaled by trailing volatility |
|
||||
| `momentum` | `lookback` | Positive trailing return |
|
||||
|
||||
Add a built-in by dropping a module in `pipeline/alpha/library/` and importing it
|
||||
from that package's `__init__.py`.
|
||||
|
||||
### Using an alpha written outside this repo
|
||||
|
||||
Write your `@register_alpha` class in any `.py` file, then register it at runtime
|
||||
with `--alpha-module` (a `.py` path or an importable dotted module). See the
|
||||
worked example in `examples/alphas/mean_reversion.py`:
|
||||
|
||||
```bash
|
||||
python3 cli.py alpha compute \
|
||||
--alpha-module examples/alphas/mean_reversion.py \
|
||||
--alpha-type mean_reversion --alpha-name mr20 \
|
||||
--param window=20 \
|
||||
--data-path <data>.pq
|
||||
```
|
||||
|
||||
`mean_reversion` declares a `window` param (not `lookback`); `--param window=20`
|
||||
supplies it and the unrelated `--lookback`/`--vol-window` defaults are ignored.
|
||||
|
||||
## Parquet schemas
|
||||
|
||||
The column contracts in `pipeline/common/schema.py` are the only interface
|
||||
between phases (data is stored long/tidy):
|
||||
|
||||
- **data** (`DATA_COLUMNS`): `symbol_id, symbol_name, date, open, high, low, close, volume, amount`
|
||||
- **alpha** (`ALPHA_COLUMNS`): `symbol_id, date, alpha_name, weight`
|
||||
- **combo** (`COMBO_COLUMNS`): `symbol_id, date, combo_name, weight`
|
||||
|
||||
The data phase writes a month-partitioned dataset, so reading the dataset
|
||||
directory yields an extra `month` (`YYYY-MM`) partition column on top of
|
||||
`DATA_COLUMNS`; the alpha phase pivots by name and ignores it.
|
||||
|
||||
## Layout
|
||||
|
||||
- `data/` — unified downloader (akshare -> baostock fallback) and data schema
|
||||
- `backtest/` — config, pandas->backtrader feed adapter, and `BacktestRunner`
|
||||
- `strategies/` — example `SmaCross` strategy
|
||||
- `analysis/` — performance reporting (sharpe, drawdown, returns, trades)
|
||||
- `cli.py` — entry point wiring the three phases together
|
||||
- `pipeline/data/` — universe resolution + download → `data/daily_bars/{universe}/month=YYYY-MM/*.pq`
|
||||
- `pipeline/alpha/` — `base.py` (`BaseAlpha`), `registry.py` (factory + plugin loader),
|
||||
`library/` (built-in alphas), `compute.py` (`compute_alpha` / `evaluate_alpha`)
|
||||
- `pipeline/combo/` — alpha combination → `combos/*.pq`
|
||||
- `pipeline/common/schema.py` — parquet column contracts
|
||||
- `data/downloader.py`, `data/universe.py` — baostock/akshare download + constituents
|
||||
- `examples/alphas/` — example external alpha(s)
|
||||
|
||||
## Roadmap (not yet implemented)
|
||||
|
||||
The pipeline currently ends at `combo`, and `alpha eval` only interprets a weight
|
||||
series as a portfolio for quick scoring (return / Sharpe / turnover / drawdown).
|
||||
It is **not** a true backtest — there is no transaction-cost, slippage, or
|
||||
execution modeling. The following phases are planned but not built yet:
|
||||
|
||||
- [ ] **Portfolio construction** — turn combo weights into target positions
|
||||
(gross/net exposure caps, per-name and sector limits, capital allocation,
|
||||
rebalance schedule).
|
||||
- [ ] **Backtesting** — event-driven simulation over the constructed positions
|
||||
with realistic fills, transaction costs, slippage, and borrow constraints;
|
||||
richer P&L / risk attribution than `alpha eval`.
|
||||
- [ ] **Forward / paper trading** — run the same construction logic on live
|
||||
daily data, track simulated fills and a running P&L without real capital.
|
||||
|
||||
Until these land, treat `alpha eval` as a fast sanity check on a weight series,
|
||||
not a performance estimate.
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
"""Performance analysis and reporting for backtest results."""
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
import pandas as pd # noqa: E402
|
||||
|
||||
|
||||
def print_results(results: list, initial_cash: float = 1_000_000.0) -> dict[str, Any]:
|
||||
"""Print and return key performance metrics from a backtrader run result."""
|
||||
if not results:
|
||||
print("No results to report.")
|
||||
return {}
|
||||
|
||||
result = results[0]
|
||||
report = {}
|
||||
|
||||
# Sharpe ratio
|
||||
sharpe = result.analyzers.sharpe.get_analysis()
|
||||
report["sharpe"] = sharpe.get("sharperatio", "N/A")
|
||||
|
||||
# Drawdown
|
||||
dd = result.analyzers.drawdown.get_analysis()
|
||||
report["max_drawdown"] = dd.get("max", {}).get("drawdown", "N/A")
|
||||
report["max_drawdown_len"] = dd.get("max", {}).get("len", "N/A")
|
||||
|
||||
# Returns
|
||||
rets = result.analyzers.returns.get_analysis()
|
||||
report["total_return"] = rets.get("rtot", "N/A")
|
||||
report["avg_return"] = rets.get("ravg", "N/A")
|
||||
|
||||
# Trades
|
||||
trades = result.analyzers.trades.get_analysis()
|
||||
report["total_trades"] = trades.get("total", {}).get("total", 0)
|
||||
report["won_trades"] = trades.get("won", {}).get("total", 0)
|
||||
report["lost_trades"] = trades.get("lost", {}).get("total", 0)
|
||||
|
||||
# Print
|
||||
print("=" * 50)
|
||||
print("BACKTEST RESULTS")
|
||||
print("=" * 50)
|
||||
print(f"Sharpe Ratio: {report['sharpe']}")
|
||||
print(f"Total Return: {report['total_return']:.4%}" if isinstance(report['total_return'], float) else f"Total Return: {report['total_return']}")
|
||||
print(f"Max Drawdown: {report['max_drawdown']:.2%}" if isinstance(report['max_drawdown'], float) else f"Max Drawdown: {report['max_drawdown']}")
|
||||
print(f"Max DD Length: {report['max_drawdown_len']}")
|
||||
print(f"Total Trades: {report['total_trades']}")
|
||||
print(f"Won/Lost: {report['won_trades']}/{report['lost_trades']}")
|
||||
print("=" * 50)
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def plot_accumulated_pnl(
|
||||
results: list, output_path: str = "reports/pnl.png", initial_cash: float = 1_000_000.0
|
||||
) -> str:
|
||||
"""Plot accumulated portfolio value from a backtest run.
|
||||
|
||||
Reads the per-day TimeReturn analyzer attached by ``BacktestRunner`` and
|
||||
compounds it into an equity curve.
|
||||
|
||||
Args:
|
||||
results: The list returned by ``cerebro.run()``.
|
||||
output_path: Destination PNG path.
|
||||
initial_cash: Starting portfolio value for scaling the curve.
|
||||
|
||||
Returns:
|
||||
The path the chart was written to.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
|
||||
series = pd.Series(dtype=float)
|
||||
if results:
|
||||
tr = results[0].analyzers.timereturn.get_analysis()
|
||||
series = pd.Series(tr).sort_index()
|
||||
|
||||
fig, ax = plt.subplots(figsize=(10, 5))
|
||||
if len(series):
|
||||
equity = (1.0 + series).cumprod() * initial_cash
|
||||
ax.plot(equity.index, equity.values, color="C0")
|
||||
ax.set_title("Accumulated Portfolio Value")
|
||||
ax.set_xlabel("Date")
|
||||
ax.set_ylabel("Value")
|
||||
ax.grid(True, alpha=0.3)
|
||||
fig.autofmt_xdate()
|
||||
fig.tight_layout()
|
||||
fig.savefig(output_path, dpi=100)
|
||||
plt.close(fig)
|
||||
return output_path
|
||||
|
||||
|
||||
def plot_ic(signal_eval: dict, output_path: str = "reports/ic.png") -> str:
|
||||
"""Plot the per-period rank IC time series from a signal evaluation.
|
||||
|
||||
Args:
|
||||
signal_eval: Dict returned by ``evaluate_cross_sectional`` (expects a
|
||||
``rank_ic_series`` pandas Series).
|
||||
output_path: Destination PNG path.
|
||||
|
||||
Returns:
|
||||
The path the chart was written to.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
|
||||
rank_ic = signal_eval.get("rank_ic_series", pd.Series(dtype=float))
|
||||
|
||||
fig, ax = plt.subplots(figsize=(10, 5))
|
||||
if len(rank_ic):
|
||||
ax.bar(rank_ic.index, rank_ic.values, width=1.0, color="C1", alpha=0.6, label="Rank IC")
|
||||
ax.axhline(rank_ic.mean(), color="C7", linestyle="--", label=f"mean={rank_ic.mean():.3f}")
|
||||
cum_mean = rank_ic.expanding().mean()
|
||||
ax.plot(cum_mean.index, cum_mean.values, color="red", linewidth=1.5, label="Cumulative mean IC")
|
||||
ax.legend()
|
||||
ax.set_title("Cross-Sectional Rank IC")
|
||||
ax.set_xlabel("Date")
|
||||
ax.set_ylabel("Rank IC")
|
||||
ax.grid(True, alpha=0.3)
|
||||
fig.autofmt_xdate()
|
||||
fig.tight_layout()
|
||||
fig.savefig(output_path, dpi=100)
|
||||
plt.close(fig)
|
||||
return output_path
|
||||
|
||||
|
||||
def dump_signals(signals_df: pd.DataFrame, output_dir: str = "results/") -> str:
|
||||
"""Save the signal matrix (date x stock) as a parquet file.
|
||||
|
||||
Args:
|
||||
signals_df: Date-indexed DataFrame of per-stock signal values.
|
||||
output_dir: Directory to write the parquet file into.
|
||||
|
||||
Returns:
|
||||
The path the parquet file was written to.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
path = os.path.join(output_dir, "signals.parquet")
|
||||
signals_df.to_parquet(path)
|
||||
return path
|
||||
|
||||
|
||||
def dump_daily_pnl(
|
||||
results: list, output_dir: str = "results/", initial_cash: float = 1_000_000.0
|
||||
) -> str:
|
||||
"""Extract the daily portfolio value from a backtest run and save as parquet.
|
||||
|
||||
Compounds the per-day TimeReturn analyzer into an equity curve.
|
||||
|
||||
Args:
|
||||
results: The list returned by ``cerebro.run()``.
|
||||
output_dir: Directory to write the parquet file into.
|
||||
initial_cash: Starting portfolio value for scaling the curve.
|
||||
|
||||
Returns:
|
||||
The path the parquet file was written to.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
series = pd.Series(dtype=float)
|
||||
if results:
|
||||
tr = results[0].analyzers.timereturn.get_analysis()
|
||||
series = pd.Series(tr).sort_index()
|
||||
|
||||
equity = (1.0 + series).cumprod() * initial_cash
|
||||
pnl_df = pd.DataFrame({"date": equity.index, "value": equity.values})
|
||||
|
||||
path = os.path.join(output_dir, "daily_pnl.parquet")
|
||||
pnl_df.to_parquet(path)
|
||||
return path
|
||||
|
||||
|
||||
def generate_report(
|
||||
results: list,
|
||||
signal_eval: dict,
|
||||
output_dir: str = "reports/",
|
||||
initial_cash: float = 1_000_000.0,
|
||||
) -> dict[str, str]:
|
||||
"""Generate the full report: PnL chart, IC chart, and a summary text file.
|
||||
|
||||
Args:
|
||||
results: The list returned by ``cerebro.run()``.
|
||||
signal_eval: Dict returned by ``evaluate_cross_sectional``.
|
||||
output_dir: Directory to write artifacts into.
|
||||
initial_cash: Starting portfolio value.
|
||||
|
||||
Returns:
|
||||
Mapping of artifact name to file path.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
pnl_path = plot_accumulated_pnl(
|
||||
results, os.path.join(output_dir, "pnl.png"), initial_cash
|
||||
)
|
||||
ic_path = plot_ic(signal_eval, os.path.join(output_dir, "ic.png"))
|
||||
|
||||
metrics = print_results(results, initial_cash)
|
||||
summary_path = os.path.join(output_dir, "summary.txt")
|
||||
with open(summary_path, "w") as f:
|
||||
f.write("BACKTEST SUMMARY\n")
|
||||
f.write("=" * 40 + "\n")
|
||||
for k, v in metrics.items():
|
||||
f.write(f"{k}: {v}\n")
|
||||
f.write("\nSIGNAL IC\n")
|
||||
f.write("=" * 40 + "\n")
|
||||
for k in ("ic_mean", "ic_std", "ir", "rank_ic_mean", "rank_ic_std", "rank_ir", "hit_rate", "n_periods"):
|
||||
if k in signal_eval:
|
||||
f.write(f"{k}: {signal_eval[k]}\n")
|
||||
|
||||
return {"pnl": pnl_path, "ic": ic_path, "summary": summary_path}
|
||||
@@ -1,14 +0,0 @@
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import date
|
||||
|
||||
|
||||
@dataclass
|
||||
class BacktestConfig:
|
||||
symbols: list[str] = field(default_factory=lambda: ["sh600000"])
|
||||
start_date: str = "2023-01-01"
|
||||
end_date: str = "2024-12-31"
|
||||
initial_cash: float = 1_000_000.0
|
||||
commission: float = 0.0003 # 0.03% for Chinese A-shares
|
||||
stamp_duty: float = 0.001 # 0.1% stamp duty on sells only (handled in strategy)
|
||||
adjust: str = "qfq"
|
||||
sizer_percent: float = 0.95 # fraction of portfolio per trade
|
||||
@@ -1,35 +0,0 @@
|
||||
"""Convert pandas DataFrames to backtrader data feeds."""
|
||||
import backtrader as bt
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def df_to_bt_feed(df: pd.DataFrame) -> bt.feeds.PandasData:
|
||||
"""Convert a standardized OHLCV DataFrame to a backtrader PandasData feed."""
|
||||
df = df.copy()
|
||||
df["date"] = pd.to_datetime(df["date"])
|
||||
df = df.set_index("date")
|
||||
df = df[["open", "high", "low", "close", "volume"]]
|
||||
return bt.feeds.PandasData(dataname=df)
|
||||
|
||||
|
||||
class SignalPandasData(bt.feeds.PandasData):
|
||||
"""PandasData feed carrying an extra ``signal`` line alongside OHLCV."""
|
||||
|
||||
lines = ("signal",)
|
||||
params = (("signal", -1),) # -1 -> match by column name
|
||||
|
||||
|
||||
def df_to_signal_feed(df: pd.DataFrame) -> "SignalPandasData":
|
||||
"""Convert an OHLCV+signal DataFrame to a SignalPandasData feed.
|
||||
|
||||
Args:
|
||||
df: DataFrame with ``date``, OHLCV columns, and a ``signal`` column.
|
||||
|
||||
Returns:
|
||||
A SignalPandasData feed (NaN signals are preserved for the strategy to skip).
|
||||
"""
|
||||
df = df.copy()
|
||||
df["date"] = pd.to_datetime(df["date"])
|
||||
df = df.set_index("date")
|
||||
df = df[["open", "high", "low", "close", "volume", "signal"]]
|
||||
return SignalPandasData(dataname=df)
|
||||
@@ -1,84 +0,0 @@
|
||||
"""BacktestRunner: orchestrates data loading, cerebro setup, and execution."""
|
||||
import logging
|
||||
import backtrader as bt
|
||||
from typing import Optional
|
||||
|
||||
from backtest.config import BacktestConfig
|
||||
from backtest.feed import df_to_bt_feed, df_to_signal_feed
|
||||
from data.downloader import download_daily
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BacktestRunner:
|
||||
"""Run backtrader backtests with Chinese equity data."""
|
||||
|
||||
def __init__(self, config: BacktestConfig):
|
||||
self.config = config
|
||||
self.cerebro = bt.Cerebro()
|
||||
self._results: Optional[list] = None
|
||||
|
||||
def add_data(self, symbol: str) -> None:
|
||||
"""Download data for a symbol and add to cerebro as a feed."""
|
||||
df = download_daily(
|
||||
symbol=symbol,
|
||||
start=self.config.start_date,
|
||||
end=self.config.end_date,
|
||||
adjust=self.config.adjust,
|
||||
)
|
||||
feed = df_to_bt_feed(df)
|
||||
self.cerebro.adddata(feed, name=symbol)
|
||||
logger.info(f"Added {symbol}: {len(df)} bars")
|
||||
|
||||
def add_signal_data(self, df, name: str) -> None:
|
||||
"""Add a pre-built OHLCV+signal DataFrame as a SignalPandasData feed."""
|
||||
feed = df_to_signal_feed(df)
|
||||
self.cerebro.adddata(feed, name=name)
|
||||
logger.info(f"Added signal feed {name}: {len(df)} bars")
|
||||
|
||||
def add_strategy(self, strategy_cls, **kwargs) -> None:
|
||||
"""Add a strategy class to cerebro."""
|
||||
self.cerebro.addstrategy(strategy_cls, **kwargs)
|
||||
|
||||
def _configure(self) -> None:
|
||||
"""Configure broker, sizer, and analyzers (independent of data feeds)."""
|
||||
self.cerebro.broker.setcash(self.config.initial_cash)
|
||||
self.cerebro.broker.setcommission(commission=self.config.commission)
|
||||
self.cerebro.addsizer(bt.sizers.PercentSizer, percents=self.config.sizer_percent * 100)
|
||||
|
||||
self.cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe", riskfreerate=0.02)
|
||||
self.cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")
|
||||
self.cerebro.addanalyzer(bt.analyzers.Returns, _name="returns")
|
||||
self.cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name="trades")
|
||||
self.cerebro.addanalyzer(
|
||||
bt.analyzers.TimeReturn, _name="timereturn", timeframe=bt.TimeFrame.Days
|
||||
)
|
||||
|
||||
def setup(self, strategy_cls, strategy_kwargs: Optional[dict] = None) -> None:
|
||||
"""Full setup: load data for all symbols, configure cerebro, add strategy."""
|
||||
for sym in self.config.symbols:
|
||||
self.add_data(sym)
|
||||
self._configure()
|
||||
self.cerebro.addstrategy(strategy_cls, **(strategy_kwargs or {}))
|
||||
|
||||
def run(self, strategy_cls, strategy_kwargs: Optional[dict] = None) -> list:
|
||||
"""Setup (downloading all symbols) and run the backtest."""
|
||||
self.setup(strategy_cls, strategy_kwargs)
|
||||
return self._execute()
|
||||
|
||||
def run_prepared(self, strategy_cls, strategy_kwargs: Optional[dict] = None) -> list:
|
||||
"""Run a backtest using feeds already added via ``add_signal_data``."""
|
||||
self._configure()
|
||||
self.cerebro.addstrategy(strategy_cls, **(strategy_kwargs or {}))
|
||||
return self._execute()
|
||||
|
||||
def _execute(self) -> list:
|
||||
start_val = self.cerebro.broker.getvalue()
|
||||
logger.info(f"Starting portfolio value: {start_val:,.2f}")
|
||||
self._results = self.cerebro.run()
|
||||
end_val = self.cerebro.broker.getvalue()
|
||||
logger.info(f"Ending portfolio value: {end_val:,.2f}")
|
||||
return self._results
|
||||
|
||||
def get_results(self) -> Optional[list]:
|
||||
return self._results
|
||||
@@ -0,0 +1,43 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Chinese Equity Quant Pipeline — decoupled phase CLI.
|
||||
|
||||
Phases:
|
||||
data — Download daily bars to parquet
|
||||
alpha — Compute alpha weights from data
|
||||
combo — Combine alphas into a single weight
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import click
|
||||
|
||||
from pipeline.data.cli import data
|
||||
from pipeline.alpha.cli import alpha
|
||||
from pipeline.combo.cli import combo
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.option(
|
||||
"--log-level", default="INFO",
|
||||
type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"], case_sensitive=False),
|
||||
help="Logging verbosity (default INFO shows download/compute progress)",
|
||||
)
|
||||
def cli(log_level):
|
||||
"""Chinese Equity Quant Pipeline.
|
||||
|
||||
Each phase is independent: read from parquet, write to parquet.
|
||||
"""
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level.upper()),
|
||||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
)
|
||||
|
||||
|
||||
cli.add_command(data)
|
||||
cli.add_command(alpha)
|
||||
cli.add_command(combo)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
+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)
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
"""Signal evaluation metrics."""
|
||||
from eval.metrics import evaluate_cross_sectional
|
||||
|
||||
__all__ = ["evaluate_cross_sectional"]
|
||||
@@ -1,90 +0,0 @@
|
||||
"""Information-coefficient metrics for alpha signals."""
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def _summarize(ic: pd.Series, rank_ic: pd.Series) -> dict[str, Any]:
|
||||
"""Aggregate per-period IC series into summary statistics."""
|
||||
ic = ic.dropna()
|
||||
rank_ic = rank_ic.dropna()
|
||||
|
||||
ic_mean = float(ic.mean()) if len(ic) else float("nan")
|
||||
ic_std = float(ic.std()) if len(ic) else float("nan")
|
||||
rank_ic_mean = float(rank_ic.mean()) if len(rank_ic) else float("nan")
|
||||
rank_ic_std = float(rank_ic.std()) if len(rank_ic) else float("nan")
|
||||
|
||||
return {
|
||||
"ic_mean": ic_mean,
|
||||
"ic_std": ic_std,
|
||||
"ir": ic_mean / ic_std if ic_std else float("nan"),
|
||||
"rank_ic_mean": rank_ic_mean,
|
||||
"rank_ic_std": rank_ic_std,
|
||||
"rank_ir": rank_ic_mean / rank_ic_std if rank_ic_std else float("nan"),
|
||||
"hit_rate": float((rank_ic > 0).mean()) if len(rank_ic) else float("nan"),
|
||||
"n_periods": int(len(rank_ic)),
|
||||
"ic_series": ic,
|
||||
"rank_ic_series": rank_ic,
|
||||
}
|
||||
|
||||
|
||||
def _cross_sectional(signals_df: pd.DataFrame, returns_df: pd.DataFrame) -> dict[str, Any]:
|
||||
"""Per-date IC across stocks (requires >= 2 stocks)."""
|
||||
dates = signals_df.index
|
||||
ic_vals, rank_ic_vals, idx = [], [], []
|
||||
for dt in dates:
|
||||
s = signals_df.loc[dt]
|
||||
r = returns_df.loc[dt]
|
||||
mask = s.notna() & r.notna()
|
||||
if mask.sum() < 2:
|
||||
continue
|
||||
sv, rv = s[mask], r[mask]
|
||||
# A degenerate (constant) vector makes correlation undefined.
|
||||
if sv.nunique() < 2 or rv.nunique() < 2:
|
||||
continue
|
||||
ic_vals.append(sv.corr(rv))
|
||||
rank_ic_vals.append(sv.corr(rv, method="spearman"))
|
||||
idx.append(dt)
|
||||
ic = pd.Series(ic_vals, index=idx, dtype=float)
|
||||
rank_ic = pd.Series(rank_ic_vals, index=idx, dtype=float)
|
||||
return _summarize(ic, rank_ic)
|
||||
|
||||
|
||||
def _rolling_single(
|
||||
signals_df: pd.DataFrame, returns_df: pd.DataFrame, window: int = 20
|
||||
) -> dict[str, Any]:
|
||||
"""Rolling time-series IC for the single-stock case.
|
||||
|
||||
With one stock there is no cross-section, so we measure how well the signal
|
||||
tracks forward returns over a trailing window instead.
|
||||
"""
|
||||
col = signals_df.columns[0]
|
||||
s = signals_df[col]
|
||||
r = returns_df[col]
|
||||
ic = s.rolling(window).corr(r)
|
||||
rank_ic = s.rank().rolling(window).corr(r.rank())
|
||||
return _summarize(ic, rank_ic)
|
||||
|
||||
|
||||
def evaluate_cross_sectional(
|
||||
signals_df: pd.DataFrame, returns_df: pd.DataFrame
|
||||
) -> dict[str, Any]:
|
||||
"""Evaluate a signal's predictive power against forward returns.
|
||||
|
||||
Args:
|
||||
signals_df: DataFrame indexed by date, one column per stock, signal values.
|
||||
returns_df: DataFrame indexed by date, one column per stock, forward returns.
|
||||
|
||||
Returns:
|
||||
Dict with ``ic_mean``, ``ic_std``, ``ir``, ``rank_ic_mean``,
|
||||
``rank_ic_std``, ``rank_ir``, ``hit_rate``, ``n_periods`` and the
|
||||
per-period ``ic_series`` / ``rank_ic_series`` (for plotting).
|
||||
"""
|
||||
cols = signals_df.columns.intersection(returns_df.columns)
|
||||
idx = signals_df.index.intersection(returns_df.index)
|
||||
signals_df = signals_df.loc[idx, cols]
|
||||
returns_df = returns_df.loc[idx, cols]
|
||||
|
||||
if len(cols) >= 2:
|
||||
return _cross_sectional(signals_df, returns_df)
|
||||
return _rolling_single(signals_df, returns_df)
|
||||
@@ -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()),
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
"""Translate signal values into position actions."""
|
||||
from portfolio.builder import PositionAction, ThresholdBuilder
|
||||
|
||||
__all__ = ["PositionAction", "ThresholdBuilder"]
|
||||
@@ -1,69 +0,0 @@
|
||||
"""Map signal values to discrete position actions."""
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class PositionAction:
|
||||
"""A target action for a single stock on a single bar."""
|
||||
|
||||
action: str # "buy", "sell", or "hold"
|
||||
size_pct: float = 0.0 # target portfolio fraction for buys
|
||||
|
||||
|
||||
class ThresholdBuilder:
|
||||
"""Open on strong positive signal, close on strong negative signal."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
buy_threshold: float = 0.02,
|
||||
sell_threshold: float = -0.02,
|
||||
size_pct: float = 0.95,
|
||||
):
|
||||
self.buy_threshold = buy_threshold
|
||||
self.sell_threshold = sell_threshold
|
||||
self.size_pct = size_pct
|
||||
|
||||
def build(self, signal_value: float, in_position: bool) -> PositionAction:
|
||||
if not in_position and signal_value >= self.buy_threshold:
|
||||
return PositionAction("buy", self.size_pct)
|
||||
if in_position and signal_value <= self.sell_threshold:
|
||||
return PositionAction("sell", 0.0)
|
||||
return PositionAction("hold", 0.0)
|
||||
|
||||
|
||||
class RankEqualWeightBuilder:
|
||||
"""Rank all stocks by signal. Buy top N% at equal weight. Sell if drops out.
|
||||
|
||||
Called once per bar with ALL stock signals. Returns per-stock actions.
|
||||
"""
|
||||
|
||||
def __init__(self, top_n: Optional[int] = None, top_pct: float = 0.2, min_signal: Optional[float] = None):
|
||||
self.top_n = top_n
|
||||
self.top_pct = top_pct
|
||||
self.min_signal = min_signal
|
||||
|
||||
def build(self, signals: dict[str, float]) -> dict[str, PositionAction]:
|
||||
# Filter by min_signal if set
|
||||
if self.min_signal is not None:
|
||||
signals = {s: v for s, v in signals.items() if v >= self.min_signal}
|
||||
|
||||
# Sort by signal descending
|
||||
ranked = sorted(signals.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Determine top N: explicit count or percentage of available stocks
|
||||
if self.top_n is not None:
|
||||
n = self.top_n
|
||||
else:
|
||||
n = max(1, int(len(signals) * self.top_pct))
|
||||
|
||||
top_symbols = set(sym for sym, _ in ranked[:n])
|
||||
size_pct = 1.0 / n if n > 0 else 0.0
|
||||
|
||||
actions = {}
|
||||
for sym in signals:
|
||||
if sym in top_symbols:
|
||||
actions[sym] = PositionAction("buy", size_pct)
|
||||
else:
|
||||
actions[sym] = PositionAction("sell", 0.0)
|
||||
return actions
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 40 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 59 KiB |
@@ -1,21 +0,0 @@
|
||||
BACKTEST SUMMARY
|
||||
========================================
|
||||
sharpe: 0.12450603119200966
|
||||
max_drawdown: 18.40327026827532
|
||||
max_drawdown_len: 251
|
||||
total_return: 0.04945463098329521
|
||||
avg_return: 0.00010217898963490745
|
||||
total_trades: 695
|
||||
won_trades: 357
|
||||
lost_trades: 333
|
||||
|
||||
SIGNAL IC
|
||||
========================================
|
||||
ic_mean: 0.03559005065789478
|
||||
ic_std: 0.2671528697174321
|
||||
ir: 0.13321979545096566
|
||||
rank_ic_mean: 0.022676047691591646
|
||||
rank_ic_std: 0.24733432522614088
|
||||
rank_ir: 0.09168176584814361
|
||||
hit_rate: 0.5315904139433552
|
||||
n_periods: 459
|
||||
@@ -4,3 +4,5 @@ baostock>=0.8.8
|
||||
pandas>=2.0.0
|
||||
matplotlib>=3.7.0
|
||||
pytest>=7.0.0
|
||||
click>=8.0.0
|
||||
pyarrow>=14.0.0
|
||||
|
||||
Binary file not shown.
Binary file not shown.
-132
@@ -1,132 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""End-to-end pipeline: universe -> signal -> cross-sectional IC
|
||||
-> multi-stock backtest (AlphaStrategy + RankEqualWeightBuilder) -> reports.
|
||||
|
||||
Usage:
|
||||
python3 run_example.py --universe hs300 --signal reversal
|
||||
python3 run_example.py --universe csi500 --signal reversal_vol
|
||||
"""
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from analysis.report import dump_daily_pnl, dump_signals, generate_report
|
||||
from backtest.config import BacktestConfig
|
||||
from backtest.runner import BacktestRunner
|
||||
from data.downloader import download_batch
|
||||
from data.universe import SYMBOLS, CSI500_SYMBOLS
|
||||
from eval.metrics import evaluate_cross_sectional
|
||||
from portfolio.builder import RankEqualWeightBuilder
|
||||
from signals.reversal import ReversalSignal
|
||||
from signals.reversal_vol import ReversalVolSignal
|
||||
from strategies.alpha_strategy import AlphaStrategy
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _forward_returns(data: dict[str, pd.DataFrame], horizon: int) -> pd.DataFrame:
|
||||
"""Build a date-indexed DataFrame of ``horizon``-day forward returns per stock."""
|
||||
forward_returns: dict[str, pd.Series] = {}
|
||||
for sym, df in data.items():
|
||||
fwd = df["close"].pct_change(horizon).shift(-horizon)
|
||||
fwd.index = pd.to_datetime(df["date"])
|
||||
forward_returns[sym] = fwd
|
||||
return pd.DataFrame(forward_returns)
|
||||
|
||||
|
||||
def main(forward_horizon: int = 5, universe: str = "csi500", signal_name: str = "reversal_vol",
|
||||
dump_dir: str = "results/"):
|
||||
universes = {"hs300": SYMBOLS, "csi500": CSI500_SYMBOLS}
|
||||
symbols = universes.get(universe, CSI500_SYMBOLS)[:30]
|
||||
|
||||
signals = {
|
||||
"reversal": ReversalSignal(lookback=5),
|
||||
"reversal_vol": ReversalVolSignal(lookback=5, vol_window=20),
|
||||
}
|
||||
signal = signals.get(signal_name, ReversalVolSignal(lookback=5, vol_window=20))
|
||||
|
||||
start, end = "2023-01-01", "2024-12-31"
|
||||
initial_cash = 1_000_000
|
||||
|
||||
logger.info(f"Universe: {universe} ({len(symbols)} stocks), Signal: {signal.name}")
|
||||
|
||||
# 1-2. Download daily data for the universe.
|
||||
data = download_batch(symbols, start, end)
|
||||
data = {s: df for s, df in data.items() if df is not None and not df.empty}
|
||||
logger.info(f"Downloaded {len(data)}/{len(symbols)} symbols")
|
||||
|
||||
# 3. Compute the signal per stock.
|
||||
signal_series: dict[str, pd.Series] = {}
|
||||
for sym, df in data.items():
|
||||
sig = signal.compute(df)
|
||||
sig.index = pd.to_datetime(df["date"])
|
||||
signal_series[sym] = sig
|
||||
|
||||
# 4. Cross-sectional IC at the matching forward horizon.
|
||||
signals_df = pd.DataFrame(signal_series)
|
||||
returns_df = _forward_returns(data, forward_horizon)
|
||||
signal_eval = evaluate_cross_sectional(signals_df, returns_df)
|
||||
|
||||
# 4b. Multi-horizon IC.
|
||||
horizon_evals = {
|
||||
h: evaluate_cross_sectional(signals_df, _forward_returns(data, h))
|
||||
for h in (1, 5, 20)
|
||||
}
|
||||
|
||||
# 5. Attach the signal column to each DataFrame and build feeds.
|
||||
config = BacktestConfig(
|
||||
symbols=list(data.keys()),
|
||||
start_date=start,
|
||||
end_date=end,
|
||||
initial_cash=initial_cash,
|
||||
sizer_percent=0.95,
|
||||
)
|
||||
runner = BacktestRunner(config)
|
||||
builder = RankEqualWeightBuilder(top_pct=0.2)
|
||||
for sym, df in data.items():
|
||||
df = df.copy()
|
||||
df["signal"] = signal.compute(df).values
|
||||
runner.add_signal_data(df, name=sym)
|
||||
|
||||
# 6. Run the multi-stock backtest.
|
||||
results = runner.run_prepared(AlphaStrategy, {"builder": builder})
|
||||
|
||||
# 7. Reports.
|
||||
artifacts = generate_report(
|
||||
results, signal_eval, output_dir="reports/", initial_cash=initial_cash
|
||||
)
|
||||
|
||||
# 7b. Dump signals and daily PnL.
|
||||
dump_signals(signals_df, dump_dir)
|
||||
dump_daily_pnl(results, dump_dir, initial_cash=initial_cash)
|
||||
|
||||
# 8. Print summary.
|
||||
print("\nSIGNAL IC")
|
||||
print("=" * 50)
|
||||
print(f"Universe: {universe} | Signal: {signal.name}")
|
||||
print(f"IC mean / std / IR: {signal_eval['ic_mean']:.4f} / "
|
||||
f"{signal_eval['ic_std']:.4f} / {signal_eval['ir']:.4f}")
|
||||
print(f"Rank IC mean / std / IR: {signal_eval['rank_ic_mean']:.4f} / "
|
||||
f"{signal_eval['rank_ic_std']:.4f} / {signal_eval['rank_ir']:.4f}")
|
||||
print(f"Hit rate: {signal_eval['hit_rate']:.2%}")
|
||||
print(f"Periods: {signal_eval['n_periods']}")
|
||||
|
||||
print("\nMULTI-HORIZON IC")
|
||||
print("=" * 50)
|
||||
print(f"{'Horizon':>8} {'Rank IC':>9} {'Rank IR':>9} {'Hit rate':>9} {'Periods':>8}")
|
||||
for h, ev in horizon_evals.items():
|
||||
print(f"{f'{h}d':>8} {ev['rank_ic_mean']:>9.4f} {ev['rank_ir']:>9.4f} "
|
||||
f"{ev['hit_rate']:>8.2%} {ev['n_periods']:>8}")
|
||||
|
||||
print(f"\nReports written to: {artifacts}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Chinese equity quant backtest")
|
||||
parser.add_argument("--universe", default="csi500", choices=["hs300", "csi500"])
|
||||
parser.add_argument("--signal", default="reversal_vol", choices=["reversal", "reversal_vol"])
|
||||
parser.add_argument("--dump-dir", default="results/")
|
||||
args = parser.parse_args()
|
||||
main(universe=args.universe, signal_name=args.signal, dump_dir=args.dump_dir)
|
||||
@@ -1,6 +0,0 @@
|
||||
"""Alpha signal abstractions."""
|
||||
from signals.base import AlphaSignal
|
||||
from signals.reversal import ReversalSignal
|
||||
from signals.reversal_vol import ReversalVolSignal
|
||||
|
||||
__all__ = ["AlphaSignal", "ReversalSignal", "ReversalVolSignal"]
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Base class for cross-sectional alpha signals."""
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class AlphaSignal(ABC):
|
||||
"""A signal that maps a single stock's OHLCV history to a per-bar score.
|
||||
|
||||
Higher scores indicate a stronger expected forward return. Implementations
|
||||
operate on one stock at a time; cross-sectional ranking happens downstream.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
"""Compute the signal for one stock.
|
||||
|
||||
Args:
|
||||
df: OHLCV DataFrame with at least a ``close`` column, ordered by date.
|
||||
|
||||
Returns:
|
||||
Signal series aligned to ``df`` (NaN where undefined).
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Human-readable signal identifier."""
|
||||
@@ -1,21 +0,0 @@
|
||||
"""Short-horizon momentum signal."""
|
||||
import pandas as pd
|
||||
|
||||
from signals.base import AlphaSignal
|
||||
|
||||
|
||||
class MomentumSignal(AlphaSignal):
|
||||
"""Positive trailing return: stocks that rose score high (momentum).
|
||||
|
||||
The signal is ``close.pct_change(lookback)`` — opposite of ReversalSignal.
|
||||
"""
|
||||
|
||||
def __init__(self, lookback: int = 5):
|
||||
self.lookback = lookback
|
||||
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
return df["close"].pct_change(self.lookback)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"momentum_{self.lookback}d"
|
||||
@@ -1,22 +0,0 @@
|
||||
"""Short-horizon reversal signal."""
|
||||
import pandas as pd
|
||||
|
||||
from signals.base import AlphaSignal
|
||||
|
||||
|
||||
class ReversalSignal(AlphaSignal):
|
||||
"""Negative trailing return: oversold stocks score high.
|
||||
|
||||
The signal is ``-close.pct_change(lookback)``, so a stock that fell over the
|
||||
lookback window gets a positive (bullish) score.
|
||||
"""
|
||||
|
||||
def __init__(self, lookback: int = 5):
|
||||
self.lookback = lookback
|
||||
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
return -df["close"].pct_change(self.lookback)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"reversal_{self.lookback}d"
|
||||
@@ -1,27 +0,0 @@
|
||||
"""Volatility-scaled short-horizon reversal signal."""
|
||||
import pandas as pd
|
||||
|
||||
from signals.base import AlphaSignal
|
||||
|
||||
|
||||
class ReversalVolSignal(AlphaSignal):
|
||||
"""Reversal score normalized by trailing volatility.
|
||||
|
||||
The raw reversal ``-close.pct_change(lookback)`` is divided by the rolling
|
||||
standard deviation of daily returns over ``vol_window``. Scaling by
|
||||
volatility damps the score of noisy, high-vol names so the signal favors
|
||||
oversold stocks whose move is large *relative* to their own volatility.
|
||||
"""
|
||||
|
||||
def __init__(self, lookback: int = 5, vol_window: int = 20):
|
||||
self.lookback = lookback
|
||||
self.vol_window = vol_window
|
||||
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
reversal = -df["close"].pct_change(self.lookback)
|
||||
vol = df["close"].pct_change().rolling(self.vol_window).std()
|
||||
return reversal / vol
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"reversal_vol_{self.lookback}d_{self.vol_window}d"
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Signal-driven multi-stock strategy."""
|
||||
import backtrader as bt
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class AlphaStrategy(bt.Strategy):
|
||||
"""Trade feeds based on precomputed ``signal`` line.
|
||||
|
||||
Supports two builder modes:
|
||||
- ThresholdBuilder: per-stock threshold (passed ``(signal_value, in_position)``)
|
||||
- RankEqualWeightBuilder: cross-sectional ranking (passed ``{symbol: signal}`` dict)
|
||||
"""
|
||||
|
||||
def __init__(self, builder):
|
||||
self.builder = builder
|
||||
|
||||
def next(self):
|
||||
# Collect all signals
|
||||
signals: dict[str, float] = {}
|
||||
for data in self.datas:
|
||||
sig = data.signal[0]
|
||||
if not pd.isna(sig):
|
||||
signals[data._name] = float(sig)
|
||||
|
||||
if not signals:
|
||||
return
|
||||
|
||||
# Detect builder type: if RankEqualWeightBuilder, use cross-sectional mode
|
||||
from portfolio.builder import RankEqualWeightBuilder
|
||||
if isinstance(self.builder, RankEqualWeightBuilder):
|
||||
actions = self.builder.build(signals)
|
||||
for data in self.datas:
|
||||
name = data._name
|
||||
if name not in actions:
|
||||
continue
|
||||
action = actions[name]
|
||||
if action.action == "buy":
|
||||
self.order_target_percent(data=data, target=action.size_pct)
|
||||
elif action.action == "sell":
|
||||
self.close(data=data)
|
||||
else:
|
||||
# Legacy per-stock ThresholdBuilder
|
||||
for data in self.datas:
|
||||
name = data._name
|
||||
if name not in signals:
|
||||
continue
|
||||
in_position = bool(self.getposition(data).size)
|
||||
action = self.builder.build(signals[name], in_position)
|
||||
if action.action == "buy":
|
||||
self.order_target_percent(data=data, target=action.size_pct)
|
||||
elif action.action == "sell":
|
||||
self.close(data=data)
|
||||
@@ -1,23 +0,0 @@
|
||||
"""Base strategy and example SMA crossover for Chinese equities."""
|
||||
import backtrader as bt
|
||||
|
||||
|
||||
class SmaCross(bt.Strategy):
|
||||
"""Simple SMA crossover strategy: buy when fast crosses above slow, sell when below."""
|
||||
|
||||
params = (
|
||||
("fast", 10),
|
||||
("slow", 30),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
self.fast_ma = bt.indicators.SMA(self.data.close, period=self.params.fast)
|
||||
self.slow_ma = bt.indicators.SMA(self.data.close, period=self.params.slow)
|
||||
self.crossover = bt.indicators.CrossOver(self.fast_ma, self.slow_ma)
|
||||
|
||||
def next(self):
|
||||
if not self.position:
|
||||
if self.crossover > 0: # fast crosses above slow
|
||||
self.buy()
|
||||
elif self.crossover < 0: # fast crosses below slow
|
||||
self.close()
|
||||
@@ -1,27 +0,0 @@
|
||||
import backtrader as bt
|
||||
|
||||
|
||||
class FiveDayReversal(bt.Strategy):
|
||||
"""Buy on 5-day oversold signal, sell on bounce or time stop."""
|
||||
|
||||
params = (
|
||||
("lookback", 5),
|
||||
("entry_threshold", -0.05), # buy when 5-day return < -5%
|
||||
("exit_threshold", 0.0), # sell when 1-day return > 0%
|
||||
("max_hold", 3), # max hold days
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
self.roc5 = bt.indicators.RateOfChange(self.data.close, period=self.params.lookback)
|
||||
self.hold_counter = 0
|
||||
|
||||
def next(self):
|
||||
if not self.position:
|
||||
if self.roc5[0] < self.params.entry_threshold: # RateOfChange returns a fraction, not %
|
||||
self.buy()
|
||||
self.hold_counter = 0
|
||||
else:
|
||||
self.hold_counter += 1
|
||||
roc1 = (self.data.close[0] / self.data.close[-1] - 1) * 100
|
||||
if roc1 > self.params.exit_threshold * 100 or self.hold_counter >= self.params.max_hold:
|
||||
self.close()
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Tests for pipeline alpha computation and combination (no network)."""
|
||||
import textwrap
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.compute import compute_alpha, evaluate_alpha
|
||||
from pipeline.alpha.registry import (
|
||||
available_alphas,
|
||||
get_alpha,
|
||||
load_alpha_module,
|
||||
register_alpha,
|
||||
)
|
||||
from pipeline.combo.combine import combine_alphas, _equal_weight
|
||||
from pipeline.common.schema import ALPHA_COLUMNS, COMBO_COLUMNS
|
||||
|
||||
|
||||
def _make_data(n_days: int = 30, symbols=("sh600000", "sz000001", "sh600519")) -> pd.DataFrame:
|
||||
"""Build a synthetic long-format DATA_COLUMNS frame with deterministic prices."""
|
||||
dates = pd.date_range("2024-01-01", periods=n_days)
|
||||
rng = np.random.default_rng(0)
|
||||
frames = []
|
||||
for i, sym in enumerate(symbols):
|
||||
# Distinct drift per symbol so the cross-section is non-degenerate.
|
||||
close = 100.0 + i * 5 + np.cumsum(rng.standard_normal(n_days))
|
||||
frames.append(pd.DataFrame({
|
||||
"symbol_id": sym,
|
||||
"symbol_name": sym,
|
||||
"date": dates,
|
||||
"open": close,
|
||||
"high": close,
|
||||
"low": close,
|
||||
"close": close,
|
||||
"volume": 1_000.0,
|
||||
"amount": 1_000.0 * close,
|
||||
}))
|
||||
return pd.concat(frames, ignore_index=True)
|
||||
|
||||
|
||||
def test_compute_alpha_schema_and_naming():
|
||||
alpha = compute_alpha(_make_data(), "rev5", "reversal", lookback=5)
|
||||
assert list(alpha.columns) == ALPHA_COLUMNS
|
||||
assert (alpha["alpha_name"] == "rev5").all()
|
||||
|
||||
|
||||
def test_reversal_sign_matches_negative_trailing_return():
|
||||
# Cross-sectional z-score preserves the sign relative to the cross-section,
|
||||
# so the stock with the most negative trailing return ranks highest.
|
||||
data = _make_data()
|
||||
alpha = compute_alpha(data, "rev5", "reversal", lookback=5)
|
||||
close = data.pivot_table(index="date", columns="symbol_id", values="close").sort_index()
|
||||
raw = -close.pct_change(5)
|
||||
last = raw.index[-1]
|
||||
expected_top = raw.loc[last].idxmax()
|
||||
got = alpha[alpha["date"] == last].set_index("symbol_id")["weight"].idxmax()
|
||||
assert got == expected_top
|
||||
|
||||
|
||||
def test_weights_are_cross_sectional_zscore():
|
||||
# Each date's weights are a z-score, so the per-date mean is ~0.
|
||||
alpha = compute_alpha(_make_data(), "rev5", "reversal", lookback=5)
|
||||
per_date_mean = alpha.groupby("date")["weight"].mean().abs()
|
||||
assert (per_date_mean < 1e-9).all()
|
||||
|
||||
|
||||
def test_evaluate_alpha_keys():
|
||||
data = _make_data()
|
||||
alpha = compute_alpha(data, "rev5", "reversal", lookback=5)
|
||||
metrics = evaluate_alpha(alpha, data)
|
||||
for key in ("cumulative_return", "sharpe_annual", "turnover_annual",
|
||||
"max_drawdown", "hit_rate", "n_dates"):
|
||||
assert key in metrics
|
||||
|
||||
|
||||
def test_equal_weight_is_mean_of_alphas():
|
||||
data = _make_data()
|
||||
a = compute_alpha(data, "rev", "reversal", lookback=5)
|
||||
b = compute_alpha(data, "mom", "momentum", lookback=5)
|
||||
combo = _equal_weight([a, b])
|
||||
# reversal = -momentum before z-scoring, but after independent per-date
|
||||
# z-scoring they are exact negatives, so the equal-weight mean is ~0.
|
||||
assert combo["weight"].abs().max() < 1e-9
|
||||
|
||||
|
||||
def test_combine_alphas_schema(tmp_path):
|
||||
data = _make_data()
|
||||
a_path = tmp_path / "a.pq"
|
||||
b_path = tmp_path / "b.pq"
|
||||
compute_alpha(data, "rev", "reversal", lookback=5).to_parquet(a_path, index=False)
|
||||
compute_alpha(data, "revvol", "reversal_vol", lookback=5, vol_window=10).to_parquet(b_path, index=False)
|
||||
combo = combine_alphas([str(a_path), str(b_path)], "eq", method="equal_weight")
|
||||
assert list(combo.columns) == COMBO_COLUMNS
|
||||
assert (combo["combo_name"] == "eq").all()
|
||||
|
||||
|
||||
# --- registry / factory -----------------------------------------------------
|
||||
|
||||
def test_builtins_are_registered():
|
||||
assert {"reversal", "reversal_vol", "momentum"} <= set(available_alphas())
|
||||
|
||||
|
||||
def test_get_alpha_filters_unaccepted_params():
|
||||
# reversal only accepts lookback; passing vol_window too must not error.
|
||||
alpha = get_alpha("reversal", lookback=7, vol_window=99)
|
||||
assert alpha.name == "reversal"
|
||||
assert alpha.lookback == 7
|
||||
assert not hasattr(alpha, "vol_window")
|
||||
|
||||
|
||||
def test_get_alpha_unknown_raises():
|
||||
with pytest.raises(KeyError):
|
||||
get_alpha("does_not_exist")
|
||||
|
||||
|
||||
def test_register_duplicate_name_raises():
|
||||
available_alphas() # ensure built-ins loaded
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
@register_alpha
|
||||
class Dup(BaseAlpha):
|
||||
name = "reversal"
|
||||
|
||||
def signal(self, close):
|
||||
return close
|
||||
|
||||
|
||||
def test_register_rejects_non_basealpha():
|
||||
with pytest.raises(TypeError):
|
||||
register_alpha(object) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# --- base class --------------------------------------------------------------
|
||||
|
||||
def test_to_weights_are_per_date_zscore():
|
||||
class _Const(BaseAlpha):
|
||||
name = "_const_test"
|
||||
|
||||
def signal(self, close):
|
||||
return close # arbitrary finite signal
|
||||
|
||||
close = _make_data().pivot_table(index="date", columns="symbol_id", values="close")
|
||||
weights = _Const().weights(close.sort_index())
|
||||
# Each date demeaned to ~0.
|
||||
assert (weights.mean(axis=1).abs() < 1e-9).all()
|
||||
|
||||
|
||||
# --- external plugin loading -------------------------------------------------
|
||||
|
||||
def test_load_external_alpha_module(tmp_path):
|
||||
module_path = tmp_path / "my_external_alpha.py"
|
||||
module_path.write_text(textwrap.dedent('''
|
||||
import pandas as pd
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
@register_alpha
|
||||
class ExternalDemoAlpha(BaseAlpha):
|
||||
name = "external_demo"
|
||||
|
||||
def __init__(self, span: int = 3):
|
||||
self.span = span
|
||||
|
||||
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
||||
return -close.pct_change(self.span)
|
||||
'''))
|
||||
|
||||
load_alpha_module(str(module_path))
|
||||
assert "external_demo" in available_alphas()
|
||||
|
||||
# The factory forwards the external alpha's own param (`span`).
|
||||
instance = get_alpha("external_demo", span=4, lookback=99)
|
||||
assert instance.span == 4
|
||||
|
||||
# And it works end-to-end through compute_alpha.
|
||||
result = compute_alpha(_make_data(), "ext", "external_demo", span=4)
|
||||
assert list(result.columns) == ALPHA_COLUMNS
|
||||
assert (result["alpha_name"] == "ext").all()
|
||||
|
||||
@@ -11,8 +11,22 @@ def test_download_single_stock():
|
||||
assert df["close"].notna().all()
|
||||
|
||||
|
||||
def test_download_baostock_fallback():
|
||||
"""Test baostock works as secondary source."""
|
||||
def test_download_baostock_primary():
|
||||
"""baostock is the primary source for 'auto'."""
|
||||
df = download_daily("sz000001", "2024-06-01", "2024-06-15", source="baostock")
|
||||
assert df is not None
|
||||
assert len(df) > 0
|
||||
|
||||
|
||||
def test_download_akshare_fallback():
|
||||
"""akshare works as the secondary source when reachable.
|
||||
|
||||
akshare is the fallback precisely because it is unreliable on some
|
||||
networks; skip rather than fail when it cannot be reached.
|
||||
"""
|
||||
try:
|
||||
df = download_daily("sh600000", "2024-01-01", "2024-01-31", source="akshare")
|
||||
except RuntimeError as e:
|
||||
pytest.skip(f"akshare unreachable on this network: {e}")
|
||||
assert df is not None
|
||||
assert len(df) > 0
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
"""Tests for cross-sectional IC evaluation."""
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from eval.metrics import evaluate_cross_sectional
|
||||
|
||||
|
||||
def test_cross_sectional_keys_present():
|
||||
dates = pd.date_range("2024-01-01", periods=10)
|
||||
cols = ["a", "b", "c"]
|
||||
rng = np.random.default_rng(0)
|
||||
signals = pd.DataFrame(rng.standard_normal((10, 3)), index=dates, columns=cols)
|
||||
returns = pd.DataFrame(rng.standard_normal((10, 3)), index=dates, columns=cols)
|
||||
res = evaluate_cross_sectional(signals, returns)
|
||||
for key in (
|
||||
"ic_mean", "ic_std", "ir", "rank_ic_mean", "rank_ic_std",
|
||||
"rank_ir", "hit_rate", "n_periods",
|
||||
):
|
||||
assert key in res
|
||||
|
||||
|
||||
def test_perfect_signal_has_positive_rank_ic():
|
||||
# When the signal equals next-period returns, rank IC should be ~1 each day.
|
||||
dates = pd.date_range("2024-01-01", periods=8)
|
||||
cols = ["a", "b", "c"]
|
||||
rng = np.random.default_rng(42)
|
||||
returns = pd.DataFrame(rng.standard_normal((8, 3)), index=dates, columns=cols)
|
||||
signals = returns.copy() # perfect foresight
|
||||
res = evaluate_cross_sectional(signals, returns)
|
||||
assert res["rank_ic_mean"] > 0.99
|
||||
assert res["hit_rate"] == 1.0
|
||||
assert res["n_periods"] == 8
|
||||
|
||||
|
||||
def test_inverted_signal_has_negative_rank_ic():
|
||||
dates = pd.date_range("2024-01-01", periods=6)
|
||||
cols = ["a", "b", "c"]
|
||||
rng = np.random.default_rng(7)
|
||||
returns = pd.DataFrame(rng.standard_normal((6, 3)), index=dates, columns=cols)
|
||||
signals = -returns # perfectly wrong
|
||||
res = evaluate_cross_sectional(signals, returns)
|
||||
assert res["rank_ic_mean"] < -0.99
|
||||
|
||||
|
||||
def test_single_stock_falls_back_to_rolling():
|
||||
dates = pd.date_range("2024-01-01", periods=40)
|
||||
rng = np.random.default_rng(1)
|
||||
signals = pd.DataFrame({"a": rng.standard_normal(40)}, index=dates)
|
||||
returns = pd.DataFrame({"a": rng.standard_normal(40)}, index=dates)
|
||||
res = evaluate_cross_sectional(signals, returns)
|
||||
# Rolling fallback still yields the standard metric keys.
|
||||
assert "rank_ic_mean" in res
|
||||
assert res["n_periods"] > 0
|
||||
@@ -1,21 +0,0 @@
|
||||
import pytest
|
||||
from backtest.config import BacktestConfig
|
||||
from backtest.runner import BacktestRunner
|
||||
from strategies.reversal import FiveDayReversal
|
||||
|
||||
|
||||
def test_reversal_smoke():
|
||||
"""Smoke test: run a minimal reversal backtest and check results exist."""
|
||||
config = BacktestConfig(
|
||||
symbols=["sh600000"],
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-03-31",
|
||||
initial_cash=100_000,
|
||||
)
|
||||
runner = BacktestRunner(config)
|
||||
results = runner.run(FiveDayReversal)
|
||||
assert results is not None
|
||||
assert len(results) == 1
|
||||
# Check analyzers exist
|
||||
sharpe = results[0].analyzers.sharpe.get_analysis()
|
||||
assert "sharperatio" in sharpe
|
||||
@@ -1,21 +0,0 @@
|
||||
import pytest
|
||||
from backtest.config import BacktestConfig
|
||||
from backtest.runner import BacktestRunner
|
||||
from strategies.base import SmaCross
|
||||
|
||||
|
||||
def test_backtest_smoke():
|
||||
"""Smoke test: run a minimal backtest and check results exist."""
|
||||
config = BacktestConfig(
|
||||
symbols=["sh600000"],
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-03-31",
|
||||
initial_cash=100_000,
|
||||
)
|
||||
runner = BacktestRunner(config)
|
||||
results = runner.run(SmaCross)
|
||||
assert results is not None
|
||||
assert len(results) == 1
|
||||
# Check analyzers exist
|
||||
sharpe = results[0].analyzers.sharpe.get_analysis()
|
||||
assert "sharperatio" in sharpe
|
||||
@@ -1,38 +0,0 @@
|
||||
"""Tests for alpha signal computation."""
|
||||
import pandas as pd
|
||||
|
||||
from signals.reversal import ReversalSignal
|
||||
|
||||
|
||||
def _make_df(closes):
|
||||
return pd.DataFrame({"close": closes})
|
||||
|
||||
|
||||
def test_reversal_name():
|
||||
assert ReversalSignal(lookback=5).name == "reversal_5d"
|
||||
assert ReversalSignal(lookback=10).name == "reversal_10d"
|
||||
|
||||
|
||||
def test_reversal_is_negative_trailing_return():
|
||||
# Monotonically rising prices -> negative (bearish) reversal signal.
|
||||
df = _make_df([10.0, 11.0, 12.0, 13.0, 14.0, 15.0])
|
||||
sig = ReversalSignal(lookback=5).compute(df)
|
||||
# First 5 values are NaN (insufficient lookback).
|
||||
assert sig.iloc[:5].isna().all()
|
||||
# 15/10 - 1 = 0.5 return -> signal = -0.5
|
||||
assert abs(sig.iloc[5] - (-0.5)) < 1e-9
|
||||
|
||||
|
||||
def test_reversal_oversold_is_positive():
|
||||
# Falling prices -> positive (bullish) reversal signal.
|
||||
df = _make_df([20.0, 18.0, 16.0, 14.0, 12.0, 10.0])
|
||||
sig = ReversalSignal(lookback=5).compute(df)
|
||||
assert sig.iloc[5] > 0
|
||||
# 10/20 - 1 = -0.5 -> signal = +0.5
|
||||
assert abs(sig.iloc[5] - 0.5) < 1e-9
|
||||
|
||||
|
||||
def test_reversal_output_length_matches_input():
|
||||
df = _make_df([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0])
|
||||
sig = ReversalSignal(lookback=3).compute(df)
|
||||
assert len(sig) == len(df)
|
||||
Reference in New Issue
Block a user