"""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)