plot_comparison.py

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import argparse
import json
import os

plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

parser = argparse.ArgumentParser(description='绘制模型性能对比图表（支持归一化）')
parser.add_argument('--配置', '-f', type=str, default='data_config.json', help='数据配置文件路径')
parser.add_argument('--输出目录', '-d', type=str, default='charts', help='输出图表目录')
parser.add_argument('--合并分组', '-m', action='store_true', help='将第一层分组合并到一张图中')
args = parser.parse_args()

def fill_merged_cells(df, file_path, sheet_name):
    """填充合并单元格：用前一个非空值向下填充"""
    try:
        wb = pd.ExcelFile(file_path).book
        ws = wb[sheet_name]
        merged_ranges = ws.merged_cells.ranges
        
        if not merged_ranges:
            return df
        
        for merged_range in merged_ranges:
            min_col, min_row = merged_range.min_col, merged_range.min_row
            max_col, max_row = merged_range.max_col, merged_range.max_row
            
            first_cell = ws.cell(min_row, min_col).value
            
            for row in range(min_row + 1, max_row + 1):
                for col in range(min_col, max_col + 1):
                    ws.cell(row, col).value = first_cell
        
        df = pd.read_excel(file_path, sheet_name=sheet_name, engine='openpyxl')
    except Exception:
        pass
    
    for col in df.columns:
        df[col] = df[col].fillna(method='ffill')
        if pd.api.types.is_numeric_dtype(df[col]):
            col_values = df[col].dropna()
            if len(col_values) > 0 and col_values.apply(lambda x: float(x).is_integer() if pd.notna(x) else True).all():
                try:
                    df[col] = df[col].astype(int)
                except (ValueError, TypeError):
                    pass
    
    return df

def load_data_from_files(config):
    all_data = []
    files_config = config.get('files', [])
    
    for file_config in files_config:
        file_path = file_config.get('file')
        sheets = file_config.get('sheets', [])
        column_mapping = file_config.get('column_mapping', {})
        column_add = file_config.get('column_add', {})
        
        if not os.path.exists(file_path):
            print(f"文件不存在: {file_path}, 跳过")
            continue
        
        xl = pd.ExcelFile(file_path)
        
        if sheets is None or (isinstance(sheets, list) and len(sheets) == 0):
            sheets = xl.sheet_names
        else:
            sheets = [s for s in sheets if s]
        
        for sheet in sheets:
            try:
                df = pd.read_excel(file_path, sheet_name=sheet)
                df = fill_merged_cells(df, file_path, sheet)
                df.columns = df.columns.str.replace('\n', '').str.strip()
                
                if column_mapping:
                    df = df.rename(columns=column_mapping)
                
                column_replace = file_config.get('column_replace', {})
                for col, replace_dict in column_replace.items():
                    if col in df.columns:
                        df[col] = df[col].replace(replace_dict)
                
                df['source_file'] = file_path
                df['source_sheet'] = sheet
                all_data.append(df)
                print(f"读取: {file_path} - {sheet}, {len(df)} 行")

                for c in column_add:
                   df[c] = column_add[c]

            except Exception as e:
                print(f"读取失败: {file_path} - {sheet}: {e}")
    
    if not all_data:
        return pd.DataFrame()
    
    combined_df = pd.concat(all_data, ignore_index=True)
    return combined_df

def apply_filter(df, filter_dict):
    for filter_col, filter_values in filter_dict.items():
        if filter_col in df.columns and filter_values:
            if isinstance(filter_values, list):
                df = df[df[filter_col].isin(filter_values)]
            else:
                df = df[df[filter_col] == filter_values]
    return df

