myquant/finance/optimizer/engine.py

"""
Optuna 优化引擎。

统一接口：optimizer.optimize(strategy_class, space, price_df, factor_df) → OptimizationResult
"""

import copy
import time

import numpy as np
import optuna
import pandas as pd

from backtest.base import BaseStrategy
from backtest.report import BacktestReport
from backtest.vectorbt.engine import VectorBTEngine
from optimizer.objectives import Objective
from optimizer.result import OptimizationResult, WalkForwardResult
from optimizer.space import SearchSpace

# 抑制 Optuna 日志
optuna.logging.set_verbosity(optuna.logging.WARNING)


class OptunaEngine:
    """
    Optuna 优化引擎。
    """

    def __init__(self, bt_engine: VectorBTEngine | None = None):
        self.bt_engine = bt_engine or VectorBTEngine()

    def optimize(
        self,
        strategy_class: type[BaseStrategy],
        search_space: SearchSpace,
        price_df: pd.DataFrame,
        factor_df: pd.DataFrame | None = None,
        metric: str = "sharpe",
        n_trials: int = 100,
        direction: str = "maximize",
        sampler: optuna.samplers.BaseSampler | None = None,
    ) -> OptimizationResult:
        """
        参数寻优。

        参数:
            strategy_class: 策略类
            search_space: 搜索空间
            price_df: 价格数据
            factor_df: 因子数据
            metric: 优化目标
            n_trials: 试验次数
            direction: 'maximize' | 'minimize'
            sampler: Optuna 采样器，默认 TPESampler
        """
        if sampler is None:
            sampler = optuna.samplers.TPESampler(seed=42)

        study = optuna.create_study(
            direction=direction,
            sampler=sampler,
        )

        objective = Objective(
            strategy_class=strategy_class,
            search_space=search_space,
            price_df=price_df,
            factor_df=factor_df,
            bt_engine=self.bt_engine,
            metric=metric,
        )

        t0 = time.time()
        study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
        elapsed = time.time() - t0

        # 用最优参数跑一次完整回测
        best_params = study.best_params
        try:
            best_strategy = strategy_class(**best_params)
        except TypeError:
            valid = {k: v for k, v in best_params.items()
                     if k in strategy_class.__init__.__code__.co_varnames}
            best_strategy = strategy_class(**valid)

        best_report = self.bt_engine.run(
            best_strategy,
            price_df,
            price_df if factor_df is None else factor_df,
        )

        # 参数重要性
        try:
            importance = optuna.importance.get_param_importances(study)
        except Exception:
            importance = {}

        # 试验记录
        trials_df = study.trials_dataframe()

        return OptimizationResult(
            best_params=study.best_params,
            best_value=study.best_value,
            metric=metric,
            best_report=best_report,
            trials_df=trials_df,
            param_importance=importance,
        )

    def optimize_walk_forward(
        self,
        strategy_class: type[BaseStrategy],
        search_space: SearchSpace,
        price_df: pd.DataFrame,
        factor_df: pd.DataFrame | None = None,
        metric: str = "sharpe",
        n_trials: int = 100,
        train_window: int = 252 * 3,
        test_window: int = 252,
    ) -> WalkForwardResult:
        """
        滚动窗口优化（Walk-Forward Analysis）。

        每一步：train_window 训练 → test_window 验证 → 滑动。
        """
        if factor_df is None:
            factor_df = price_df

        n_total = len(price_df)
        windows = []
        test_equities = []
        param_history = []

        start = 0
        while start + train_window + test_window <= n_total:
            train_slice = slice(start, start + train_window)
            test_slice = slice(start + train_window, start + train_window + test_window)

            train_price = price_df.iloc[train_slice]
            train_factor = factor_df.iloc[train_slice]
            test_price = price_df.iloc[test_slice]
            test_factor = factor_df.iloc[test_slice]

            # 训练集上优化
            opt_result = self.optimize(
                strategy_class=strategy_class,
                search_space=search_space,
                price_df=train_price,
                factor_df=train_factor,
                metric=metric,
                n_trials=n_trials,
            )

            # 测试集上验证
            try:
                test_strategy = strategy_class(**opt_result.best_params)
            except TypeError:
                valid = {k: v for k, v in opt_result.best_params.items()
                         if k in strategy_class.__init__.__code__.co_varnames}
                test_strategy = strategy_class(**valid)

            test_report = self.bt_engine.run(test_strategy, test_price, test_factor)

            if len(test_report.equity_curve) > 0:
                test_equities.append(test_report.equity_curve)

            train_idx = train_price.index
            test_idx = test_price.index
            windows.append({
                "train_start": train_idx[0] if len(train_idx) > 0 else "",
                "train_end": train_idx[-1] if len(train_idx) > 0 else "",
                "test_start": test_idx[0] if len(test_idx) > 0 else "",
                "test_end": test_idx[-1] if len(test_idx) > 0 else "",
                "best_params": opt_result.best_params,
                "best_value": opt_result.best_value,
                "test_return": test_report.total_return,
                "test_sharpe": test_report.sharpe_ratio,
                "test_mdd": test_report.max_drawdown,
            })
            param_history.append(opt_result.best_params)

            start += test_window

        # 合并测试期权益曲线
        consolidated = _merge_test_periods(test_equities, self.bt_engine.initial_capital)

        # 参数稳定性
        param_df = pd.DataFrame(param_history) if param_history else pd.DataFrame()
        if not param_df.empty:
            param_df.index.name = "window"

        return WalkForwardResult(
            windows=windows,
            consolidated_report=consolidated,
            param_stability=param_df,
        )


def _merge_test_periods(
    equity_list: list[pd.Series],
    initial_capital: float = 100_000,
) -> BacktestReport | None:
    """拼接各窗口测试期权益曲线为一个连续序列。"""
    if not equity_list:
        return None

    merged = pd.concat(equity_list)
    merged = merged.sort_index()
    merged = merged[~merged.index.duplicated()]

    # 确保 DatetimeIndex
    if not isinstance(merged.index, pd.DatetimeIndex):
        merged.index = pd.to_datetime(merged.index, format="%Y%m%d")

    dd = merged / merged.cummax() - 1
    daily_ret = merged.pct_change().dropna()
    years = max(len(daily_ret) / 252, 0.02)

    total_ret = (merged.iloc[-1] / merged.iloc[0] - 1) * 100
    cagr = ((total_ret / 100 + 1) ** (1 / years) - 1) * 100
    mdd = dd.min() * 100
    std_ret = daily_ret.std() * np.sqrt(252)
    sharpe = (daily_ret.mean() * 252) / std_ret if std_ret > 0 else 0
    calmar = cagr / abs(mdd) if abs(mdd) > 0 else 0

    try:
        monthly = merged.resample("ME").last().pct_change()
    except Exception:
        monthly = pd.Series(dtype=float)

    return BacktestReport(
        total_return=round(total_ret, 2),
        cagr=round(cagr, 2),
        max_drawdown=round(mdd, 2),
        sharpe_ratio=round(sharpe, 2),
        calmar_ratio=round(calmar, 2),
        annual_volatility=round(std_ret * 100 if std_ret != 0 else 0, 2),
        equity_curve=merged,
        drawdown_curve=dd,
        monthly_returns=monthly,
    )