chinese-equity-quant/analysis/report.py

"""Performance analysis and reporting for backtest results."""
import os
from typing import Any

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import pandas as pd  # noqa: E402


def print_results(results: list, initial_cash: float = 1_000_000.0) -> dict[str, Any]:
    """Print and return key performance metrics from a backtrader run result."""
    if not results:
        print("No results to report.")
        return {}

    result = results[0]
    report = {}

    # Sharpe ratio
    sharpe = result.analyzers.sharpe.get_analysis()
    report["sharpe"] = sharpe.get("sharperatio", "N/A")

    # Drawdown
    dd = result.analyzers.drawdown.get_analysis()
    report["max_drawdown"] = dd.get("max", {}).get("drawdown", "N/A")
    report["max_drawdown_len"] = dd.get("max", {}).get("len", "N/A")

    # Returns
    rets = result.analyzers.returns.get_analysis()
    report["total_return"] = rets.get("rtot", "N/A")
    report["avg_return"] = rets.get("ravg", "N/A")

    # Trades
    trades = result.analyzers.trades.get_analysis()
    report["total_trades"] = trades.get("total", {}).get("total", 0)
    report["won_trades"] = trades.get("won", {}).get("total", 0)
    report["lost_trades"] = trades.get("lost", {}).get("total", 0)

    # Print
    print("=" * 50)
    print("BACKTEST RESULTS")
    print("=" * 50)
    print(f"Sharpe Ratio:        {report['sharpe']}")
    print(f"Total Return:        {report['total_return']:.4%}" if isinstance(report['total_return'], float) else f"Total Return:        {report['total_return']}")
    print(f"Max Drawdown:        {report['max_drawdown']:.2%}" if isinstance(report['max_drawdown'], float) else f"Max Drawdown:        {report['max_drawdown']}")
    print(f"Max DD Length:       {report['max_drawdown_len']}")
    print(f"Total Trades:        {report['total_trades']}")
    print(f"Won/Lost:            {report['won_trades']}/{report['lost_trades']}")
    print("=" * 50)

    return report


def plot_accumulated_pnl(
    results: list, output_path: str = "reports/pnl.png", initial_cash: float = 1_000_000.0
) -> str:
    """Plot accumulated portfolio value from a backtest run.

    Reads the per-day TimeReturn analyzer attached by ``BacktestRunner`` and
    compounds it into an equity curve.

    Args:
        results: The list returned by ``cerebro.run()``.
        output_path: Destination PNG path.
        initial_cash: Starting portfolio value for scaling the curve.

    Returns:
        The path the chart was written to.
    """
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    series = pd.Series(dtype=float)
    if results:
        tr = results[0].analyzers.timereturn.get_analysis()
        series = pd.Series(tr).sort_index()

    fig, ax = plt.subplots(figsize=(10, 5))
    if len(series):
        equity = (1.0 + series).cumprod() * initial_cash
        ax.plot(equity.index, equity.values, color="C0")
    ax.set_title("Accumulated Portfolio Value")
    ax.set_xlabel("Date")
    ax.set_ylabel("Value")
    ax.grid(True, alpha=0.3)
    fig.autofmt_xdate()
    fig.tight_layout()
    fig.savefig(output_path, dpi=100)
    plt.close(fig)
    return output_path


def plot_ic(signal_eval: dict, output_path: str = "reports/ic.png") -> str:
    """Plot the per-period rank IC time series from a signal evaluation.

    Args:
        signal_eval: Dict returned by ``evaluate_cross_sectional`` (expects a
            ``rank_ic_series`` pandas Series).
        output_path: Destination PNG path.

    Returns:
        The path the chart was written to.
    """
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    rank_ic = signal_eval.get("rank_ic_series", pd.Series(dtype=float))

    fig, ax = plt.subplots(figsize=(10, 5))
    if len(rank_ic):
        ax.bar(rank_ic.index, rank_ic.values, width=1.0, color="C1", alpha=0.6, label="Rank IC")
        ax.axhline(rank_ic.mean(), color="C3", linestyle="--", label=f"mean={rank_ic.mean():.3f}")
        ax.legend()
    ax.set_title("Cross-Sectional Rank IC")
    ax.set_xlabel("Date")
    ax.set_ylabel("Rank IC")
    ax.grid(True, alpha=0.3)
    fig.autofmt_xdate()
    fig.tight_layout()
    fig.savefig(output_path, dpi=100)
    plt.close(fig)
    return output_path


def generate_report(
    results: list,
    signal_eval: dict,
    output_dir: str = "reports/",
    initial_cash: float = 1_000_000.0,
) -> dict[str, str]:
    """Generate the full report: PnL chart, IC chart, and a summary text file.

    Args:
        results: The list returned by ``cerebro.run()``.
        signal_eval: Dict returned by ``evaluate_cross_sectional``.
        output_dir: Directory to write artifacts into.
        initial_cash: Starting portfolio value.

    Returns:
        Mapping of artifact name to file path.
    """
    os.makedirs(output_dir, exist_ok=True)
    pnl_path = plot_accumulated_pnl(
        results, os.path.join(output_dir, "pnl.png"), initial_cash
    )
    ic_path = plot_ic(signal_eval, os.path.join(output_dir, "ic.png"))

    metrics = print_results(results, initial_cash)
    summary_path = os.path.join(output_dir, "summary.txt")
    with open(summary_path, "w") as f:
        f.write("BACKTEST SUMMARY\n")
        f.write("=" * 40 + "\n")
        for k, v in metrics.items():
            f.write(f"{k}: {v}\n")
        f.write("\nSIGNAL IC\n")
        f.write("=" * 40 + "\n")
        for k in ("ic_mean", "ic_std", "ir", "rank_ic_mean", "rank_ic_std", "rank_ir", "hit_rate", "n_periods"):
            if k in signal_eval:
                f.write(f"{k}: {signal_eval[k]}\n")

    return {"pnl": pnl_path, "ic": ic_path, "summary": summary_path}