okx-codex-trader/scripts/stress_long_short_fusion.py

from __future__ import annotations

import argparse
import sys
from itertools import product
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from scripts.search_long_short_fusion import (
    OUTPUT_DIR,
    PREFIX,
    INITIAL_EQUITY,
    build_components,
    component_returns,
    horizon_return_fields,
    markdown_table,
    metrics,
)


STRESS_PREFIX = "fusion-stress"


def stressed_equity(
    components: dict[str, pd.Series],
    weights: dict[str, float],
    *,
    fee_single_side: float,
    slippage: float,
    funding_8h: float,
) -> pd.Series:
    returns = pd.DataFrame({name: component_returns(series) for name, series in components.items()}).dropna()
    combined = sum(returns[name] * weights.get(name, 0.0) for name in returns.columns)
    short_weight = sum(weight for name, weight in weights.items() if "short" in name)
    active_short = returns[[name for name in returns.columns if "short" in name]].abs().sum(axis=1) > 0.0
    turnover_proxy = sum(abs(weight) for weight in weights.values())
    base_fee = 0.0004
    extra_fee = max(0.0, fee_single_side - base_fee) * turnover_proxy / 30.0
    extra_slippage = slippage * turnover_proxy / 30.0
    funding_drag = active_short.astype(float) * short_weight * max(0.0, -funding_8h) * 3.0
    stressed = combined - extra_fee - extra_slippage - funding_drag
    equity = INITIAL_EQUITY * (1.0 + stressed).cumprod()
    equity.name = "equity"
    return equity


def parse_weights(row: pd.Series) -> dict[str, float]:
    long_key = str(row["long_variant"])
    return {
        long_key: float(row["long_weight"]),
        "btc_risk_short": float(row["btc_risk_short_weight"]),
        "eth_4h_vol_short": float(row["eth_4h_vol_short_weight"]),
        "btc_4h_vol_short": float(row["btc_4h_vol_short_weight"]),
        "eth_4h_vol_short_gated": float(row["eth_4h_vol_short_gated_weight"]),
        "btc_4h_vol_short_gated": float(row["btc_4h_vol_short_gated_weight"]),
    }


def report_text(command: str, paths: list[Path], robustness: pd.DataFrame, worst: pd.DataFrame) -> str:
    return "\n".join(
        [
            "# Long Short Fusion Stress Test",
            "",
            f"Run command: `{command}`",
            "",
            "Output files:",
            *[f"- `{path}`" for path in paths],
            "",
            "Stress grid: fee 0.04/0.08/0.10% single-side, slippage 0/0.05/0.10%, funding every 8h -0.005/0/+0.005%. Negative funding is short pay.",
            "",
            "## Robustness Summary",
            "",
            markdown_table(robustness),
            "",
            "## Worst Cases",
            "",
            markdown_table(worst),
            "",
        ]
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--years", type=float, default=8.0)
    parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR)
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    selected = pd.read_csv(args.output_dir / f"{PREFIX}-selected.csv")
    components_full = build_components(args.years)
    components = {name: item[1] for name, item in components_full.items()}
    rows = []
    for _, selected_row in selected.iterrows():
        weights = parse_weights(selected_row)
        for fee, slippage, funding in product((0.0004, 0.0008, 0.0010), (0.0, 0.0005, 0.0010), (-0.00005, 0.0, 0.00005)):
            equity = stressed_equity(
                components,
                weights,
                fee_single_side=fee,
                slippage=slippage,
                funding_8h=funding,
            )
            rows.append(
                {
                    "name": selected_row["name"],
                    "fee_single_side": fee,
                    "slippage": slippage,
                    "funding_8h": funding,
                    "long_variant": selected_row["long_variant"],
                    "long_weight": selected_row["long_weight"],
                    "short_exposure": selected_row["short_exposure"],
                    **horizon_return_fields(equity),
                    **metrics(equity),
                }
            )
    totals = pd.DataFrame(rows)
    robustness_rows = []
    for name, group in totals.groupby("name"):
        worst = group.sort_values(["total_return", "max_drawdown"], ascending=[True, False]).iloc[0]
        robustness_rows.append(
            {
                "name": name,
                "stress_cases": len(group),
                "positive_cases": int((group["total_return"] > 0.0).sum()),
                "worst_total_return": float(group["total_return"].min()),
                "worst_annualized_return": float(group["annualized_return"].min()),
                "worst_max_drawdown": float(group["max_drawdown"].max()),
                "worst_calmar": float(group["calmar"].min()),
                "worst_return_3y": float(group["h3y_return"].min()),
                "worst_return_1y": float(group["h1y_return"].min()),
                "worst_return_6m": float(group["h6m_return"].min()),
                "worst_return_3m": float(group["h3m_return"].min()),
                "worst_fee_single_side": float(worst["fee_single_side"]),
                "worst_slippage": float(worst["slippage"]),
                "worst_funding_8h": float(worst["funding_8h"]),
            }
        )
    robustness = pd.DataFrame(robustness_rows).sort_values(["worst_calmar", "worst_total_return"], ascending=[False, False])
    worst = totals.sort_values(["total_return", "max_drawdown"], ascending=[True, False]).head(20)

    totals_path = args.output_dir / f"{STRESS_PREFIX}-totals.csv"
    robustness_path = args.output_dir / f"{STRESS_PREFIX}-robustness.csv"
    worst_path = args.output_dir / f"{STRESS_PREFIX}-worst.csv"
    report_path = args.output_dir / f"{STRESS_PREFIX}-report.md"
    totals.to_csv(totals_path, index=False)
    robustness.to_csv(robustness_path, index=False)
    worst.to_csv(worst_path, index=False)
    report_path.write_text(
        report_text(
            f"rtk .venv/bin/python scripts/stress_long_short_fusion.py --years {args.years}",
            [totals_path, robustness_path, worst_path, report_path],
            robustness,
            worst,
        ),
        encoding="utf-8",
    )
    print(report_path)
    print(robustness.head(10).to_string(index=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())