Direction9/experiments/visualization_direction9.py

#!/usr/bin/env python
"""
方向 9 实验可视化模块

功能:
- 绘制 MMLU vs r_eff 散点图
- 绘制 RS_cross 跨域分层图
- 绘制 SPD 模型间谱距离热力图
- 绘制逐层谱收敛轮廓 (LCP)
- 绘制能力涌现 r_eff 曲线

用法:
    python experiments/visualization_direction9.py --input-dir output/cross_model
"""

import os
import json
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Any
from datetime import datetime

import matplotlib
matplotlib.use('Agg')  # 非交互式后端
import matplotlib.pyplot as plt
import seaborn as sns


class Direction9Visualizer:
    """方向 9 实验可视化器"""

    def __init__(self, output_dir: str = "experiments/output/figures"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # 设置样式
        sns.set_style("whitegrid")
        sns.set_context("paper", font_scale=1.2)

        # 中文支持
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False

    def plot_mmlu_vs_r_eff(
        self,
        r_effs: Dict[str, Dict[str, float]],
        model_mmlu: Dict[str, float],
        output_name: str = "mmlu_vs_r_eff.png"
    ) -> str:
        """
        绘制 MMLU vs r_eff 散点图

        Args:
            r_effs: {dataset: {model: r_eff}}
            model_mmlu: {model: mmlu_score}
        """
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        axes = axes.flatten()

        datasets = list(r_effs.keys())
        colors = {'llama3.2-3b-instruct': '#1f77b4', 'Mistral-7B-v0.3': '#ff7f0e',
                  'llama3-8b': '#2ca02c', 'Qwen2.5-7B-Instruct': '#d62728',
                  'gemma-2-9b-it': '#9467bd'}

        for idx, dataset in enumerate(datasets):
            ax = axes[idx]
            data = r_effs[dataset]

            models = list(data.keys())
            mmlu_vals = [model_mmlu.get(m, 0) for m in models]
            r_eff_vals = [data[m] for m in models]

            for model in models:
                ax.scatter(
                    model_mmlu.get(model, 0),
                    data[model],
                    c=[colors.get(model, '#333333')],
                    s=100, alpha=0.7,
                    label=model
                )

            # 拟合趋势线
            if len(mmlu_vals) > 1:
                z = np.polyfit(mmlu_vals, r_eff_vals, 1)
                p = np.poly1d(z)
                ax.plot(sorted(mmlu_vals), p(sorted(mmlu_vals)),
                       '--', color='gray', alpha=0.5,
                       label=f'Trend (τ={np.corrcoef(mmlu_vals, r_eff_vals)[0,1]:.2f})')

            ax.set_xlabel('MMLU Score (%)')
            ax.set_ylabel('Effective Rank (r_eff)')
            ax.set_title(f'{dataset}')
            ax.legend(fontsize=8)

        plt.tight_layout()
        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)

    def plot_rs_cross_comparison(
        self,
        rs_cross: Dict[str, float],
        output_name: str = "rs_cross_comparison.png"
    ) -> str:
        """绘制 RS_cross 跨域分层比较图"""
        fig, ax = plt.subplots(figsize=(10, 6))

        datasets = list(rs_cross.keys())
        values = list(rs_cross.values())

        # 按 RS_cross 排序
        sorted_idx = np.argsort(values)
        datasets = [datasets[i] for i in sorted_idx]
        values = [values[i] for i in sorted_idx]

        colors = plt.cm.OrRd(np.linspace(0.3, 0.9, len(datasets)))

        bars = ax.bar(range(len(datasets)), values, color=colors)

        ax.set_xticks(range(len(datasets)))
        ax.set_xticklabels(datasets, rotation=45, ha='right')
        ax.set_ylabel('RS_cross (跨模型谱方差)')
        ax.set_title('跨域谱分层比较\n(越小表示跨模型谱收敛越好)')

        # 添加数值标签
        for bar, val in zip(bars, values):
            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                   f'{val:.3f}', ha='center', va='bottom', fontsize=10)

        plt.tight_layout()
        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)

    def plot_spd_heatmap(
        self,
        spd_matrix: List[List[float]],
        model_names: List[str],
        dataset_name: str = "all",
        output_name: str = "spd_heatmap.png"
    ) -> str:
        """绘制 SPD 模型间谱距离热力图"""
        spd = np.array(spd_matrix)

        fig, ax = plt.subplots(figsize=(10, 8))

        # 使用 seaborn 热力图
        sns.heatmap(spd, annot=True, fmt='.2f', cmap='YlOrRd',
                   xticklabels=model_names, yticklabels=model_names,
                   ax=ax, cbar_kws={'label': 'Wasserstein-1 Distance'})

        ax.set_title(f'SPD 模型间谱距离矩阵\n({dataset_name})')
        ax.set_xlabel('Model')
        ax.set_ylabel('Model')

        plt.tight_layout()
        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)

    def plot_lcp_profile(
        self,
        lcp_data: Dict[str, Any],
        output_name: str = "lcp_profile.png"
    ) -> str:
        """绘制逐层谱收敛轮廓 (LCP)"""
        layer_indices = lcp_data.get('layer_indices', [])
        lcp_values = lcp_data.get('lcp_values', [])
        dataset = lcp_data.get('dataset', 'unknown')

        fig, ax = plt.subplots(figsize=(10, 6))

        ax.plot(layer_indices, lcp_values, 'o-', linewidth=2, markersize=8)
        ax.fill_between(layer_indices, lcp_values, alpha=0.3)

        ax.set_xlabel('Layer Index')
        ax.set_ylabel('RS_cross')
        ax.set_title(f'逐层谱收敛轮廓\n({dataset})')
        ax.grid(True, alpha=0.3)

        # 标注最小值和最大值
        min_idx = np.argmin(lcp_values)
        max_idx = np.argmax(lcp_values)
        ax.annotate(f'Min: {lcp_values[min_idx]:.3f}',
                   xy=(layer_indices[min_idx], lcp_values[min_idx]),
                   xytext=(10, 10), textcoords='offset points',
                   bbox=dict(boxstyle='round', fc='w', ec='gray'))

        plt.tight_layout()
        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)

    def plot_emergence_curve(
        self,
        emergence_results: Dict[str, Any],
        output_name: str = "emergence_curves.png"
    ) -> str:
        """绘制能力涌现 r_eff 曲线"""
        n_abilities = len(emergence_results)
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        axes = axes.flatten()

        colors = plt.cm.viridis(np.linspace(0.1, 0.9, n_abilities))

        for idx, (ability, data) in enumerate(emergence_results.items()):
            ax = axes[idx]

            r_eff_curve = data.get('r_eff_curve', [])
            model_order = data.get('model_order', [])
            model_mmlu = data.get('model_mmlu', [])
            max_jump = data.get('max_jump', {})
            changepoints = data.get('changepoints', [])

            ax.plot(range(len(r_eff_curve)), r_eff_curve, 'o-',
                   linewidth=2, markersize=8, color=colors[idx])

            # 标注变化点
            for cp in changepoints:
                ax.axvline(x=cp, color='red', linestyle='--', alpha=0.5)
                ax.scatter([cp], [r_eff_curve[cp] if cp < len(r_eff_curve) else 0],
                          color='red', s=100, zorder=5)

            # 标注最大跃变点
            jump_idx = max_jump.get('model_index', 0)
            if jump_idx < len(r_eff_curve):
                ax.scatter([jump_idx], [r_eff_curve[jump_idx]],
                          color='orange', s=150, marker='*', zorder=5)

            ax.set_xlabel('Model Index (sorted by MMLU)')
            ax.set_ylabel('Effective Rank (r_eff)')
            ax.set_title(f'{ability}\nJump@{max_jump.get("model_name", "N/A")} '
                        f'(MMLU={max_jump.get("mmlu_at_jump", 0):.1f}%)')
            ax.set_xticks(range(len(model_order)))
            ax.set_xticklabels([m.split('-')[0] for m in model_order],
                            rotation=45, ha='right', fontsize=8)
            ax.grid(True, alpha=0.3)

        plt.tight_layout()
        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)

    def create_summary_dashboard(
        self,
        e1_result: Dict[str, Any],
        e3_result: Dict[str, Any],
        output_name: str = "summary_dashboard.png"
    ) -> str:
        """创建实验摘要仪表盘"""
        fig = plt.figure(figsize=(16, 12))

        # 子图布局
        gs = plt.GridSpec(2, 2, figure=fig)
        ax1 = fig.add_subplot(gs[0, 0])
        ax2 = fig.add_subplot(gs[0, 1])
        ax3 = fig.add_subplot(gs[1, :])

        # 1. RS_cross 柱状图
        rs_cross = e1_result.get('rs_cross', {})
        if rs_cross:
            datasets = list(rs_cross.keys())
            values = list(rs_cross.values())
            ax1.bar(range(len(datasets)), values, color='steelblue')
            ax1.set_xticks(range(len(datasets)))
            ax1.set_xticklabels(datasets, rotation=45, ha='right')
            ax1.set_ylabel('RS_cross')
            ax1.set_title('跨域谱分层')

        # 2. 能力涌现显著性
        if e3_result:
            abilities = list(e3_result.keys())
            significance = [e3_result[a].get('jump_significance', 0) for a in abilities]
            colors = ['green' if s > 2 else 'orange' if s > 1 else 'red' for s in significance]
            ax2.bar(range(len(abilities)), significance, color=colors)
            ax2.set_xticks(range(len(abilities)))
            ax2.set_xticklabels(abilities, rotation=45, ha='right')
            ax2.set_ylabel('Jump Significance (σ)')
            ax2.set_title('能力涌现显著性\n(绿>2σ, 橙>1σ, 红<1σ)')
            ax2.axhline(y=2, color='green', linestyle='--', alpha=0.5)
            ax2.axhline(y=1, color='orange', linestyle='--', alpha=0.5)

        # 3. 模型 r_eff 对比表
        ax3.axis('off')
        r_effs = e1_result.get('r_effs', {})
        if r_effs:
            table_data = []
            headers = ['Dataset'] + list(list(r_effs.values())[0].keys())
            for ds, r_dict in r_effs.items():
                row = [ds] + [f'{r_dict.get(m, 0):.1f}' for m in headers[1:]]
                table_data.append(row)

            table = ax3.table(
                cellText=table_data,
                colLabels=headers,
                loc='center',
                cellLoc='center',
                colColours=['lightblue'] * len(headers)
            )
            table.auto_set_font_size(False)
            table.set_fontsize(10)
            table.scale(1.2, 1.5)

        plt.suptitle(f'方向 9 实验摘要仪表盘\n生成时间：{datetime.now().strftime("%Y-%m-%d %H:%M")}',
                    fontsize=14, y=0.98)
        plt.tight_layout()

        output_path = self.output_dir / output_name
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"已保存：{output_path}")
        return str(output_path)


