sensitivity_analysis.py

from __future__ import annotations

import os
from dataclasses import asdict, replace

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from Backtest_utils import BacktestConfig
from main_update import load_dual_clock_inputs, run_backtest


DATA_DIR = "data"
OUTPUT_EXCEL = os.path.join(DATA_DIR, "sensitivity_analysis.xlsx")
HEATMAP_PNG = os.path.join(DATA_DIR, "sensitivity_heatmaps.png")


TOP_K_GRID = [40, 50, 60]
BETA_LOOKBACK_GRID = [48, 60, 72]
W_MAX_GRID = [0.05, 0.10, 0.15]
LAMBDA_SIGNAL_GRID = [0.01, 0.05, 0.10]


def evaluate_config(
    price_df: pd.DataFrame,
    returns_df: pd.DataFrame,
    ff5_realized_df: pd.DataFrame,
    benchmark_df: pd.DataFrame,
    config: BacktestConfig,
    ff5_signal_df: pd.DataFrame,
) -> dict:
    detail_df, metrics_df = run_backtest(
        price_df=price_df,
        returns_df=returns_df,
        ff5_realized_df=ff5_realized_df,
        benchmark_df=benchmark_df,
        config=config,
        ff5_signal_df=ff5_signal_df,
        show_progress=False,
    )

    if metrics_df.empty:
        result = {
            "months": 0,
            "cagr": np.nan,
            "annual_vol": np.nan,
            "sharpe": np.nan,
            "max_drawdown": np.nan,
            "hit_rate": np.nan,
            "annual_turnover": np.nan,
            "ending_nav": np.nan,
        }
    else:
        metrics = metrics_df.loc["signal_tilt"]
        result = {
            "months": metrics["months"],
            "cagr": metrics["cagr"],
            "annual_vol": metrics["annual_vol"],
            "sharpe": metrics["sharpe"],
            "max_drawdown": metrics["max_drawdown"],
            "hit_rate": metrics["hit_rate"],
            "annual_turnover": metrics["annual_turnover"],
            "ending_nav": detail_df["signal_tilt_nav"].iloc[-1] if not detail_df.empty else np.nan,
        }

    result.update(asdict(config))
    return result


def run_heatmap_grid(
    price_df: pd.DataFrame,
    returns_df: pd.DataFrame,
    ff5_realized_df: pd.DataFrame,
    benchmark_df: pd.DataFrame,
    base_config: BacktestConfig,
    ff5_signal_df: pd.DataFrame,
    row_param: str,
    row_values: list,
    col_param: str,
    col_values: list,
) -> pd.DataFrame:
    rows = []
    for row_value in row_values:
        for col_value in col_values:
            config = replace(base_config, **{row_param: row_value, col_param: col_value})
            result = evaluate_config(price_df, returns_df, ff5_realized_df, benchmark_df, config, ff5_signal_df)
            result["grid_name"] = f"{row_param}_vs_{col_param}"
            result["row_param"] = row_param
            result["row_value"] = row_value
            result["col_param"] = col_param
            result["col_value"] = col_value
            rows.append(result)

    return pd.DataFrame(rows)


def draw_heatmap(
    ax,
    heatmap_df: pd.DataFrame,
    title: str,
    row_param: str,
    col_param: str,
    metric_name: str,
    value_fmt: str,
):
    pivot_df = heatmap_df.pivot(index="row_value", columns="col_value", values=metric_name)
    image = ax.imshow(pivot_df.values, cmap="YlGnBu", aspect="auto")
    ax.set_xticks(range(len(pivot_df.columns)))
    ax.set_xticklabels([f"{value:.2f}" if isinstance(value, float) else str(value) for value in pivot_df.columns])
    ax.set_yticks(range(len(pivot_df.index)))
    ax.set_yticklabels([f"{value:.2f}" if isinstance(value, float) else str(value) for value in pivot_df.index])
    ax.set_xlabel(col_param)
    ax.set_ylabel(row_param)
    ax.set_title(title)

    for i in range(pivot_df.shape[0]):
        for j in range(pivot_df.shape[1]):
            value = pivot_df.iloc[i, j]
            text = "nan" if pd.isna(value) else format(value, value_fmt)
            ax.text(j, i, text, ha="center", va="center", color="black", fontsize=8)

    plt.colorbar(image, ax=ax, fraction=0.046, pad=0.04)


