#!/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()