Add minute bar feature pipeline
This commit is contained in:
Yuxuan Yan
@@ -4,6 +4,7 @@
|
||||
Phases:
|
||||
data — Download daily bars to parquet
|
||||
alpha — Compute alpha weights from data
|
||||
feature — Compute daily features from minute bars
|
||||
combo — Combine alphas into a single weight
|
||||
portfolio — Build tradable positions and simulate execution
|
||||
"""
|
||||
@@ -14,6 +15,7 @@ import click
|
||||
|
||||
from pipeline.data.cli import data
|
||||
from pipeline.alpha.cli import alpha
|
||||
from pipeline.features.cli import feature
|
||||
from pipeline.combo.cli import combo
|
||||
from pipeline.portfolio.cli import portfolio
|
||||
from tools.pqcat import pqcat
|
||||
@@ -40,6 +42,7 @@ def cli(log_level):
|
||||
|
||||
cli.add_command(data)
|
||||
cli.add_command(alpha)
|
||||
cli.add_command(feature)
|
||||
cli.add_command(combo)
|
||||
cli.add_command(portfolio)
|
||||
cli.add_command(pqcat)
|
||||
|
||||
@@ -31,11 +31,68 @@ _BATCH_COLUMNS = [
|
||||
"peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
|
||||
]
|
||||
|
||||
# Raw Baostock minute bars. The ``time`` field is usually compact
|
||||
# YYYYMMDDHHMMSSmmm; parsing below also tolerates HH:MM:SS strings in tests.
|
||||
_MINUTE_FIELDS = "date,time,code,open,high,low,close,volume,amount,adjustflag"
|
||||
_MINUTE_NUMERIC = ["open", "high", "low", "close", "volume", "amount"]
|
||||
_MINUTE_COLUMNS = [
|
||||
"symbol", "datetime", "date", "time", "frequency",
|
||||
"open", "high", "low", "close", "volume", "amount", "vwap", "adjustflag",
|
||||
]
|
||||
_MINUTE_FREQUENCIES = {"5", "15", "30", "60"}
|
||||
|
||||
|
||||
class _SessionLost(Exception):
|
||||
"""baostock reported the session was dropped (``用户未登录``)."""
|
||||
|
||||
|
||||
def _normalize_minute_frequency(frequency: str | int) -> tuple[str, str]:
|
||||
"""Return Baostock frequency and partition label for a minute interval."""
|
||||
raw = str(frequency).strip().lower()
|
||||
if raw.endswith("m"):
|
||||
raw = raw[:-1]
|
||||
if raw not in _MINUTE_FREQUENCIES:
|
||||
raise ValueError(
|
||||
f"Unsupported minute frequency '{frequency}'. "
|
||||
f"Expected one of {sorted(_MINUTE_FREQUENCIES)} minutes."
|
||||
)
|
||||
return raw, f"{raw}m"
|
||||
|
||||
|
||||
def _parse_minute_datetime(date: pd.Series, time: pd.Series) -> pd.Series:
|
||||
"""Parse Baostock minute timestamps into pandas datetimes."""
|
||||
date_dt = pd.to_datetime(date, errors="coerce")
|
||||
date_compact = date_dt.dt.strftime("%Y%m%d")
|
||||
time_text = time.astype(str).str.strip()
|
||||
time_digits = time_text.str.replace(r"\D", "", regex=True)
|
||||
|
||||
full_digits = time_digits.str.slice(0, 14)
|
||||
from_full = pd.to_datetime(full_digits, format="%Y%m%d%H%M%S", errors="coerce")
|
||||
|
||||
from_short = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]")
|
||||
short_time = time_digits.str.len().between(1, 6)
|
||||
if short_time.any():
|
||||
short_digits = (
|
||||
time_digits.loc[short_time]
|
||||
.str.pad(6, side="right", fillchar="0")
|
||||
.str.slice(0, 6)
|
||||
)
|
||||
from_short.loc[short_time] = pd.to_datetime(
|
||||
date_compact.loc[short_time] + short_digits,
|
||||
format="%Y%m%d%H%M%S",
|
||||
errors="coerce",
|
||||
)
|
||||
|
||||
from_text = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]")
|
||||
text_time = time_text.str.contains(":", regex=False)
|
||||
if text_time.any():
|
||||
from_text.loc[text_time] = pd.to_datetime(
|
||||
date.astype(str).loc[text_time] + " " + time_text.loc[text_time],
|
||||
errors="coerce",
|
||||
)
|
||||
return from_full.fillna(from_short).fillna(from_text)
|
||||
|
||||
|
||||
def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
|
||||
"""Download daily bars from akshare. Returns DataFrame with OHLCV columns."""
|
||||
try:
|
||||
@@ -239,3 +296,104 @@ def download_daily_batch(
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def download_minute_batch(
|
||||
symbols: Iterable[str],
|
||||
start: str,
|
||||
end: str,
|
||||
frequency: str | int = 5,
|
||||
relogin_every: int = 200,
|
||||
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
|
||||
"""Download raw Baostock minute bars for many symbols.
|
||||
|
||||
Minute bars are intentionally unadjusted (`adjustflag='3'`) because the
|
||||
output is raw intraday market data for downstream feature aggregation, not a
|
||||
tradable daily price series.
|
||||
|
||||
Args:
|
||||
symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
|
||||
start, end: ``YYYY-MM-DD`` bounds.
|
||||
frequency: Baostock minute frequency. ``5``/``"5"``/``"5m"`` all mean
|
||||
5-minute bars.
|
||||
relogin_every: Proactively refresh the baostock session every N symbols.
|
||||
|
||||
Yields:
|
||||
``(symbol, df)`` where ``df`` has raw minute bars or ``None`` when no
|
||||
data is available.
|
||||
"""
|
||||
query_frequency, frequency_label = _normalize_minute_frequency(frequency)
|
||||
adjustflag = _BAOSTOCK_ADJUST["none"]
|
||||
|
||||
def _relogin() -> None:
|
||||
try:
|
||||
bs.logout()
|
||||
except Exception:
|
||||
pass
|
||||
bs.login()
|
||||
|
||||
def _fetch(symbol: str) -> Optional[pd.DataFrame]:
|
||||
"""One Baostock minute query; returns df, None, or raises _SessionLost."""
|
||||
code = f"{symbol[:2]}.{symbol[2:]}"
|
||||
rs = bs.query_history_k_data_plus(
|
||||
code=code,
|
||||
fields=_MINUTE_FIELDS,
|
||||
start_date=start,
|
||||
end_date=end,
|
||||
frequency=query_frequency,
|
||||
adjustflag=adjustflag,
|
||||
)
|
||||
if rs.error_code != "0":
|
||||
if "未登录" in (rs.error_msg or ""):
|
||||
raise _SessionLost(rs.error_msg)
|
||||
logger.warning("baostock minute error for %s: %s", symbol, rs.error_msg)
|
||||
return None
|
||||
rows = []
|
||||
while rs.next():
|
||||
rows.append(rs.get_row_data())
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
df = pd.DataFrame(rows, columns=_MINUTE_FIELDS.split(","))
|
||||
df[_MINUTE_NUMERIC] = df[_MINUTE_NUMERIC].apply(pd.to_numeric, errors="coerce")
|
||||
df["datetime"] = _parse_minute_datetime(df["date"], df["time"])
|
||||
bad_timestamps = df["datetime"].isna()
|
||||
if bad_timestamps.any():
|
||||
raise ValueError(
|
||||
f"Could not parse {int(bad_timestamps.sum())} minute timestamp(s)"
|
||||
)
|
||||
df["date"] = pd.to_datetime(df["date"], errors="coerce")
|
||||
df["time"] = df["datetime"].dt.strftime("%H:%M:%S")
|
||||
df["frequency"] = frequency_label
|
||||
df["vwap"] = (df["amount"] / df["volume"]).where(df["volume"] > 0)
|
||||
df["symbol"] = symbol
|
||||
return df[_MINUTE_COLUMNS].sort_values("datetime").reset_index(drop=True)
|
||||
|
||||
bs.login()
|
||||
try:
|
||||
for i, symbol in enumerate(symbols):
|
||||
if i and relogin_every and i % relogin_every == 0:
|
||||
_relogin()
|
||||
|
||||
df: Optional[pd.DataFrame] = None
|
||||
for attempt in (1, 2):
|
||||
try:
|
||||
df = _fetch(symbol)
|
||||
break
|
||||
except _SessionLost:
|
||||
if attempt == 1:
|
||||
_relogin()
|
||||
continue
|
||||
logger.warning("baostock minute session lost for %s after relogin", symbol)
|
||||
except Exception as e:
|
||||
logger.warning("baostock minute download failed for %s: %s", symbol, e)
|
||||
break
|
||||
|
||||
if df is not None and not df.empty:
|
||||
yield symbol, df
|
||||
else:
|
||||
yield symbol, None
|
||||
finally:
|
||||
try:
|
||||
bs.logout()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
# Minute Bar Data Notes
|
||||
|
||||
The minute-bar path downloads raw Baostock intraday bars and stores them as a
|
||||
Hive-partitioned dataset:
|
||||
|
||||
```bash
|
||||
uv run python cli.py data download-minute \
|
||||
--universe sh600000 \
|
||||
--start-date 2024-01-02 --end-date 2024-01-05 \
|
||||
--frequency 5
|
||||
```
|
||||
|
||||
The default layout is:
|
||||
|
||||
```text
|
||||
data/minute_bars/{universe}/frequency=5m/month=YYYY-MM/*.pq
|
||||
```
|
||||
|
||||
Feature plugins can aggregate those bars to daily `symbol_id,date` feature
|
||||
files, for example:
|
||||
|
||||
```bash
|
||||
uv run python cli.py feature compute \
|
||||
--minute-path data/minute_bars/sh600000 \
|
||||
--daily-path data/daily_bars/sh600000 \
|
||||
--feature-type minute_daily_summary \
|
||||
--feature-name minute_summary
|
||||
```
|
||||
|
||||
## Daily vs Minute Reconciliation
|
||||
|
||||
Baostock's daily raw bars and 5-minute raw bars are close, but they should not
|
||||
be treated as perfectly reconstructible from each other.
|
||||
|
||||
When checking consistency, compare daily raw bars (`data download --adjust none`)
|
||||
against minute bars on the same raw price scale. The minute aggregation should
|
||||
use:
|
||||
|
||||
- `open`: first minute open
|
||||
- `high`: max minute high
|
||||
- `low`: min minute low
|
||||
- `close`: last minute close
|
||||
- `volume`: sum minute volume
|
||||
- `amount`: sum minute amount
|
||||
- `vwap`: `sum(amount) / sum(volume)`
|
||||
|
||||
In a sanity check for `sh600000` from `2024-01-02` through `2024-01-05`, Baostock
|
||||
returned 4 daily rows and 192 5-minute bars, exactly 48 bars per day. Open, low,
|
||||
and close matched daily exactly on all 4 days. High matched on 3 of 4 days; on
|
||||
`2024-01-04`, the daily high was `6.67` while the max 5-minute high was `6.66`.
|
||||
Minute-summed volume and amount were higher than daily by roughly `0.16%` to
|
||||
`1.23%`. VWAP remained very close, with max relative difference around
|
||||
`0.0043%`.
|
||||
|
||||
This appears to be a source-level Baostock reconciliation caveat, not a parser
|
||||
or ordering issue: the minute bars covered the regular `09:35:00` through
|
||||
`15:00:00` range and sorted correctly by timestamp.
|
||||
|
||||
Practical guidance:
|
||||
|
||||
- Use tolerance-based daily-vs-minute checks; do not require exact equality for
|
||||
high, volume, or amount.
|
||||
- Expect open/close alignment to be a stronger sanity check than exact volume
|
||||
reconstruction.
|
||||
- Use minute-derived values as separate daily features, not as replacements for
|
||||
the canonical daily bar dataset unless a strategy explicitly wants that
|
||||
source convention.
|
||||
+19
-5
@@ -5,7 +5,7 @@ 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
|
||||
from abc import ABC
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
@@ -15,15 +15,14 @@ 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.
|
||||
key) and implement either :meth:`signal` or :meth:`signal_from_data`.
|
||||
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.
|
||||
|
||||
@@ -34,6 +33,21 @@ class BaseAlpha(ABC):
|
||||
A wide DataFrame aligned to ``close`` where higher values indicate a
|
||||
stronger long. Use NaN where the signal is undefined.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} must implement signal() or signal_from_data()"
|
||||
)
|
||||
|
||||
def signal_from_data(
|
||||
self,
|
||||
data: pd.DataFrame,
|
||||
close: pd.DataFrame,
|
||||
) -> pd.DataFrame:
|
||||
"""Compute the raw signal from long daily data plus wide closes.
|
||||
|
||||
Feature-aware alphas can override this to pivot joined feature columns
|
||||
from ``data``. The default preserves the existing close-only alpha API.
|
||||
"""
|
||||
return self.signal(close)
|
||||
|
||||
def to_weights(self, signal: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Cross-sectionally z-score a signal into signed position weights.
|
||||
|
||||
@@ -56,6 +56,10 @@ def list_(alpha_modules):
|
||||
@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(
|
||||
"--feature-path", "feature_paths", multiple=True,
|
||||
help="Daily feature parquet file/dataset to left-join on symbol_id,date (repeatable)",
|
||||
)
|
||||
@click.option(
|
||||
"--alpha-module", "alpha_modules", multiple=True,
|
||||
help="External module(s) to import so their alphas register (dotted path or .py file)",
|
||||
@@ -74,7 +78,7 @@ def list_(alpha_modules):
|
||||
help="Most-liquid names kept per date when --liquid-universe is set",
|
||||
)
|
||||
def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window,
|
||||
alpha_modules, extra_params, liquid_universe, universe_top_n):
|
||||
feature_paths, alpha_modules, extra_params, liquid_universe, universe_top_n):
|
||||
"""Compute one alpha from raw data and save as parquet."""
|
||||
for spec in alpha_modules:
|
||||
load_alpha_module(spec)
|
||||
@@ -100,6 +104,7 @@ def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window,
|
||||
alpha_name=alpha_name,
|
||||
alpha_type=alpha_type,
|
||||
universe=universe,
|
||||
feature_paths=feature_paths,
|
||||
**params,
|
||||
)
|
||||
|
||||
|
||||
@@ -7,12 +7,15 @@ through :mod:`pipeline.alpha.registry`.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.alpha.registry import get_alpha
|
||||
from pipeline.common.schema import ALPHA_COLUMNS
|
||||
from pipeline.features.compute import FEATURE_KEY_COLUMNS, read_feature_frames, validate_feature_frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -33,6 +36,38 @@ def _pivot_open(df: pd.DataFrame) -> pd.DataFrame:
|
||||
return pivot.sort_index()
|
||||
|
||||
|
||||
def join_feature_frames(
|
||||
data: pd.DataFrame,
|
||||
feature_frames: Iterable[pd.DataFrame],
|
||||
) -> pd.DataFrame:
|
||||
"""Left-join validated daily feature frames onto long daily data."""
|
||||
out = data.copy()
|
||||
out["date"] = pd.to_datetime(out["date"])
|
||||
existing = set(out.columns)
|
||||
joined_cols: list[str] = []
|
||||
|
||||
for frame in feature_frames:
|
||||
features = validate_feature_frame(frame)
|
||||
feature_cols = [col for col in features.columns if col not in FEATURE_KEY_COLUMNS]
|
||||
overlap = sorted(existing.intersection(feature_cols))
|
||||
if overlap:
|
||||
raise ValueError(
|
||||
f"Feature columns conflict with existing daily data columns: {overlap}"
|
||||
)
|
||||
out = out.merge(
|
||||
features,
|
||||
on=FEATURE_KEY_COLUMNS,
|
||||
how="left",
|
||||
validate="many_to_one",
|
||||
)
|
||||
existing.update(feature_cols)
|
||||
joined_cols.extend(feature_cols)
|
||||
|
||||
if joined_cols:
|
||||
logger.info("Joined feature columns into daily data: %s", joined_cols)
|
||||
return out
|
||||
|
||||
|
||||
def _forward_open_to_open_returns(open_: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Return earned by a close-formed signal after next-open execution.
|
||||
|
||||
@@ -105,6 +140,8 @@ def compute_alpha(
|
||||
alpha_name: str,
|
||||
alpha_type: str,
|
||||
universe: dict | None = None,
|
||||
feature_paths: Iterable[str | Path] | None = None,
|
||||
feature_frames: Iterable[pd.DataFrame] | None = None,
|
||||
**params,
|
||||
) -> pd.DataFrame:
|
||||
"""Compute alpha weights from raw data.
|
||||
@@ -118,6 +155,10 @@ def compute_alpha(
|
||||
:func:`investable_universe_mask`) *before* it is turned into
|
||||
weights, so unheld names get weight 0. Keys are forwarded as keyword
|
||||
arguments to :func:`investable_universe_mask`.
|
||||
feature_paths: Optional parquet files/datasets keyed by ``symbol_id``
|
||||
and ``date``. Their numeric feature columns are left-joined onto
|
||||
``data`` before alpha logic runs.
|
||||
feature_frames: Optional in-memory feature frames with the same schema.
|
||||
**params: Constructor parameters for the alpha (e.g. ``lookback``,
|
||||
``vol_window``). Only the params the alpha's ``__init__`` accepts are
|
||||
used; extras are ignored.
|
||||
@@ -128,12 +169,20 @@ def compute_alpha(
|
||||
Raises:
|
||||
KeyError: If ``alpha_type`` is not registered.
|
||||
"""
|
||||
feature_inputs: list[pd.DataFrame] = []
|
||||
if feature_paths:
|
||||
feature_inputs.extend(read_feature_frames(feature_paths))
|
||||
if feature_frames:
|
||||
feature_inputs.extend(feature_frames)
|
||||
if feature_inputs:
|
||||
data = join_feature_frames(data, feature_inputs)
|
||||
|
||||
alpha = get_alpha(alpha_type, **params)
|
||||
close = _pivot_close(data)
|
||||
signal = alpha.signal_from_data(data, close)
|
||||
if universe is None:
|
||||
weights = alpha.weights(close)
|
||||
weights = alpha.to_weights(signal)
|
||||
else:
|
||||
signal = alpha.signal(close)
|
||||
mask = investable_universe_mask(data, signal, **universe)
|
||||
weights = alpha.to_weights(signal.where(mask))
|
||||
|
||||
|
||||
@@ -26,6 +26,24 @@ DATA_COLUMNS: Final[list[str]] = [
|
||||
"pcfNcfTTM", # float64: P/CF (net cash flow, TTM)
|
||||
]
|
||||
|
||||
# Required columns for raw intraday minute bar parquet files.
|
||||
MINUTE_BAR_COLUMNS: Final[list[str]] = [
|
||||
"symbol_id", # str: internal code like 'sh600000'
|
||||
"symbol_name", # str: stock name like '浦发银行'
|
||||
"datetime", # datetime64: intraday bar timestamp
|
||||
"date", # date component, aligned with daily DATA_COLUMNS date
|
||||
"time", # str: HH:MM:SS bar time
|
||||
"frequency", # str: e.g. '5m'
|
||||
"open", # float64
|
||||
"high", # float64
|
||||
"low", # float64
|
||||
"close", # float64
|
||||
"volume", # float64 (shares)
|
||||
"amount", # float64 (turnover in yuan, raw/unadjusted)
|
||||
"vwap", # float64: amount / volume
|
||||
"adjustflag", # str: baostock adjustment flag; '3' for raw/unadjusted
|
||||
]
|
||||
|
||||
# Required columns for alpha parquet files.
|
||||
# Alphas are position WEIGHTS: positive=long, negative=short.
|
||||
ALPHA_COLUMNS: Final[list[str]] = [
|
||||
|
||||
+35
-1
@@ -3,7 +3,7 @@
|
||||
import click
|
||||
from datetime import date
|
||||
|
||||
from pipeline.data.downloader import download_universe
|
||||
from pipeline.data.downloader import download_minute_universe, download_universe
|
||||
|
||||
|
||||
@click.group(name="data")
|
||||
@@ -42,3 +42,37 @@ def download(universe, start_date, end_date, output_dir, symbols, chunk_size, ad
|
||||
f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}"
|
||||
)
|
||||
click.echo(f"Dataset: {stats['dataset_path']}")
|
||||
|
||||
|
||||
@data.command("download-minute")
|
||||
@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/minute_bars", help="Root for the partitioned dataset")
|
||||
@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
|
||||
@click.option("--chunk-size", default=100, type=int, help="Symbols per durability flush")
|
||||
@click.option("--frequency", default="5", help="Minute frequency: 5, 15, 30, or 60")
|
||||
def download_minute(universe, start_date, end_date, output_dir, symbols, chunk_size, frequency):
|
||||
"""Download raw Baostock minute bars into a partitioned parquet dataset.
|
||||
|
||||
Writes ``{output_dir}/{universe}/frequency=5m/month=YYYY-MM/*.pq`` for the
|
||||
default 5-minute frequency.
|
||||
"""
|
||||
stats = download_minute_universe(
|
||||
universe=universe,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
output_dir=output_dir,
|
||||
max_symbols=symbols,
|
||||
chunk_size=chunk_size,
|
||||
frequency=frequency,
|
||||
)
|
||||
click.echo(
|
||||
f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
|
||||
f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}, "
|
||||
f"frequency={stats['frequency']}"
|
||||
)
|
||||
click.echo(f"Dataset: {stats['dataset_path']}")
|
||||
|
||||
+126
-2
@@ -11,9 +11,13 @@ 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.downloader import (
|
||||
_normalize_minute_frequency,
|
||||
download_daily_batch,
|
||||
download_minute_batch,
|
||||
)
|
||||
from data.universe import get_all_stocks, get_hs300_stocks, get_zz500_stocks
|
||||
from pipeline.common.schema import DATA_COLUMNS
|
||||
from pipeline.common.schema import DATA_COLUMNS, MINUTE_BAR_COLUMNS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -89,6 +93,25 @@ def _write_month_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: s
|
||||
)
|
||||
|
||||
|
||||
def _write_minute_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: str) -> None:
|
||||
"""Append rows to a Hive-partitioned minute dataset.
|
||||
|
||||
Layout: ``frequency=5m/month=YYYY-MM/*.pq``.
|
||||
"""
|
||||
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=["frequency", "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",
|
||||
@@ -177,3 +200,104 @@ def download_universe(
|
||||
"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()),
|
||||
}
|
||||
|
||||
|
||||
def download_minute_universe(
|
||||
universe: str = "csi500",
|
||||
start_date: str = "2017-01-01",
|
||||
end_date: str = "2026-12-31",
|
||||
output_dir: str = "data/minute_bars",
|
||||
max_symbols: int = 0,
|
||||
chunk_size: int = 100,
|
||||
frequency: str | int = 5,
|
||||
) -> dict:
|
||||
"""Download raw minute bars into a frequency/month-partitioned dataset.
|
||||
|
||||
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}/frequency=5m/month=YYYY-MM``
|
||||
is written.
|
||||
max_symbols: Cap on symbols (0 = all).
|
||||
chunk_size: Symbols per durability flush.
|
||||
frequency: Minute interval. ``5``/``"5"``/``"5m"`` are 5-minute bars.
|
||||
|
||||
Returns:
|
||||
Stats dict with dataset path, row count, symbol count, date range, and
|
||||
frequency label.
|
||||
"""
|
||||
_, frequency_label = _normalize_minute_frequency(frequency)
|
||||
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(
|
||||
"Minute universe %s: %d symbols, %s → %s, frequency=%s",
|
||||
universe,
|
||||
n_requested,
|
||||
start_date,
|
||||
end_date,
|
||||
frequency,
|
||||
)
|
||||
|
||||
base_dir = Path(output_dir) / universe
|
||||
target_frequency_dir = base_dir / f"frequency={frequency_label}"
|
||||
if target_frequency_dir.exists():
|
||||
shutil.rmtree(target_frequency_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_minute_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 minute chunk %d: %d rows (%d symbols done)",
|
||||
chunk_idx,
|
||||
len(chunk),
|
||||
succeeded,
|
||||
)
|
||||
buffer = []
|
||||
chunk_idx += 1
|
||||
|
||||
for i, (symbol, df) in enumerate(
|
||||
download_minute_batch(symbols, start_date, end_date, frequency=frequency), start=1
|
||||
):
|
||||
if df is None:
|
||||
logger.warning(" %s: no minute data", symbol)
|
||||
else:
|
||||
df["symbol_id"] = symbol
|
||||
df["symbol_name"] = names.get(symbol, symbol)
|
||||
buffer.append(df[MINUTE_BAR_COLUMNS])
|
||||
succeeded += 1
|
||||
if len(buffer) >= chunk_size:
|
||||
flush()
|
||||
if i % 100 == 0:
|
||||
logger.info("Minute progress: %d/%d symbols", i, n_requested)
|
||||
flush()
|
||||
|
||||
if succeeded == 0:
|
||||
raise RuntimeError("No minute data downloaded for any symbol")
|
||||
|
||||
return {
|
||||
"dataset_path": str(base_dir),
|
||||
"frequency": frequency_label,
|
||||
"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()),
|
||||
}
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Daily feature plugin package."""
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Base class for daily feature plugins."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class BaseFeature(ABC):
|
||||
"""Aggregate raw minute bars into daily, symbol-keyed feature columns."""
|
||||
|
||||
#: Unique registry key. Every concrete feature must set this to a non-empty str.
|
||||
name: str = ""
|
||||
|
||||
@abstractmethod
|
||||
def compute(
|
||||
self,
|
||||
minute: pd.DataFrame,
|
||||
daily: pd.DataFrame | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""Compute daily features.
|
||||
|
||||
Args:
|
||||
minute: Raw minute bars with ``symbol_id`` and ``date`` keys.
|
||||
daily: Optional daily data frame for calendar alignment or
|
||||
reference daily columns.
|
||||
|
||||
Returns:
|
||||
DataFrame with ``symbol_id``, ``date``, and one or more numeric
|
||||
feature columns.
|
||||
"""
|
||||
|
||||
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,108 @@
|
||||
"""CLI for daily feature computation."""
|
||||
|
||||
import os
|
||||
|
||||
import click
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.features.compute import compute_feature
|
||||
from pipeline.features.registry import available_features, load_feature_module
|
||||
|
||||
|
||||
@click.group(name="feature")
|
||||
def feature():
|
||||
"""Compute daily feature parquet files from minute bars."""
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@feature.command("list")
|
||||
@click.option(
|
||||
"--feature-module", "feature_modules", multiple=True,
|
||||
help="External module(s) to import first (dotted path or .py file)",
|
||||
)
|
||||
def list_(feature_modules):
|
||||
"""List the registered feature types."""
|
||||
for spec in feature_modules:
|
||||
load_feature_module(spec)
|
||||
for name in available_features():
|
||||
click.echo(name)
|
||||
|
||||
|
||||
@feature.command("compute")
|
||||
@click.option("--minute-path", required=True, help="Path to minute parquet dataset/file")
|
||||
@click.option("--daily-path", default=None, help="Optional daily data parquet for alignment")
|
||||
@click.option("--feature-type", required=True, help="Registry key of the feature class")
|
||||
@click.option("--feature-name", required=True, help="Name for this feature run/output file")
|
||||
@click.option("--output-dir", default="features", help="Directory to save feature parquet")
|
||||
@click.option(
|
||||
"--feature-module", "feature_modules", multiple=True,
|
||||
help="External module(s) to import so their features register (dotted path or .py file)",
|
||||
)
|
||||
@click.option(
|
||||
"--param", "extra_params", multiple=True,
|
||||
help="Extra feature constructor param as name=value (repeatable)",
|
||||
)
|
||||
def compute(
|
||||
minute_path,
|
||||
daily_path,
|
||||
feature_type,
|
||||
feature_name,
|
||||
output_dir,
|
||||
feature_modules,
|
||||
extra_params,
|
||||
):
|
||||
"""Compute one daily feature file from raw minute bars."""
|
||||
for spec in feature_modules:
|
||||
load_feature_module(spec)
|
||||
|
||||
options = available_features()
|
||||
if feature_type not in options:
|
||||
raise click.BadParameter(
|
||||
f"Unknown feature-type '{feature_type}'. Available: {options}. "
|
||||
f"Use --feature-module to register an external feature.",
|
||||
param_hint="--feature-type",
|
||||
)
|
||||
|
||||
minute = pd.read_parquet(minute_path)
|
||||
click.echo(f"Loaded minute bars: {len(minute):,} rows from {minute_path}")
|
||||
|
||||
daily = None
|
||||
if daily_path:
|
||||
daily = pd.read_parquet(daily_path)
|
||||
click.echo(f"Loaded daily data: {len(daily):,} rows from {daily_path}")
|
||||
|
||||
result = compute_feature(
|
||||
minute=minute,
|
||||
daily=daily,
|
||||
feature_type=feature_type,
|
||||
**_parse_params(extra_params),
|
||||
)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
out_path = f"{output_dir}/{feature_name}.pq"
|
||||
result.to_parquet(out_path, index=False)
|
||||
feature_cols = [col for col in result.columns if col not in ("symbol_id", "date")]
|
||||
click.echo(
|
||||
f"Saved feature: {out_path} ({len(result):,} rows, "
|
||||
f"{len(feature_cols)} columns)"
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Feature computation and validation."""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
import pandas as pd
|
||||
from pandas.api.types import is_numeric_dtype
|
||||
|
||||
from pipeline.features.registry import get_feature
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FEATURE_KEY_COLUMNS = ["symbol_id", "date"]
|
||||
|
||||
|
||||
def validate_feature_frame(features: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Validate and normalize a daily feature frame.
|
||||
|
||||
A valid feature frame is keyed by unique ``symbol_id,date`` rows and has at
|
||||
least one numeric feature column beyond those keys.
|
||||
"""
|
||||
duplicated = features.columns[features.columns.duplicated()].tolist()
|
||||
if duplicated:
|
||||
raise ValueError(f"Feature output has duplicate columns: {duplicated}")
|
||||
|
||||
missing = [col for col in FEATURE_KEY_COLUMNS if col not in features.columns]
|
||||
if missing:
|
||||
raise ValueError(f"Feature output missing required columns: {missing}")
|
||||
|
||||
out = features.copy()
|
||||
out["date"] = pd.to_datetime(out["date"])
|
||||
|
||||
if out.duplicated(FEATURE_KEY_COLUMNS).any():
|
||||
raise ValueError("Feature output has duplicate symbol_id,date rows")
|
||||
|
||||
feature_cols = [col for col in out.columns if col not in FEATURE_KEY_COLUMNS]
|
||||
if not feature_cols:
|
||||
raise ValueError("Feature output must include at least one feature column")
|
||||
|
||||
non_numeric = [col for col in feature_cols if not is_numeric_dtype(out[col])]
|
||||
if non_numeric:
|
||||
raise ValueError(f"Feature columns must be numeric: {non_numeric}")
|
||||
|
||||
out = out[FEATURE_KEY_COLUMNS + feature_cols].copy()
|
||||
return out.sort_values(FEATURE_KEY_COLUMNS).reset_index(drop=True)
|
||||
|
||||
|
||||
def compute_feature(
|
||||
minute: pd.DataFrame,
|
||||
feature_type: str,
|
||||
daily: pd.DataFrame | None = None,
|
||||
**params,
|
||||
) -> pd.DataFrame:
|
||||
"""Compute one registered feature from raw minute bars."""
|
||||
feature = get_feature(feature_type, **params)
|
||||
result = validate_feature_frame(feature.compute(minute=minute, daily=daily))
|
||||
feature_cols = [col for col in result.columns if col not in FEATURE_KEY_COLUMNS]
|
||||
logger.info(
|
||||
"Feature '%s' (%r): %d symbols × %d dates, columns=%s",
|
||||
feature_type,
|
||||
feature,
|
||||
result["symbol_id"].nunique(),
|
||||
result["date"].nunique(),
|
||||
feature_cols,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def read_feature_frames(feature_paths: Iterable[str | Path]) -> list[pd.DataFrame]:
|
||||
"""Read and validate feature parquet files."""
|
||||
return [
|
||||
validate_feature_frame(pd.read_parquet(path))
|
||||
for path in feature_paths
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Built-in feature library."""
|
||||
|
||||
from pipeline.features.library import minute_daily_summary # noqa: F401
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Daily summary features derived from raw minute bars."""
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.features.base import BaseFeature
|
||||
from pipeline.features.registry import register_feature
|
||||
|
||||
|
||||
@register_feature
|
||||
class MinuteDailySummaryFeature(BaseFeature):
|
||||
"""Aggregate intraday bars into daily summary columns."""
|
||||
|
||||
name = "minute_daily_summary"
|
||||
|
||||
def compute(
|
||||
self,
|
||||
minute: pd.DataFrame,
|
||||
daily: pd.DataFrame | None = None,
|
||||
) -> pd.DataFrame:
|
||||
minute = minute.copy()
|
||||
minute["date"] = pd.to_datetime(minute["date"])
|
||||
sort_cols = ["symbol_id", "date"]
|
||||
if "datetime" in minute.columns:
|
||||
minute["datetime"] = pd.to_datetime(minute["datetime"])
|
||||
sort_cols.append("datetime")
|
||||
elif "time" in minute.columns:
|
||||
sort_cols.append("time")
|
||||
minute = minute.sort_values(sort_cols)
|
||||
|
||||
grouped = minute.groupby(["symbol_id", "date"], sort=True)
|
||||
summary = grouped.agg(
|
||||
minute_bar_count=("close", "count"),
|
||||
first_open=("open", "first"),
|
||||
last_close=("close", "last"),
|
||||
high=("high", "max"),
|
||||
low=("low", "min"),
|
||||
volume_sum=("volume", "sum"),
|
||||
amount_sum=("amount", "sum"),
|
||||
)
|
||||
summary["minute_intraday_return"] = (
|
||||
summary["last_close"] / summary["first_open"] - 1.0
|
||||
)
|
||||
summary["minute_intraday_range"] = summary["high"] / summary["low"] - 1.0
|
||||
summary["minute_vwap"] = (
|
||||
summary["amount_sum"] / summary["volume_sum"].where(summary["volume_sum"] > 0)
|
||||
)
|
||||
summary = summary.reset_index()
|
||||
|
||||
if daily is not None:
|
||||
daily_keys = daily[["symbol_id", "date"]].copy()
|
||||
daily_keys["date"] = pd.to_datetime(daily_keys["date"])
|
||||
daily_keys = daily_keys.drop_duplicates(["symbol_id", "date"])
|
||||
result = daily_keys.merge(summary, on=["symbol_id", "date"], how="left")
|
||||
if "close" in daily.columns:
|
||||
daily_close = daily[["symbol_id", "date", "close"]].copy()
|
||||
daily_close["date"] = pd.to_datetime(daily_close["date"])
|
||||
daily_close = daily_close.drop_duplicates(["symbol_id", "date"])
|
||||
result = result.merge(
|
||||
daily_close.rename(columns={"close": "daily_close"}),
|
||||
on=["symbol_id", "date"],
|
||||
how="left",
|
||||
)
|
||||
reference_close = result["daily_close"].fillna(result["last_close"])
|
||||
else:
|
||||
reference_close = result["last_close"]
|
||||
else:
|
||||
result = summary
|
||||
reference_close = result["last_close"]
|
||||
|
||||
result["minute_vwap_deviation"] = (
|
||||
result["minute_vwap"] / reference_close.replace(0.0, np.nan) - 1.0
|
||||
)
|
||||
return result[
|
||||
[
|
||||
"symbol_id",
|
||||
"date",
|
||||
"minute_bar_count",
|
||||
"minute_intraday_return",
|
||||
"minute_intraday_range",
|
||||
"minute_vwap",
|
||||
"minute_vwap_deviation",
|
||||
]
|
||||
]
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Registry and factory for daily feature plugins."""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import Optional, Type
|
||||
|
||||
from pipeline.features.base import BaseFeature
|
||||
|
||||
_REGISTRY: dict[str, Type[BaseFeature]] = {}
|
||||
_builtins_loaded = False
|
||||
|
||||
|
||||
def register_feature(cls: Type[BaseFeature]) -> Type[BaseFeature]:
|
||||
"""Class decorator that registers a feature under ``BaseFeature.name``."""
|
||||
if not (isinstance(cls, type) and issubclass(cls, BaseFeature)):
|
||||
raise TypeError(f"{cls!r} is not a BaseFeature 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"Feature name '{key}' already registered by {existing.__name__}"
|
||||
)
|
||||
_REGISTRY[key] = cls
|
||||
return cls
|
||||
|
||||
|
||||
def available_features() -> list[str]:
|
||||
"""Sorted names of all registered features (built-ins are loaded lazily)."""
|
||||
_ensure_builtins()
|
||||
return sorted(_REGISTRY)
|
||||
|
||||
|
||||
def get_feature(name: str, **params) -> BaseFeature:
|
||||
"""Instantiate a registered feature by name.
|
||||
|
||||
Only parameters accepted by the feature class's ``__init__`` are forwarded.
|
||||
"""
|
||||
_ensure_builtins()
|
||||
if name not in _REGISTRY:
|
||||
raise KeyError(f"Unknown feature '{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_feature_module(spec: str) -> None:
|
||||
"""Import an external module so its ``@register_feature`` classes register."""
|
||||
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"Feature 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 feature 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[BaseFeature]) -> 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.features.library # noqa: F401
|
||||
_builtins_loaded = True
|
||||
@@ -347,3 +347,69 @@ def test_universe_filter_does_not_corrupt_signal_history():
|
||||
held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique())
|
||||
# The two most liquid names (highest amount) are sh600519, sz300750.
|
||||
assert held == {"sh600519", "sz300750"}
|
||||
|
||||
|
||||
# --- feature-aware alpha integration ----------------------------------------
|
||||
|
||||
def test_compute_alpha_without_feature_path_matches_empty_feature_paths():
|
||||
data = _make_data()
|
||||
|
||||
base = compute_alpha(data, "rev5", "reversal", lookback=5)
|
||||
with_empty_features = compute_alpha(
|
||||
data,
|
||||
"rev5",
|
||||
"reversal",
|
||||
lookback=5,
|
||||
feature_paths=[],
|
||||
)
|
||||
|
||||
pd.testing.assert_frame_equal(base, with_empty_features)
|
||||
|
||||
|
||||
def test_feature_aware_alpha_reads_joined_feature_column(tmp_path):
|
||||
module_path = tmp_path / "feature_aware_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 FeatureAwareAlpha(BaseAlpha):
|
||||
name = "feature_aware_test_alpha"
|
||||
|
||||
def signal_from_data(
|
||||
self,
|
||||
data: pd.DataFrame,
|
||||
close: pd.DataFrame,
|
||||
) -> pd.DataFrame:
|
||||
signal = data.pivot_table(
|
||||
index="date",
|
||||
columns="symbol_id",
|
||||
values="toy_feature",
|
||||
aggfunc="first",
|
||||
)
|
||||
return signal.reindex(index=close.index, columns=close.columns)
|
||||
'''))
|
||||
|
||||
data = _make_data()
|
||||
feature = data[["symbol_id", "date"]].copy()
|
||||
feature["toy_feature"] = feature["symbol_id"].map({
|
||||
"sh600000": 1.0,
|
||||
"sz000001": 2.0,
|
||||
"sh600519": 3.0,
|
||||
})
|
||||
feature_path = tmp_path / "toy_feature.pq"
|
||||
feature.to_parquet(feature_path, index=False)
|
||||
|
||||
load_alpha_module(str(module_path))
|
||||
result = compute_alpha(
|
||||
data,
|
||||
"feature_run",
|
||||
"feature_aware_test_alpha",
|
||||
feature_paths=[str(feature_path)],
|
||||
)
|
||||
|
||||
assert list(result.columns) == ALPHA_COLUMNS
|
||||
assert (result["alpha_name"] == "feature_run").all()
|
||||
last = result[result["date"] == result["date"].max()]
|
||||
assert last.set_index("symbol_id")["weight"].idxmax() == "sh600519"
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Tests for minute-derived daily feature plugins."""
|
||||
|
||||
import textwrap
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from pipeline.features.compute import compute_feature, validate_feature_frame
|
||||
from pipeline.features.registry import (
|
||||
available_features,
|
||||
get_feature,
|
||||
load_feature_module,
|
||||
)
|
||||
|
||||
|
||||
def _minute_bars() -> pd.DataFrame:
|
||||
return pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000", "sz000001"],
|
||||
"symbol_name": ["PF Bank", "PF Bank", "Ping An"],
|
||||
"datetime": pd.to_datetime([
|
||||
"2024-01-02 09:35:00",
|
||||
"2024-01-02 09:40:00",
|
||||
"2024-01-02 09:35:00",
|
||||
]),
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-02"]),
|
||||
"time": ["09:35:00", "09:40:00", "09:35:00"],
|
||||
"frequency": ["5m", "5m", "5m"],
|
||||
"open": [10.0, 10.5, 20.0],
|
||||
"high": [11.0, 12.0, 21.0],
|
||||
"low": [9.0, 10.0, 19.0],
|
||||
"close": [10.5, 11.0, 20.5],
|
||||
"volume": [100.0, 300.0, 200.0],
|
||||
"amount": [1000.0, 3300.0, 4100.0],
|
||||
"vwap": [10.0, 11.0, 20.5],
|
||||
"adjustflag": ["3", "3", "3"],
|
||||
})
|
||||
|
||||
|
||||
def test_built_in_minute_daily_summary():
|
||||
daily = pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sz000001", "sh600000"],
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-03"]),
|
||||
"close": [11.0, 20.5, 12.0],
|
||||
})
|
||||
|
||||
result = compute_feature(
|
||||
minute=_minute_bars(),
|
||||
daily=daily,
|
||||
feature_type="minute_daily_summary",
|
||||
)
|
||||
|
||||
assert "minute_daily_summary" in available_features()
|
||||
row = result[
|
||||
(result["symbol_id"] == "sh600000")
|
||||
& (result["date"] == pd.Timestamp("2024-01-02"))
|
||||
].iloc[0]
|
||||
assert row["minute_bar_count"] == 2
|
||||
assert np.isclose(row["minute_intraday_return"], 11.0 / 10.0 - 1.0)
|
||||
assert np.isclose(row["minute_intraday_range"], 12.0 / 9.0 - 1.0)
|
||||
assert np.isclose(row["minute_vwap"], 4300.0 / 400.0)
|
||||
assert np.isclose(row["minute_vwap_deviation"], (4300.0 / 400.0) / 11.0 - 1.0)
|
||||
|
||||
missing = result[
|
||||
(result["symbol_id"] == "sh600000")
|
||||
& (result["date"] == pd.Timestamp("2024-01-03"))
|
||||
].iloc[0]
|
||||
assert pd.isna(missing["minute_vwap"])
|
||||
|
||||
|
||||
def test_load_external_feature_module_and_filter_params(tmp_path):
|
||||
module_path = tmp_path / "external_feature.py"
|
||||
module_path.write_text(textwrap.dedent('''
|
||||
import pandas as pd
|
||||
from pipeline.features.base import BaseFeature
|
||||
from pipeline.features.registry import register_feature
|
||||
|
||||
@register_feature
|
||||
class ExternalVolumeFeature(BaseFeature):
|
||||
name = "external_volume_feature"
|
||||
|
||||
def __init__(self, scale: float = 1.0):
|
||||
self.scale = scale
|
||||
|
||||
def compute(self, minute: pd.DataFrame, daily=None) -> pd.DataFrame:
|
||||
out = (
|
||||
minute.groupby(["symbol_id", "date"], as_index=False)["volume"]
|
||||
.sum()
|
||||
.rename(columns={"volume": "scaled_volume"})
|
||||
)
|
||||
out["scaled_volume"] *= self.scale
|
||||
return out
|
||||
'''))
|
||||
|
||||
load_feature_module(str(module_path))
|
||||
assert "external_volume_feature" in available_features()
|
||||
|
||||
instance = get_feature("external_volume_feature", scale=2.0, ignored=99)
|
||||
assert instance.scale == 2.0
|
||||
assert not hasattr(instance, "ignored")
|
||||
|
||||
result = compute_feature(
|
||||
minute=_minute_bars(),
|
||||
feature_type="external_volume_feature",
|
||||
scale=2.0,
|
||||
ignored=99,
|
||||
)
|
||||
row = result[result["symbol_id"] == "sh600000"].iloc[0]
|
||||
assert np.isclose(row["scaled_volume"], 800.0)
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_missing_keys():
|
||||
with pytest.raises(ValueError, match="missing required"):
|
||||
validate_feature_frame(pd.DataFrame({"symbol_id": ["sh600000"], "x": [1.0]}))
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_duplicate_keys_after_date_normalization():
|
||||
with pytest.raises(ValueError, match="duplicate symbol_id,date"):
|
||||
validate_feature_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000"],
|
||||
"date": ["2024-01-02", pd.Timestamp("2024-01-02")],
|
||||
"x": [1.0, 2.0],
|
||||
}))
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_duplicate_columns():
|
||||
bad = pd.DataFrame(
|
||||
[["sh600000", pd.Timestamp("2024-01-02"), 1.0, 2.0]],
|
||||
columns=["symbol_id", "date", "dup", "dup"],
|
||||
)
|
||||
with pytest.raises(ValueError, match="duplicate columns"):
|
||||
validate_feature_frame(bad)
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_non_numeric_feature_columns():
|
||||
with pytest.raises(ValueError, match="numeric"):
|
||||
validate_feature_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"date": [pd.Timestamp("2024-01-02")],
|
||||
"bad": ["not numeric"],
|
||||
}))
|
||||
@@ -0,0 +1,215 @@
|
||||
"""Tests for raw Baostock minute bar download plumbing."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import data.downloader as low_level_downloader
|
||||
import pipeline.data.downloader as pipeline_downloader
|
||||
from data.downloader import download_minute_batch
|
||||
from pipeline.common.schema import MINUTE_BAR_COLUMNS
|
||||
from pipeline.data.downloader import download_minute_universe
|
||||
|
||||
|
||||
class _FakeResult:
|
||||
def __init__(self, rows, error_code="0", error_msg=""):
|
||||
self.rows = rows
|
||||
self.error_code = error_code
|
||||
self.error_msg = error_msg
|
||||
self._idx = -1
|
||||
|
||||
def next(self):
|
||||
self._idx += 1
|
||||
return self._idx < len(self.rows)
|
||||
|
||||
def get_row_data(self):
|
||||
return self.rows[self._idx]
|
||||
|
||||
|
||||
def test_download_minute_batch_maps_and_parses_baostock_rows(monkeypatch):
|
||||
rows = [
|
||||
[
|
||||
"2024-01-02",
|
||||
"20240102093500000",
|
||||
"sh.600000",
|
||||
"10",
|
||||
"11",
|
||||
"9",
|
||||
"10.5",
|
||||
"1000",
|
||||
"10500",
|
||||
"3",
|
||||
],
|
||||
[
|
||||
"2024-01-02",
|
||||
"20240102094000000",
|
||||
"sh.600000",
|
||||
"10.5",
|
||||
"12",
|
||||
"10",
|
||||
"11",
|
||||
"2000",
|
||||
"22000",
|
||||
"3",
|
||||
],
|
||||
]
|
||||
calls = []
|
||||
|
||||
def fake_query(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return _FakeResult(rows)
|
||||
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
fake_query,
|
||||
)
|
||||
|
||||
[(symbol, df)] = list(
|
||||
download_minute_batch(
|
||||
["sh600000"],
|
||||
"2024-01-02",
|
||||
"2024-01-02",
|
||||
frequency=5,
|
||||
)
|
||||
)
|
||||
|
||||
assert symbol == "sh600000"
|
||||
assert calls[0]["code"] == "sh.600000"
|
||||
assert calls[0]["frequency"] == "5"
|
||||
assert calls[0]["adjustflag"] == "3"
|
||||
assert df is not None
|
||||
assert df["datetime"].iloc[0] == pd.Timestamp("2024-01-02 09:35:00")
|
||||
assert df["time"].tolist() == ["09:35:00", "09:40:00"]
|
||||
assert (df["frequency"] == "5m").all()
|
||||
assert np.isclose(df["open"].iloc[0], 10.0)
|
||||
assert np.isclose(df["vwap"].iloc[0], 10.5)
|
||||
assert pd.api.types.is_numeric_dtype(df["volume"])
|
||||
|
||||
|
||||
def test_download_minute_batch_empty_result_yields_none(monkeypatch):
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
lambda **kwargs: _FakeResult([]),
|
||||
)
|
||||
|
||||
assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
|
||||
("sh600000", None)
|
||||
]
|
||||
|
||||
|
||||
def test_download_minute_batch_rejects_unparsed_timestamps(monkeypatch):
|
||||
bad_rows = [[
|
||||
"2024-01-02",
|
||||
"not-a-time",
|
||||
"sh.600000",
|
||||
"10",
|
||||
"11",
|
||||
"9",
|
||||
"10.5",
|
||||
"1000",
|
||||
"10500",
|
||||
"3",
|
||||
]]
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
lambda **kwargs: _FakeResult(bad_rows),
|
||||
)
|
||||
|
||||
assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
|
||||
("sh600000", None)
|
||||
]
|
||||
|
||||
|
||||
def test_download_minute_universe_writes_frequency_month_partitions(tmp_path, monkeypatch):
|
||||
minute = pd.DataFrame({
|
||||
"symbol": ["sh600000", "sh600000"],
|
||||
"datetime": pd.to_datetime(["2024-01-02 09:35:00", "2024-01-02 09:40:00"]),
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02"]),
|
||||
"time": ["09:35:00", "09:40:00"],
|
||||
"frequency": ["5m", "5m"],
|
||||
"open": [10.0, 10.5],
|
||||
"high": [11.0, 12.0],
|
||||
"low": [9.0, 10.0],
|
||||
"close": [10.5, 11.0],
|
||||
"volume": [1000.0, 2000.0],
|
||||
"amount": [10500.0, 22000.0],
|
||||
"vwap": [10.5, 11.0],
|
||||
"adjustflag": ["3", "3"],
|
||||
})
|
||||
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"_resolve_universe",
|
||||
lambda universe, max_symbols=0: pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"symbol_name": ["PF Bank"],
|
||||
}),
|
||||
)
|
||||
|
||||
def fake_batch(symbols, start, end, frequency=5):
|
||||
assert symbols == ["sh600000"]
|
||||
assert frequency == "5"
|
||||
yield "sh600000", minute
|
||||
|
||||
monkeypatch.setattr(pipeline_downloader, "download_minute_batch", fake_batch)
|
||||
|
||||
preserved = tmp_path / "toy" / "frequency=15m" / "month=2024-01" / "old.pq"
|
||||
preserved.parent.mkdir(parents=True)
|
||||
preserved_minute = minute.copy()
|
||||
preserved_minute["frequency"] = "15m"
|
||||
preserved_minute["symbol_id"] = "sh600000"
|
||||
preserved_minute["symbol_name"] = "PF Bank"
|
||||
preserved_minute[MINUTE_BAR_COLUMNS].to_parquet(preserved, index=False)
|
||||
|
||||
stats = download_minute_universe(
|
||||
universe="toy",
|
||||
start_date="2024-01-02",
|
||||
end_date="2024-01-02",
|
||||
output_dir=str(tmp_path),
|
||||
chunk_size=1,
|
||||
frequency="5",
|
||||
)
|
||||
|
||||
dataset_path = Path(stats["dataset_path"])
|
||||
assert (dataset_path / "frequency=5m" / "month=2024-01").is_dir()
|
||||
assert preserved.exists()
|
||||
out = pd.read_parquet(dataset_path / "frequency=5m")
|
||||
assert (set(MINUTE_BAR_COLUMNS) - {"frequency"}) <= set(out.columns)
|
||||
assert set(out["symbol_id"]) == {"sh600000"}
|
||||
assert set(out["symbol_name"]) == {"PF Bank"}
|
||||
assert stats["n_rows"] == 2
|
||||
|
||||
|
||||
def test_download_minute_universe_raises_when_all_symbols_empty(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"_resolve_universe",
|
||||
lambda universe, max_symbols=0: pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"symbol_name": ["PF Bank"],
|
||||
}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"download_minute_batch",
|
||||
lambda symbols, start, end, frequency=5: iter([("sh600000", None)]),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="No minute data"):
|
||||
download_minute_universe(
|
||||
universe="toy",
|
||||
start_date="2024-01-02",
|
||||
end_date="2024-01-02",
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
Reference in New Issue
Block a user