def plot_heatmaps(all_results: dict[str, pd.DataFrame]):
    fig, axes = plt.subplots(2, 2, figsize=(15, 11))
    plot_specs = [
        ("top_k_vs_beta_lookback", axes[0, 0], "Sharpe: top_k vs beta_lookback", "top_k", "beta_lookback", "sharpe", ".2f"),
        ("lambda_signal_vs_w_max", axes[0, 1], "Sharpe: lambda_signal vs w_max", "lambda_signal", "w_max", "sharpe", ".2f"),
        ("top_k_vs_beta_lookback", axes[1, 0], "CAGR: top_k vs beta_lookback", "top_k", "beta_lookback", "cagr", ".1%"),
        ("lambda_signal_vs_w_max", axes[1, 1], "CAGR: lambda_signal vs w_max", "lambda_signal", "w_max", "cagr", ".1%"),
    ]

    for grid_name, axis, title, row_param, col_param, metric_name, value_fmt in plot_specs:
        draw_heatmap(
            ax=axis,
            heatmap_df=all_results[grid_name],
            title=title,
            row_param=row_param,
            col_param=col_param,
            metric_name=metric_name,
            value_fmt=value_fmt,
        )

    plt.tight_layout()
    plt.savefig(HEATMAP_PNG, dpi=200, bbox_inches="tight")
    plt.show()


def export_results(base_config: BacktestConfig, all_results: dict[str, pd.DataFrame]):
    config_df = pd.DataFrame([asdict(base_config)])
    with pd.ExcelWriter(OUTPUT_EXCEL) as writer:
        config_df.to_excel(writer, sheet_name="base_config", index=False)
        for sheet_name, result_df in all_results.items():
            result_df.to_excel(writer, sheet_name=sheet_name[:31], index=False)


def print_summary(all_results: dict[str, pd.DataFrame]):
    for grid_name, result_df in all_results.items():
        best_sharpe = result_df.sort_values("sharpe", ascending=False).head(1)
        print(f"\nBest Sharpe in {grid_name}:")
        print(
            best_sharpe[
                ["row_param", "row_value", "col_param", "col_value", "sharpe", "cagr", "max_drawdown"]
            ].to_string(index=False)
        )


def main():
    base_config = BacktestConfig()
    price_df, returns_df, ff5_realized_df, ff5_signal_df, benchmark_df = load_dual_clock_inputs()

    all_results = {
        "top_k_vs_beta_lookback": run_heatmap_grid(
            price_df=price_df,
            returns_df=returns_df,
            ff5_realized_df=ff5_realized_df,
            benchmark_df=benchmark_df,
            base_config=base_config,
            ff5_signal_df=ff5_signal_df,
            row_param="top_k",
            row_values=TOP_K_GRID,
            col_param="beta_lookback",
            col_values=BETA_LOOKBACK_GRID,
        ),
        "lambda_signal_vs_w_max": run_heatmap_grid(
            price_df=price_df,
            returns_df=returns_df,
            ff5_realized_df=ff5_realized_df,
            benchmark_df=benchmark_df,
            base_config=base_config,
            ff5_signal_df=ff5_signal_df,
            row_param="lambda_signal",
            row_values=LAMBDA_SIGNAL_GRID,
            col_param="w_max",
            col_values=W_MAX_GRID,
        ),
    }

    export_results(base_config, all_results)
    print_summary(all_results)
    plot_heatmaps(all_results)
    print(f"\nSensitivity results saved to: {OUTPUT_EXCEL}")
    print(f"Heatmaps saved to: {HEATMAP_PNG}")


if __name__ == "__main__":
    main()