def main():
    import argparse

    parser = argparse.ArgumentParser(description="方向 9 实验可视化")
    parser.add_argument("--input-dir", type=str,
                       default="experiments/output/cross_model",
                       help="实验结果输入目录")
    parser.add_argument("--output-dir", type=str,
                       default="experiments/output/figures",
                       help="图片输出目录")

    args = parser.parse_args()

    visualizer = Direction9Visualizer(args.output_dir)

    # 加载 E1 结果
    e1_result_file = Path(args.input_dir) / "e1_result.json"
    if e1_result_file.exists():
        with open(e1_result_file) as f:
            e1_result = json.load(f)

        # 绘制 MMLU vs r_eff
        # 需要 model_mmlu 数据
        model_mmlu = {
            "llama3.2-3b-instruct": 58.0,
            "Mistral-7B-v0.3": 62.5,
            "llama3-8b": 68.4,
            "Qwen2.5-7B-Instruct": 86.3,
            "gemma-2-9b-it": 82.0
        }

        # 绘制 RS_cross 比较图
        if 'rs_cross' in e1_result:
            visualizer.plot_rs_cross_comparison(e1_result['rs_cross'])

        # 绘制 SPD 热力图
        if 'spd_matrices' in e1_result:
            for dataset, spd in e1_result['spd_matrices'].items():
                model_names = list(e1_result['r_effs'].get(dataset, {}).keys())
                if model_names:
                    visualizer.plot_spd_heatmap(spd, model_names, dataset)

    # 加载 E3 结果
    e3_result_file = Path(args.input_dir).parent / "emergence/emergence_results.json"
    if e3_result_file.exists():
        with open(e3_result_file) as f:
            e3_result = json.load(f)
        visualizer.plot_emergence_curve(e3_result)

    print("\n可视化完成!")


if __name__ == "__main__":
    main()