royy/chinese-equity-quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
31baa18ce5b33bf16e006457cf722ec00eba6860
chinese-equity-quant/tests/test_workflow_sanity.py
T
Yuxuan Yan 31baa18ce5 Add offline workflow and coverage tests
2026-06-16 17:37:16 +08:00

509 lines
17 KiB
Python
Raw Blame History

"""Verbose offline checks for the daily research workflow."""
from __future__ import annotations
from pathlib import Path
import numpy as np
import pandas as pd
from pipeline.alpha.compute import compute_alpha
from pipeline.combo.combine import combine_alphas
from pipeline.common.schema import (
ALPHA_COLUMNS,
COMBO_COLUMNS,
FILL_COLUMNS,
PNL_COLUMNS,
POSITION_COLUMNS,
)
from pipeline.portfolio.constraints import (
PriceLimitConstraint,
SuspensionConstraint,
VolumeCapConstraint,
)
from pipeline.portfolio.construct import construct_positions
from pipeline.portfolio.research import evaluate_portfolio
from pipeline.portfolio.simulator import ReferenceSimulator
from tests.helpers import (
GENERATED_SYMBOLS,
generated_sessions,
make_generated_alpha_weights,
make_generated_combo_weights,
make_generated_daily_bars,
)
FIXTURE_PATH = Path(__file__).parent / "fixtures" / "daily_bars_real_2024_01_sample.pq"
def _assert_sorted_by_symbol_date(frame: pd.DataFrame) -> None:
expected = frame.sort_values(["symbol_id", "date"]).reset_index(drop=True)
pd.testing.assert_frame_equal(frame.reset_index(drop=True), expected)
def _assert_metric_dict_is_finite(metrics: dict[str, float]) -> None:
for key in (
"cumulative_return",
"sharpe_annual",
"turnover_annual",
"max_drawdown",
"hit_rate",
"n_dates",
):
assert key in metrics
assert np.isfinite(metrics[key])
assert "ic" not in metrics
assert "rank_ic" not in metrics
assert "ir" not in metrics
def test_tiny_workflow_golden_outputs_are_stable(tmp_path):
dates = pd.to_datetime(["2024-01-02", "2024-01-03", "2024-01-04"])
daily_bars = pd.DataFrame([
{
"symbol_id": "sh600000",
"symbol_name": "A",
"date": dates[0],
"open": 10.0,
"high": 10.0,
"low": 10.0,
"close": 10.0,
"preclose": 10.0,
"volume": 1_000_000.0,
"amount": 10_000_000.0,
"vwap": 10.0,
"turn": 1.0,
"pctChg": 0.0,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
{
"symbol_id": "sz000001",
"symbol_name": "B",
"date": dates[0],
"open": 20.0,
"high": 20.0,
"low": 20.0,
"close": 20.0,
"preclose": 20.0,
"volume": 1_000_000.0,
"amount": 20_000_000.0,
"vwap": 20.0,
"turn": 1.0,
"pctChg": 0.0,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
{
"symbol_id": "sh600000",
"symbol_name": "A",
"date": dates[1],
"open": 10.0,
"high": 12.0,
"low": 10.0,
"close": 12.0,
"preclose": 10.0,
"volume": 1_000_000.0,
"amount": 10_000_000.0,
"vwap": 10.0,
"turn": 1.0,
"pctChg": 20.0,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
{
"symbol_id": "sz000001",
"symbol_name": "B",
"date": dates[1],
"open": 20.0,
"high": 20.0,
"low": 18.0,
"close": 18.0,
"preclose": 20.0,
"volume": 1_000_000.0,
"amount": 20_000_000.0,
"vwap": 20.0,
"turn": 1.0,
"pctChg": -10.0,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
{
"symbol_id": "sh600000",
"symbol_name": "A",
"date": dates[2],
"open": 12.0,
"high": 13.0,
"low": 12.0,
"close": 13.0,
"preclose": 12.0,
"volume": 1_000_000.0,
"amount": 12_000_000.0,
"vwap": 12.0,
"turn": 1.0,
"pctChg": 8.33,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
{
"symbol_id": "sz000001",
"symbol_name": "B",
"date": dates[2],
"open": 18.0,
"high": 21.0,
"low": 18.0,
"close": 21.0,
"preclose": 18.0,
"volume": 1_000_000.0,
"amount": 18_000_000.0,
"vwap": 18.0,
"turn": 1.0,
"pctChg": 16.67,
"tradestatus": 1,
"isST": 0,
"peTTM": 1.0,
"pbMRQ": 1.0,
"psTTM": 1.0,
"pcfNcfTTM": 1.0,
},
])
alpha = pd.DataFrame({
"symbol_id": ["sh600000", "sz000001", "sh600000", "sz000001"],
"date": [dates[0], dates[0], dates[1], dates[1]],
"alpha_name": ["gold_alpha"] * 4,
"weight": [1.0, -1.0, -1.0, 1.0],
})
alpha_path = tmp_path / "gold_alpha.pq"
alpha.to_parquet(alpha_path, index=False)
combo = combine_alphas([str(alpha_path)], "gold_combo")
positions = construct_positions(combo, daily_bars, booksize=20_000.0, portfolio_name="gold_port")
fills, pnl = ReferenceSimulator().run(positions, daily_bars)
expected_combo = pd.DataFrame({
"symbol_id": ["sh600000", "sh600000", "sz000001", "sz000001"],
"date": [dates[0], dates[1], dates[0], dates[1]],
"combo_name": ["gold_combo"] * 4,
"weight": [1.0, -1.0, -1.0, 1.0],
})
expected_positions = pd.DataFrame({
"symbol_id": ["sh600000", "sh600000", "sz000001", "sz000001"],
"date": [dates[0], dates[1], dates[0], dates[1]],
"portfolio_name": ["gold_port"] * 4,
"target_weight": [0.5, -0.5, -0.5, 0.5],
"target_value": [10000.0, -10000.0, -10000.0, 10000.0],
"target_shares": [1000.0, -10000.0 / 12.0, -500.0, 10000.0 / 18.0],
"position_shares": [1000, -833, -500, 556],
"position_value": [10000.0, -9996.0, -10000.0, 10008.0],
"price": [10.0, 12.0, 20.0, 18.0],
})
expected_fills = pd.DataFrame({
"symbol_id": ["sh600000", "sz000001", "sh600000", "sz000001"],
"date": [dates[1], dates[1], dates[2], dates[2]],
"portfolio_name": ["gold_port"] * 4,
"prev_shares": [0, 0, 1000, -500],
"target_shares": [1000, -500, -833, 556],
"traded_shares": [1000, -500, -1833, 1056],
"realized_shares": [1000, -500, -833, 556],
"blocked": [0, 0, 0, 0],
"trade_cost": [0.0, 0.0, 0.0, 0.0],
})
expected_pnl = pd.DataFrame({
"date": [dates[1], dates[2]],
"portfolio_name": ["gold_port", "gold_port"],
"gross_exposure": [21000.0, 22505.0],
"net_exposure": [3000.0, 847.0],
"pnl": [3000.0, 835.0],
"cost": [0.0, 0.0],
"turnover": [1.0, 2.0502],
"n_positions": [2, 2],
})
pd.testing.assert_frame_equal(combo, expected_combo)
pd.testing.assert_frame_equal(positions, expected_positions)
pd.testing.assert_frame_equal(fills, expected_fills)
pd.testing.assert_frame_equal(pnl, expected_pnl)
def test_generated_alpha_combo_portfolio_execution_workflow(tmp_path):
daily_bars = make_generated_daily_bars()
computed_alpha = compute_alpha(
data=daily_bars,
alpha_name="generated_reversal_3d",
alpha_type="reversal",
lookback=3,
)
assert list(computed_alpha.columns) == ALPHA_COLUMNS
assert not computed_alpha.empty
assert set(computed_alpha["symbol_id"]).issubset(set(GENERATED_SYMBOLS))
assert computed_alpha["date"].min() > daily_bars["date"].min()
assert computed_alpha["weight"].notna().all()
assert computed_alpha["weight"].abs().sum() > 0.0
assert {"ic", "rank_ic", "ir"}.isdisjoint(computed_alpha.columns)
_assert_sorted_by_symbol_date(computed_alpha)
alpha_a = make_generated_alpha_weights("alpha_a", zero_date_index=2)
alpha_b = make_generated_alpha_weights(
"alpha_b",
scale=0.5,
offset=0.25,
zero_date_index=2,
)
alpha_a_path = tmp_path / "alpha_a.pq"
alpha_b_path = tmp_path / "alpha_b.pq"
alpha_a.to_parquet(alpha_a_path, index=False)
alpha_b.to_parquet(alpha_b_path, index=False)
identity_combo = combine_alphas([str(alpha_a_path)], "identity_combo")
assert list(identity_combo.columns) == COMBO_COLUMNS
assert (identity_combo["combo_name"] == "identity_combo").all()
pd.testing.assert_frame_equal(
identity_combo[["symbol_id", "date", "weight"]],
alpha_a[["symbol_id", "date", "weight"]],
)
equal_combo = combine_alphas([str(alpha_a_path), str(alpha_b_path)], "equal_combo")
expected_equal_weights = (
pd.concat([alpha_a, alpha_b], ignore_index=True)
.groupby(["symbol_id", "date"], as_index=False)["weight"]
.mean()
.sort_values(["symbol_id", "date"])
.reset_index(drop=True)
)
pd.testing.assert_frame_equal(
equal_combo[["symbol_id", "date", "weight"]],
expected_equal_weights,
)
portfolio_weights = make_generated_combo_weights("workflow_combo", zero_date_index=2)
positions = construct_positions(
weights_df=portfolio_weights,
data_df=daily_bars,
booksize=2_000_000.0,
portfolio_name="workflow_portfolio",
)
assert list(positions.columns) == POSITION_COLUMNS
assert not positions.empty
assert (positions["portfolio_name"] == "workflow_portfolio").all()
assert pd.api.types.is_integer_dtype(positions["position_shares"])
assert np.allclose(
positions["position_value"],
positions["position_shares"].astype(float) * positions["price"].fillna(0.0),
)
target_gross_by_date = positions.groupby("date")["target_weight"].apply(lambda s: s.abs().sum())
nonzero_target_dates = target_gross_by_date[target_gross_by_date > 0.0]
assert np.allclose(nonzero_target_dates, 1.0)
nonzero_share_counts = positions.loc[positions["position_shares"] != 0, "position_shares"].abs()
assert (nonzero_share_counts >= 100).all()
zero_gross_date = generated_sessions(10)[2]
previous_date = generated_sessions(10)[1]
zero_gross_positions = positions[positions["date"] == zero_gross_date].set_index("symbol_id")
previous_positions = positions[positions["date"] == previous_date].set_index("symbol_id")
common_symbols = zero_gross_positions.index.intersection(previous_positions.index)
assert not common_symbols.empty
assert (zero_gross_positions.loc[common_symbols, "target_weight"] == 0.0).all()
pd.testing.assert_series_equal(
zero_gross_positions.loc[common_symbols, "position_shares"],
previous_positions.loc[common_symbols, "position_shares"],
check_names=False,
)
simulator = ReferenceSimulator(
constraints=[
SuspensionConstraint(),
PriceLimitConstraint(),
VolumeCapConstraint(max_frac=0.02),
],
cost_bps=5,
slippage_bps=5,
)
fills, pnl = simulator.run(positions, daily_bars)
assert list(fills.columns) == FILL_COLUMNS
assert list(pnl.columns) == PNL_COLUMNS
assert not fills.empty
assert not pnl.empty
assert (fills["realized_shares"] == fills["prev_shares"] + fills["traded_shares"]).all()
assert fills["blocked"].sum() > 0
fill_prices = fills.merge(
daily_bars[["symbol_id", "date", "open"]],
on=["symbol_id", "date"],
how="left",
validate="many_to_one",
)
expected_trade_cost = (
fill_prices["traded_shares"].abs()
* fill_prices["open"].fillna(0.0)
* 10
/ 10_000
)
assert np.allclose(fill_prices["trade_cost"], expected_trade_cost)
cost_by_date = fills.groupby("date")["trade_cost"].sum()
assert np.allclose(
pnl.set_index("date")["cost"],
cost_by_date.reindex(pnl["date"], fill_value=0.0),
)
booksize_used_by_simulator = positions.groupby("date")["target_value"].apply(lambda s: s.abs().sum()).max()
traded_value_by_date = (
fill_prices.assign(traded_value=fill_prices["traded_shares"].abs() * fill_prices["open"])
.groupby("date")["traded_value"]
.sum()
)
assert np.allclose(
pnl.set_index("date")["turnover"],
traded_value_by_date.reindex(pnl["date"], fill_value=0.0) / booksize_used_by_simulator,
)
metrics = evaluate_portfolio(positions, daily_bars)
_assert_metric_dict_is_finite(metrics)
def test_generated_workflow_outputs_keep_parquet_schema_contracts(tmp_path):
daily_bars = make_generated_daily_bars(n_sessions=10, include_missing=False)
alpha = compute_alpha(
data=daily_bars,
alpha_name="schema_reversal",
alpha_type="reversal",
lookback=3,
)
alpha_path = tmp_path / "schema_reversal.pq"
alpha.to_parquet(alpha_path, index=False)
combo = combine_alphas([str(alpha_path)], "schema_combo")
positions = construct_positions(
weights_df=combo,
data_df=daily_bars,
booksize=1_000_000.0,
portfolio_name="schema_portfolio",
)
fills, pnl = ReferenceSimulator(cost_bps=5, slippage_bps=5).run(positions, daily_bars)
assert list(alpha.columns) == ALPHA_COLUMNS
assert pd.api.types.is_object_dtype(alpha["symbol_id"])
assert pd.api.types.is_datetime64_any_dtype(alpha["date"])
assert pd.api.types.is_object_dtype(alpha["alpha_name"])
assert pd.api.types.is_float_dtype(alpha["weight"])
assert not alpha.isna().any().any()
assert np.isfinite(alpha["weight"]).all()
assert list(combo.columns) == COMBO_COLUMNS
assert pd.api.types.is_object_dtype(combo["symbol_id"])
assert pd.api.types.is_datetime64_any_dtype(combo["date"])
assert pd.api.types.is_object_dtype(combo["combo_name"])
assert pd.api.types.is_float_dtype(combo["weight"])
assert not combo.isna().any().any()
assert np.isfinite(combo["weight"]).all()
assert list(positions.columns) == POSITION_COLUMNS
assert pd.api.types.is_integer_dtype(positions["position_shares"])
assert pd.api.types.is_datetime64_any_dtype(positions["date"])
assert not positions.isna().any().any()
position_numeric_columns = [
"target_weight",
"target_value",
"target_shares",
"position_value",
"price",
]
assert np.isfinite(positions[position_numeric_columns]).all().all()
assert list(fills.columns) == FILL_COLUMNS
assert pd.api.types.is_integer_dtype(fills["prev_shares"])
assert pd.api.types.is_integer_dtype(fills["target_shares"])
assert pd.api.types.is_integer_dtype(fills["traded_shares"])
assert pd.api.types.is_integer_dtype(fills["realized_shares"])
assert pd.api.types.is_integer_dtype(fills["blocked"])
assert not fills.isna().any().any()
assert np.isfinite(fills["trade_cost"]).all()
assert list(pnl.columns) == PNL_COLUMNS
assert pd.api.types.is_integer_dtype(pnl["n_positions"])
assert not pnl.isna().any().any()
pnl_numeric_columns = [
"gross_exposure",
"net_exposure",
"pnl",
"cost",
"turnover",
]
assert np.isfinite(pnl[pnl_numeric_columns]).all().all()
def test_frozen_real_fixture_runs_high_level_workflow(tmp_path):
real_daily_bars = pd.read_parquet(FIXTURE_PATH)
assert real_daily_bars.shape == (36, 19)
assert set(real_daily_bars["symbol_id"]) == set(GENERATED_SYMBOLS)
assert real_daily_bars["date"].min() == pd.Timestamp("2024-01-02")
assert real_daily_bars["date"].max() == pd.Timestamp("2024-01-12")
assert real_daily_bars.groupby("date")["symbol_id"].nunique().eq(4).all()
reversal_alpha = compute_alpha(
data=real_daily_bars,
alpha_name="real_reversal_3d",
alpha_type="reversal",
lookback=3,
)
reversal_vol_alpha = compute_alpha(
data=real_daily_bars,
alpha_name="real_reversal_vol_3d",
alpha_type="reversal_vol",
lookback=3,
vol_window=3,
)
reversal_path = tmp_path / "real_reversal.pq"
reversal_vol_path = tmp_path / "real_reversal_vol.pq"
reversal_alpha.to_parquet(reversal_path, index=False)
reversal_vol_alpha.to_parquet(reversal_vol_path, index=False)
combo = combine_alphas([str(reversal_path), str(reversal_vol_path)], "real_equal_combo")
positions = construct_positions(
weights_df=combo,
data_df=real_daily_bars,
booksize=1_000_000.0,
portfolio_name="real_fixture_portfolio",
)
fills, pnl = ReferenceSimulator(cost_bps=5, slippage_bps=5).run(positions, real_daily_bars)
metrics = evaluate_portfolio(positions, real_daily_bars)
assert not reversal_alpha.empty
assert not reversal_vol_alpha.empty
assert not combo.empty
assert not positions.empty
assert not fills.empty
assert not pnl.empty
assert np.isfinite(combo["weight"]).all()
assert np.isfinite(positions["target_weight"]).all()
assert np.isfinite(pnl[["gross_exposure", "net_exposure", "pnl", "cost", "turnover"]]).all().all()
_assert_metric_dict_is_finite(metrics)
Reference in New Issue View Git Blame Copy Permalink