visualization.py

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd

class YieldVisualization:
    def __init__(self):
        """
        初始化可视化类
        """
        self.fig = None

    def create_yield_comparison_chart(self, strategy_returns, benchmark_returns, strategy_name="策略", benchmark_name="基准"):
        """
        创建收益率对比图

        参数:
        strategy_returns: 策略收益率数据 (pandas Series或list)
        benchmark_returns: 基准收益率数据 (pandas Series或list)
        strategy_name: 策略名称
        benchmark_name: 基准名称

        返回:
        plotly.graph_objects.Figure: 图表对象
        """
        # 创建图表
        self.fig = go.Figure()

        # 添加策略收益率曲线
        self.fig.add_trace(
            go.Scatter(
                x=list(range(len(strategy_returns))),
                y=strategy_returns,
                mode='lines',
                name=strategy_name,
                line=dict(color='blue', width=2)
            )
        )

        # 添加基准收益率曲线
        self.fig.add_trace(
            go.Scatter(
                x=list(range(len(benchmark_returns))),
                y=benchmark_returns,
                mode='lines',
                name=benchmark_name,
                line=dict(color='red', width=2, dash='dash')
            )
        )

        # 设置图表布局
        self.fig.update_layout(
            title=f"{strategy_name} vs {benchmark_name} 收益率对比",
            xaxis_title="时间",
            yaxis_title="累计收益率",
            legend_title="图例",
            hovermode='x unified',
            template='plotly_white'
        )

        return self.fig

    def add_interactive_features(self):
        """
        添加交互功能
        """
        if self.fig is not None:
            # 添加时间范围选择器
            self.fig.update_xaxes(
                rangeselector=dict(
                    buttons=list([
                        dict(count=1, label="1M", step="month", stepmode="backward"),
                        dict(count=6, label="6M", step="month", stepmode="backward"),
                        dict(count=1, label="YTD", step="year", stepmode="todate"),
                        dict(count=1, label="1Y", step="year", stepmode="backward"),
                        dict(step="all")
                    ])
                ),
                rangeslider_visible=True
            )

    def save_chart(self, filename="yield_comparison.html"):
        """
        保存图表为HTML文件

        参数:
        filename: 保存的文件名
        """
        if self.fig is not None:
            self.fig.write_html(filename)
            print(f"图表已保存为 {filename}")

    def show_chart(self):
        """
        显示图表
        """
        if self.fig is not None:
            self.fig.show()