chinese-equity-quant/tests/helpers.py

"""Shared deterministic test data for offline workflow tests."""

from __future__ import annotations

import numpy as np
import pandas as pd

from pipeline.common.schema import (
    ALPHA_COLUMNS,
    COMBO_COLUMNS,
    DATA_COLUMNS,
    MINUTE_BAR_COLUMNS,
)


GENERATED_SYMBOLS: tuple[str, ...] = (
    "sh600000",
    "sz000001",
    "sh600519",
    "sz300750",
)

GENERATED_SYMBOL_NAMES: dict[str, str] = {
    "sh600000": "PF Bank",
    "sz000001": "Ping An Bank",
    "sh600519": "Kweichow Moutai",
    "sz300750": "CATL",
}


def generated_sessions(n_sessions: int = 12) -> pd.DatetimeIndex:
    """Return a fixed business-day calendar used by generated fixtures."""
    return pd.bdate_range("2024-01-02", periods=n_sessions)


def make_generated_daily_bars(
    n_sessions: int = 12,
    include_missing: bool = True,
) -> pd.DataFrame:
    """Build daily bars with explicit edge cases and no randomness.

    The panel covers four A-share symbols and includes a suspended row, an ST
    flag, a zero-volume row, a missing symbol-date, and limit-style open/close
    moves. Values are deterministic so tests can assert exact identities.
    """
    dates = generated_sessions(n_sessions)
    base_close = {
        "sh600000": 10.00,
        "sz000001": 15.00,
        "sh600519": 1200.00,
        "sz300750": 180.00,
    }
    returns = {
        "sh600000": [0.000, 0.012, -0.006, 0.018, 0.100, -0.014, 0.006, 0.000, 0.008, -0.011, 0.004, 0.009],
        "sz000001": [0.000, -0.008, 0.011, -0.004, 0.006, 0.000, -0.012, 0.009, 0.005, -0.007, 0.010, -0.003],
        "sh600519": [0.000, 0.006, 0.004, -0.010, 0.012, -0.006, 0.005, 0.003, -0.009, 0.007, -0.004, 0.006],
        "sz300750": [0.000, -0.010, 0.014, 0.006, -0.008, 0.011, -0.004, 0.009, -0.200, 0.012, -0.006, 0.008],
    }
    base_volume = {
        "sh600000": 1_200_000.0,
        "sz000001": 900_000.0,
        "sh600519": 80_000.0,
        "sz300750": 240_000.0,
    }

    rows: list[dict[str, object]] = []
    for sym in GENERATED_SYMBOLS:
        closes = [base_close[sym]]
        pattern = returns[sym]
        for step in range(1, n_sessions):
            ret = pattern[step % len(pattern)]
            closes.append(closes[-1] * (1.0 + ret))
        closes_arr = np.asarray(closes, dtype=float)
        precloses = np.concatenate([[closes_arr[0]], closes_arr[:-1]])

        for i, date in enumerate(dates):
            preclose = float(precloses[i])
            close = float(closes_arr[i])
            open_price = preclose * (1.0 + 0.25 * (close / preclose - 1.0))
            high = max(open_price, close) * 1.01
            low = min(open_price, close) * 0.99
            volume = base_volume[sym] + 10_000.0 * i
            tradestatus = 1
            is_st = 0

            if sym == "sh600000" and i == 4:
                open_price = preclose * 1.10
                close = open_price
                high = open_price
                low = open_price
            if sym == "sh600000" and i == 7:
                volume = 0.0
            if sym == "sz000001" and i == 5:
                open_price = preclose
                close = preclose
                high = preclose
                low = preclose
                volume = 0.0
                tradestatus = 0
            if sym == "sh600519" and i == 7:
                is_st = 1
            if sym == "sz300750" and i == 8:
                open_price = preclose * 0.80
                close = open_price
                high = open_price
                low = open_price

            amount = volume * ((open_price + close) / 2.0)
            vwap = amount / volume if volume > 0 else np.nan
            pct_chg = (close / preclose - 1.0) * 100.0 if preclose else 0.0
            rows.append({
                "symbol_id": sym,
                "symbol_name": GENERATED_SYMBOL_NAMES[sym],
                "date": date,
                "open": open_price,
                "high": high,
                "low": low,
                "close": close,
                "preclose": preclose,
                "volume": volume,
                "amount": amount,
                "vwap": vwap,
                "turn": volume / 1_000_000.0,
                "pctChg": pct_chg,
                "tradestatus": tradestatus,
                "isST": is_st,
                "peTTM": 8.0 + i,
                "pbMRQ": 1.0 + 0.05 * i,
                "psTTM": 2.0 + 0.03 * i,
                "pcfNcfTTM": 5.0 + 0.1 * i,
            })

    result = pd.DataFrame(rows)
    if include_missing and n_sessions > 6:
        missing_mask = (
            (result["symbol_id"] == "sz300750")
            & (result["date"] == dates[6])
        )
        result = result.loc[~missing_mask].copy()

    result = result[DATA_COLUMNS]
    return result.sort_values(["date", "symbol_id"]).reset_index(drop=True)