def parse_metric_cols(metric_cols, key_cols):
    """解析 metric_cols，分离普通列和需要归一化的列"""
    normal_metrics = []
    normalize_configs = []
    
    for m in metric_cols:
        if isinstance(m, dict):
            for col_name, base_value in m.items():
                if isinstance(base_value, list):
                    base_dict = {}
                    for i, k in enumerate(key_cols):
                        if i < len(base_value):
                            base_dict[k] = base_value[i]
                    normalize_configs.append({
                        'column': col_name,
                        'base_value': base_dict
                    })
                elif isinstance(base_value, dict):
                    normalize_configs.append({
                        'column': col_name,
                        'base_value': base_value
                    })
                else:
                    parts = str(base_value).split('_')
                    base_dict = {}
                    for i, k in enumerate(key_cols):
                        if i < len(parts):
                            base_dict[k] = parts[i]
                    normalize_configs.append({
                        'column': col_name,
                        'base_value': base_dict
                    })
        elif isinstance(m, list):
            for item in m:
                if isinstance(item, dict):
                    for col_name, base_value in item.items():
                        if isinstance(base_value, list):
                            base_dict = {}
                            for i, k in enumerate(key_cols):
                                if i < len(base_value):
                                    base_dict[k] = base_value[i]
                            normalize_configs.append({
                                'column': col_name,
                                'base_value': base_dict
                            })
                        else:
                            normalize_configs.append({
                                'column': col_name,
                                'base_value': base_value
                            })
                else:
                    normal_metrics.append(item)
        else:
            normal_metrics.append(m)
    
    return normal_metrics, normalize_configs

def apply_normalization(df, key_cols, normalize_configs, group_cols):
    """对指定列进行归一化处理（每个分组内独立归一化）"""
    df = df.copy()
    
    for config in normalize_configs:
        col = config['column']
        base_value = config['base_value']
        
        if col not in df.columns:
            continue
        
        if not key_cols:
            print(f"  警告: key_cols为空，跳过归一化")
            continue
        
        valid_key_cols = [k for k in key_cols if k in df.columns]
        if not valid_key_cols:
            print(f"  警告: 未找到有效的key_cols {key_cols}，跳过归一化")
            continue
        
        if isinstance(base_value, dict):
            base_dict = base_value
        else:
            base_dict = {}
            base_parts = base_value.split('_')
            for i, k in enumerate(valid_key_cols):
                if i < len(base_parts):
                    base_dict[k] = base_parts[i]
        
        effective_group_cols = [c for c in group_cols if c and c in df.columns]
        
        if effective_group_cols:
            def normalize_group(group):
                base_mask = pd.Series(True, index=group.index)
                for k, v in base_dict.items():
                    if k in group.columns:
                        base_mask = base_mask & (group[k] == v)
                
                if base_mask.any():
                    base_val = group.loc[base_mask, col].mean()
                    if base_val == 0 or pd.isna(base_val):
                        valid_data = group[group[col].notna() & (group[col] != 0)]
                        if len(valid_data) > 0:
                            base_val = valid_data[col].iloc[0]
                        else:
                            return group
                else:
                    valid_data = group[group[col].notna() & (group[col] != 0)]
                    if len(valid_data) > 0:
                        base_val = valid_data[col].iloc[0]
                    else:
                        return group
                
                if base_val == 0 or pd.isna(base_val):
                    return group
                group[col] = (group[col] / base_val) * 100
                return group
            df = df.groupby(effective_group_cols, group_keys=False).apply(normalize_group)
            print(f"  归一化列 '{col}': 每个分组内基准值 {base_dict} = 100%")
        else:
            base_mask = pd.Series(True, index=df.index)
            for k, v in base_dict.items():
                if k in df.columns:
                    base_mask = base_mask & (df[k] == v)
            if not base_mask.any():
                print(f"  警告: 未找到基准值 {base_dict}，跳过归一化")
                continue
            base_values = df.loc[base_mask, col].mean()
            if base_values == 0 or pd.isna(base_values):
                valid_data = df[df[col].notna() & (df[col] != 0)]
                if len(valid_data) > 0:
                    base_values = valid_data[col].iloc[0]
                else:
                    print(f"  警告: 无有效数据，跳过归一化")
                    continue
            df[col] = (df[col] / base_values) * 100
            print(f"  归一化列 '{col}': 基准值 {base_dict} = 100%")
    
    return df

