visualize_results.py

import json
import matplotlib.pyplot as plt
import numpy as np
import os

def load_all_results():
    """Load results from all experiments"""
    all_results = {}
    
    # Load baseline results
    try:
        with open('results/baseline_results.json', 'r') as f:
            baseline_results = json.load(f)
            baseline_success = sum(1 for r in baseline_results if r.get('execution_success', False))
            baseline_total = len(baseline_results)
            baseline_rate = baseline_success / baseline_total * 100
            all_results['Baseline'] = {
                'success': baseline_success,
                'total': baseline_total,
                'rate': baseline_rate
            }
    except (FileNotFoundError, json.JSONDecodeError):
        print("No baseline results found")
    
    # Load self-eval results
    try:
        with open('results/self_eval_results.json', 'r') as f:
            self_eval_results = json.load(f)
            self_eval_success = 0
            for r in self_eval_results:
                if r.get('improved_success', False) or r.get('execution_success', False):
                    self_eval_success += 1
            self_eval_total = len(self_eval_results)
            self_eval_rate = self_eval_success / self_eval_total * 100
            all_results['Self-Eval'] = {
                'success': self_eval_success,
                'total': self_eval_total,
                'rate': self_eval_rate
            }
    except (FileNotFoundError, json.JSONDecodeError):
        print("No self-evaluation results found")
    
    # Load post-trained results
    percentages = [0.1, 1, 5, 10, 30, 50, 100] 
    for percentage in percentages:
        result_path = f'results/post_trained_{percentage}pct/humaneval_results.json'
        try:
            with open(result_path, 'r') as f:
                post_trained_results = json.load(f)
                success = sum(1 for r in post_trained_results if r.get('execution_success', False))
                total = len(post_trained_results)
                rate = success / total * 100
                all_results[f'Post-Trained {percentage}%'] = {
                    'success': success,
                    'total': total,
                    'rate': rate
                }
        except (FileNotFoundError, json.JSONDecodeError):
            print(f"No post-trained {percentage}% results found")
    
    # Load combined approach results
    for percentage in percentages:
        result_path = f'results/combined_{percentage}pct_results.json'
        try:
            with open(result_path, 'r') as f:
                combined_results = json.load(f)
                initial_success = sum(1 for r in combined_results if r.get('initial_success', False))
                final_success = initial_success
                final_success += sum(1 for r in combined_results 
                                   if not r.get('initial_success', False) and r.get('final_success', False))
                total = len(combined_results)
                rate = final_success / total * 100
                all_results[f'Combined {percentage}%'] = {
                    'success': final_success,
                    'total': total,
                    'rate': rate
                }
        except (FileNotFoundError, json.JSONDecodeError):
            print(f"No combined {percentage}% results found")
    
    return all_results

def create_visualizations():
    """Generate a comprehensive set of visualizations"""
    results = load_all_results()
    
    # Save the consolidated results
    with open('results/consolidated_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    
    # Create main comparison chart
    create_comparison_chart(results)
    
    # Create post-training progression chart
    create_progression_chart(results)

def create_comparison_chart(results):
    """Create a bar chart comparing all approaches"""
    # Prepare data
    labels = list(results.keys())
    success_rates = [results[key]['rate'] for key in labels]
    
    # Set up custom colors
    color_map = {
        'Baseline': 'gray',
        'Self-Eval': 'lightblue'
    }
    
    # Colors for post-trained models (red gradient)
    post_trained_keys = [k for k in labels if 'Post-Trained' in k]
    post_trained_count = len(post_trained_keys)
    if post_trained_count > 0:
        for i, key in enumerate(post_trained_keys):
            color_map[key] = plt.cm.Reds(0.4 + (0.6 * i / max(post_trained_count-1, 1)))
    
    # Colors for combined models (blue gradient)
    combined_keys = [k for k in labels if 'Combined' in k]
    combined_count = len(combined_keys)
    if combined_count > 0:
        for i, key in enumerate(combined_keys):
            color_map[key] = plt.cm.Blues(0.4 + (0.6 * i / max(combined_count-1, 1)))
    
    # Get colors for each bar
    colors = [color_map.get(label, 'lightgray') for label in labels]
    
    # Create the plot
    plt.figure(figsize=(12, 8))
    bars = plt.bar(labels, success_rates, color=colors)
    
    # Add value labels on bars
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.5,
                 f'{height:.1f}%', ha='center', va='bottom', fontsize=9)
    
    # Formatting
    plt.xlabel('Approach', fontsize=12)
    plt.ylabel('Success Rate (%)', fontsize=12)
    plt.title('HumanEval Success Rates by Approach', fontsize=14, fontweight='bold')
    plt.xticks(rotation=45, ha='right')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    
    # Save figure
    os.makedirs('results/figures', exist_ok=True)
    plt.savefig('results/figures/all_approaches_comparison.png', dpi=300)
    plt.savefig('results/figures/all_approaches_comparison.pdf')

def create_progression_chart(results):
    """Create line chart showing progression with increasing data percentages"""
    # Extract data for progression
    # Add 0.1% to the list of percentages
    percentages = [0.1, 1, 5, 10, 30, 50, 100]
    post_trained_rates = []
    combined_rates = []
    
    for p in percentages:
        if f'Post-Trained {p}%' in results:
            post_trained_rates.append(results[f'Post-Trained {p}%']['rate'])
        else:
            post_trained_rates.append(None)
            
        if f'Combined {p}%' in results:
            combined_rates.append(results[f'Combined {p}%']['rate'])
        else:
            combined_rates.append(None)
    
    # Filter out None values while preserving corresponding percentages
    valid_post = [(p, r) for p, r in zip(percentages, post_trained_rates) if r is not None]
    valid_combined = [(p, r) for p, r in zip(percentages, combined_rates) if r is not None]
    
    # Create plot
    plt.figure(figsize=(10, 6))
    
    # Plot baseline as horizontal line if available
    if 'Baseline' in results:
        plt.axhline(y=results['Baseline']['rate'], color='gray', linestyle='--', 
                   label=f"Baseline ({results['Baseline']['rate']:.1f}%)")
    
    # Plot self-eval as horizontal line if available
    if 'Self-Eval' in results:
        plt.axhline(y=results['Self-Eval']['rate'], color='lightblue', linestyle='--', 
                   label=f"Self-Eval ({results['Self-Eval']['rate']:.1f}%)")
    
    # Plot post-trained and combined lines
    if valid_post:
        post_p, post_r = zip(*valid_post)
        plt.plot(post_p, post_r, 'o-', color='red', linewidth=2, 
                label='Post-Trained', marker='o', markersize=8)
        
    if valid_combined:
        comb_p, comb_r = zip(*valid_combined)
        plt.plot(comb_p, comb_r, 's-', color='blue', linewidth=2, 
                label='Combined Approach', marker='s', markersize=8)
    
    # Formatting
    plt.xlabel('Percentage of Data Used (%)', fontsize=12)
    plt.ylabel('Success Rate (%)', fontsize=12)
    plt.title('Model Performance vs. Training Data Size', fontsize=14, fontweight='bold')
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # X-axis formatting
    plt.xticks(percentages)
    plt.xlim(0.5, 105)  # Set limits to make the plot look better
    
    # Add legend
    plt.legend(loc='best', fontsize=10)
    
    # Save figure
    plt.tight_layout()
    plt.savefig('results/figures/performance_progression.png', dpi=300)
    plt.savefig('results/figures/performance_progression.pdf')

if __name__ == "__main__":
    create_visualizations()