def make_generated_minute_bars(
    daily: pd.DataFrame | None = None,
) -> pd.DataFrame:
    """Expand generated daily bars into a tiny deterministic intraday panel."""
    daily = make_generated_daily_bars() if daily is None else daily.copy()
    rows: list[dict[str, object]] = []
    bar_times = ("09:35:00", "10:30:00", "14:55:00")

    for daily_row in daily.sort_values(["date", "symbol_id"]).itertuples(index=False):
        if int(getattr(daily_row, "tradestatus", 1)) == 0:
            continue

        volume = float(daily_row.volume)
        volume_slices = [0.25 * volume, 0.35 * volume, 0.40 * volume]
        prices = np.linspace(float(daily_row.open), float(daily_row.close), len(bar_times))
        for j, time_text in enumerate(bar_times):
            dt = pd.Timestamp(daily_row.date) + pd.Timedelta(time_text)
            open_price = prices[j - 1] if j else float(daily_row.open)
            close_price = float(prices[j])
            high = max(open_price, close_price) * 1.002
            low = min(open_price, close_price) * 0.998
            minute_volume = float(volume_slices[j])
            amount = minute_volume * ((open_price + close_price) / 2.0)
            rows.append({
                "symbol_id": daily_row.symbol_id,
                "symbol_name": daily_row.symbol_name,
                "datetime": dt,
                "date": pd.Timestamp(daily_row.date).normalize(),
                "time": time_text,
                "frequency": "5m",
                "open": open_price,
                "high": high,
                "low": low,
                "close": close_price,
                "volume": minute_volume,
                "amount": amount,
                "vwap": amount / minute_volume if minute_volume > 0 else np.nan,
                "adjustflag": "3",
            })

    return pd.DataFrame(rows, columns=MINUTE_BAR_COLUMNS)


def make_generated_derived_features(
    daily: pd.DataFrame | None = None,
) -> pd.DataFrame:
    """Return numeric daily derived values, including NaN and infinity cells."""
    daily = make_generated_daily_bars() if daily is None else daily.copy()
    keys = (
        daily[["symbol_id", "date"]]
        .drop_duplicates()
        .sort_values(["date", "symbol_id"])
        .reset_index(drop=True)
    )
    date_rank = keys["date"].rank(method="dense").astype(float)
    symbol_rank = keys["symbol_id"].map({
        "sh600000": 1.0,
        "sz000001": 2.0,
        "sh600519": 3.0,
        "sz300750": 4.0,
    })
    out = keys.copy()
    out["toy_feature"] = symbol_rank + date_rank / 100.0
    out["finite_feature"] = symbol_rank * date_rank
    out["nan_feature"] = out["toy_feature"]
    out["inf_feature"] = out["toy_feature"]
    if len(out) >= 2:
        out.loc[out.index[0], "nan_feature"] = np.nan
        out.loc[out.index[1], "inf_feature"] = np.inf
    return out


def make_generated_alpha_weights(
    alpha_name: str = "alpha_a",
    *,
    scale: float = 1.0,
    offset: float = 0.0,
    zero_date_index: int | None = None,
    n_sessions: int = 10,
) -> pd.DataFrame:
    """Create a deterministic long alpha grid with optional zero-gross date."""
    dates = generated_sessions(n_sessions)
    even = np.array([1.20, -0.80, 0.40, -0.80], dtype=float)
    odd = np.array([-0.60, 1.10, -0.90, 0.40], dtype=float)
    rows: list[dict[str, object]] = []
    for i, date in enumerate(dates):
        vector = even.copy() if i % 2 == 0 else odd.copy()
        vector = vector + offset
        vector = vector - vector.mean()
        if zero_date_index is not None and i == zero_date_index:
            vector = np.zeros_like(vector)
        for sym, weight in zip(GENERATED_SYMBOLS, scale * vector):
            rows.append({
                "symbol_id": sym,
                "date": date,
                "alpha_name": alpha_name,
                "weight": float(weight),
            })
    result = pd.DataFrame(rows, columns=ALPHA_COLUMNS)
    return result.sort_values(["symbol_id", "date"]).reset_index(drop=True)


def make_generated_combo_weights(
    combo_name: str = "combo",
    *,
    zero_date_index: int | None = 2,
    n_sessions: int = 10,
) -> pd.DataFrame:
    """Create deterministic combo weights for portfolio construction tests."""
    alpha = make_generated_alpha_weights(
        "combo_source",
        zero_date_index=zero_date_index,
        n_sessions=n_sessions,
    )
    combo = alpha.rename(columns={"alpha_name": "combo_name"}).copy()
    combo["combo_name"] = combo_name
    return combo[COMBO_COLUMNS].sort_values(["symbol_id", "date"]).reset_index(drop=True)