def generate_chart(df_subset, output_path, colkey, outer_group_cols, inner_group_cols, metric_cols, 
                   normalize_configs=None, merge_groups=False):
    df_subset = df_subset.copy()
    compare_col = "ColKey"
    df_subset[compare_col] = df_subset[colkey].apply(lambda x: '_'.join(x.dropna().astype(str)), axis=1)
    
    all_group_cols = outer_group_cols + inner_group_cols
    if all_group_cols:
        df_grouped = df_subset[all_group_cols + [compare_col] + metric_cols].groupby(all_group_cols + [compare_col]).mean().reset_index()
    else:
        df_grouped = df_subset[[compare_col] + metric_cols].groupby([compare_col]).mean().reset_index()
        df_grouped[compare_col] = df_grouped.index
    
    if len(df_grouped) == 0:
        print(f"    无数据，跳过")
        return False
    
    if outer_group_cols:
        outer_values = df_grouped.groupby(outer_group_cols).size().reset_index()
    else:
        outer_values = pd.DataFrame({'': ['all']})
    
    n_outer = len(outer_values)
    engine_values = df_grouped[compare_col].unique()
    n_engines = len(engine_values)
    
    color_palette = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D', '#3B1F2B', '#95C623', '#7B2D26']
    colors = [color_palette[i % len(color_palette)] for i in range(n_engines)]
    
    normalized_cols = [c['column'] for c in (normalize_configs or [])]
    
    if merge_groups and n_outer > 1:
        fig, axes = plt.subplots(1, 4, figsize=(8 * n_outer + 20, 10))
        
        bar_width = 0.12
        bar_spacing = 0.05
        group_gap = 3
        
        x_labels_all = None
        
        for col, metric in enumerate(metric_cols):
            ax = axes[col]
            
            current_x = 0
            
            for row_idx, (_, outer_row) in enumerate(outer_values.iterrows()):
                df_outer = df_grouped.copy()
                for gcol in outer_group_cols:
                    df_outer = df_outer[df_outer[gcol] == outer_row[gcol]]
                
                outer_label_value = '-'.join([str(outer_row[gcol]) for gcol in outer_group_cols])
                
                pt = df_outer.pivot_table(
                    index=inner_group_cols, 
                    columns=compare_col, 
                    values=metric
                ).fillna(0)
                
                n_bars_per_group = len(pt)
                group_width = n_bars_per_group * n_engines * (bar_width + bar_spacing) + group_gap
                group_center = current_x + group_width / 2
                
                x_labels = ['/'.join([str(v) for v in idx]) for idx in pt.index]
                if x_labels_all is None:
                    x_labels_all = x_labels
                
                x = np.arange(len(x_labels)) * (n_engines * (bar_width + bar_spacing)) + current_x
                
                for i, engine in enumerate(engine_values):
                    if engine in pt.columns:
                        values = pt[engine].values
                        offset = i * bar_width
                        label = f"{engine} ({outer_label_value})"
                        bars = ax.bar(x + offset, values, bar_width, label=label, color=colors[i], edgecolor='white', linewidth=0.5)
                        
                        for bar, val in zip(bars, values):
                            if val > 0:
                                y_pos = bar.get_height() + bar.get_height()*0.02 if bar.get_height() > 0 else 1
                                if metric in normalized_cols:
                                    ax.text(bar.get_x() + bar.get_width()/2, y_pos, 
                                            f'{val:.0f}%', ha='center', va='bottom', fontsize=5, fontweight='bold')
                                else:
                                    ax.text(bar.get_x() + bar.get_width()/2, y_pos, 
                                            f'{val:.1f}', ha='center', va='bottom', fontsize=5, fontweight='bold')
                
                ax.axvline(x=current_x + n_bars_per_group * n_engines * (bar_width + bar_spacing) + group_gap/2, color='gray', linestyle='--', linewidth=1)
                
                ax.text(group_center, ax.get_ylim()[1] * 0.95, outer_label_value, 
                        ha='center', va='top', fontsize=9, fontweight='bold', 
                        bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
                
                current_x = current_x + n_bars_per_group * n_engines * (bar_width + bar_spacing) + group_gap
            
            total_inner_labels = len(x_labels_all)
            inner_positions = []
            inner_labels = []
            for gi in range(len(outer_values)):
                base_x = gi * (total_inner_labels * n_engines * (bar_width + bar_spacing) + group_gap)
                for xi in range(total_inner_labels):
                    center_pos = base_x + xi * n_engines * (bar_width + bar_spacing) + (n_engines * bar_width + (n_engines-1) * bar_spacing) / 2
                    inner_positions.append(center_pos)
                    inner_labels.append(x_labels_all[xi])
            
            ax.set_xticks(inner_positions)
            ax.set_xticklabels(inner_labels, rotation=45, ha='right', fontsize=6)
            
            ax.set_xlabel('/'.join(inner_group_cols), fontsize=9)
            if metric in normalized_cols:
                ax.set_ylabel(f'{metric} (%)', fontsize=10)
                ax.set_title(f'{metric} (归一化)', fontsize=12, fontweight='bold')
            else:
                ax.set_ylabel(metric, fontsize=10)
                ax.set_title(f'{metric}', fontsize=12, fontweight='bold')
            ax.grid(axis='y', alpha=0.3, linestyle='--')
            ax.legend(fontsize=5, loc='upper right', framealpha=0.9, ncol=1)
    else:
        fig, axes = plt.subplots(n_outer, 4, figsize=(24, 5 * n_outer))
        
        if n_outer == 1:
            axes = axes.reshape(1, -1)
        
        bar_width = 0.2
        
        outer_label = '/'.join(outer_group_cols) if outer_group_cols else '全部'
        
        for row_idx, (_, outer_row) in enumerate(outer_values.iterrows()):
            df_outer = df_grouped.copy()
            for col in outer_group_cols:
                df_outer = df_outer[df_outer[col] == outer_row[col]]
            
            outer_label_value = '-'.join([str(outer_row[col]) for col in outer_group_cols])
            
            for col, metric in enumerate(metric_cols):
                ax = axes[row_idx, col]
                
                pt = df_outer.pivot_table(
                    index=inner_group_cols, 
                    columns=compare_col, 
                    values=metric
                ).fillna(0)
                
                x_labels = ['/'.join([str(v) for v in idx]) for idx in pt.index]
                x = np.arange(len(x_labels))
                
                for i, engine in enumerate(engine_values):
                    if engine in pt.columns:
                        values = pt[engine].values
                        offset = (i - n_engines/2 + 0.5) * bar_width
                        bars = ax.bar(x + offset, values, bar_width, label=engine, color=colors[i], edgecolor='white', linewidth=0.5)
                        
                        for bar, val in zip(bars, values):
                            if val > 0:
                                y_pos = bar.get_height() + bar.get_height()*0.02 if bar.get_height() > 0 else 1
                                if metric in normalized_cols:
                                    ax.text(bar.get_x() + bar.get_width()/2, y_pos, 
                                            f'{val:.0f}%', ha='center', va='bottom', fontsize=7, fontweight='bold')
                                else:
                                    ax.text(bar.get_x() + bar.get_width()/2, y_pos, 
                                            f'{val:.1f}', ha='center', va='bottom', fontsize=7, fontweight='bold')
                
                ax.set_xlabel('/'.join(inner_group_cols), fontsize=9)
                if metric in normalized_cols:
                    ax.set_ylabel(f'{metric} (%)', fontsize=10)
                    ax.set_title(f'{outer_label}={outer_label_value} - {metric} (归一化)', fontsize=11, fontweight='bold')
                else:
                    ax.set_ylabel(metric, fontsize=10)
                    ax.set_title(f'{outer_label}={outer_label_value} - {metric}', fontsize=11, fontweight='bold')
                ax.set_xticks(x)
                ax.set_xticklabels(x_labels, rotation=45, ha='right', fontsize=7)
                ax.grid(axis='y', alpha=0.3, linestyle='--')
                
                ax.legend(fontsize=6, loc='upper right', framealpha=0.9, ncol=1)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
    plt.close()
    return True

print(f"从配置文件加载数据: {args.配置}")

with open(args.配置, 'r', encoding='utf-8') as f:
    config = json.load(f)

df = load_data_from_files(config)

if df.empty:
    print("未加载到数据")
    exit(1)

print(f"\n可用列名: {df.columns.tolist()}")

col_mapping = {}
for std_col, alt_cols in [
    ('模型', ['模型', 'model', 'Model']),
    ('卡类型', ['卡类型', 'card_type', '卡']),
    ('卡数', ['卡数', 'num_cards', '卡数', 'GPU数量']),
    ('vLLM版本', ['vLLM版本', 'vllm_version', 'vLLM版本']),
    ('V0/V1 Engine', ['V0/V1 Engine', 'Engine', 'engine']),
    ('输入长度(tokens)', ['输入长度(tokens)', 'input_length', 'input length', '输入长度']),
    ('输出长度(tokens)', ['输出长度(tokens)', 'output_length', 'output length', '输出长度']),
    ('并发数', ['并发数', 'concurrency', '并发', 'num_concurrent']),
    ('平均首字延时TTFT(ms)', ['平均首字延时TTFT(ms)', 'ttft', 'TTFT', '首字延时']),
    ('平均生成时间TPOT(ms)', ['平均生成时间TPOT(ms)', 'tpot', 'TPOT', '生成时间']),
    ('生成吞吐量(tokens/s)', ['生成吞吐量(tokens/s)', 'gen_throughput', '生成吞吐']),
    ('总吞吐量(tokens/s)', ['总吞吐量(tokens/s)', 'total_throughput', '总吞吐'])
]:
    for alt in alt_cols:
        if alt in df.columns:
            col_mapping[std_col] = alt

print(f"\n列映射: {col_mapping}")
df_renamed = df.rename(columns=col_mapping)

filter_config = config.get('filter', {})
df_renamed = apply_filter(df_renamed, filter_config)

print(f"过滤后数据量: {len(df_renamed)}")


dist_cols_config = config.get('dist_cols', ['模型', '卡数'])
dist_cols = [col_mapping.get(c, c) for c in dist_cols_config]
dist_cols = [c for c in dist_cols if c in df_renamed.columns]

os.makedirs(args.输出目录, exist_ok=True)

group_cols = config.get('group_cols', [[], []])
if isinstance(group_cols[0], list):
    outer_group = group_cols[0] if len(group_cols) > 0 else []
    inner_group = group_cols[1] if len(group_cols) > 1 else []
else:
    outer_group = []
    inner_group = group_cols
key_cols = config.get('key_cols', [])
if len(key_cols) == 0:
    print(f"column key error")

metric_cols = config.get('metric_cols', [
    '平均首字延时TTFT(ms)', 
    '平均生成时间TPOT(ms)', 
    '生成吞吐量(tokens/s)', 
    '总吞吐量(tokens/s)' 
    ])

normal_metrics, normalize_configs = parse_metric_cols(metric_cols, key_cols)
all_metric_cols = normal_metrics + [c['column'] for c in normalize_configs]

print(f"\n普通指标: {normal_metrics}")
print(f"归一化配置: {normalize_configs}")

dist_combinations = df_renamed.groupby(dist_cols).size().reset_index()
print(f"\n将生成 {len(dist_combinations)} 个图表...")

chart_count = 0

for idx, (_, dist_row) in enumerate(dist_combinations.iterrows()):
    df_subset = df_renamed.copy()
    for dist_col in dist_cols:
        df_subset = df_subset[df_subset[dist_col] == dist_row[dist_col]]
    
    filter_parts = []
    for dist_col in dist_cols:
        val = dist_row[dist_col]
        safe_col_name = dist_col.replace('/', '_').replace('\\', '_')[:10]
        filter_parts.append(f"{safe_col_name}_{val}")
    
    output_filename = '_'.join(filter_parts) + ".png"
    output_path = os.path.join(args.输出目录, output_filename)
    
    orig_count = len(df_subset)
    df_grouped = df_subset.groupby(key_cols).size().reset_index(name='count')
    grouped_count = len(df_grouped)
    print(f"[{idx+1}/{len(dist_combinations)}] {output_filename}: 原始{orig_count}行 -> 分组后{grouped_count}组")
    print(f"    分组详情: {df_grouped[key_cols].values.tolist()}")
    
    for c in all_metric_cols:
        if c in df_subset.columns:
            numeric_vals = pd.to_numeric(df_subset[c], errors='coerce')
            if numeric_vals.notna().any():
                sample = numeric_vals.dropna().iloc[0]
                if isinstance(sample, (int, np.integer)) and not pd.isna(numeric_vals).any():
                    df_subset[c] = numeric_vals
                else:
                    df_subset[c] = numeric_vals.fillna(0)
    
    if normalize_configs:
        group_cols = outer_group + inner_group
        df_subset = apply_normalization(df_subset, key_cols, normalize_configs, group_cols)
    
    success = generate_chart(df_subset, output_path, key_cols, outer_group, inner_group, all_metric_cols, 
                             normalize_configs=normalize_configs, merge_groups=args.合并分组)
    
    if success:
        chart_count += 1

print(f"\n完成！共生成 {chart_count} 个图表，保存到目录: {args.输出目录}")