plot_audio.py

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

# --- Configuration ---
ROOT_DIR = "runs"      # The parent folder containing the logs
PLOTS_DIR = "aplots"    # Directory to save the generated plots

# Updated: Only audio datasets, removed image datasets
DATASETS = ["ljspeech", "musdb18"]

# Updated: Removed "diffusion" as per instruction
MODELS = ["diffusion", "vae"]

# Visualization Settings
COLORS = {
    "gan": "#ff7f0e",        # Orange
    "vae": "#d62728"         # Red
}

# Config defining all possible metrics to track
METRICS_CONFIG = [
    # Updated title to remove Diffusion reference
    {"key": "loss", "title": "Total Loss (GAN:Avg, VAE:ELBO)", "ylabel": "Loss/ELBO", "log_scale": True},
    {"key": "g_loss", "title": "Generator Loss", "ylabel": "G Loss", "log_scale": False},
    {"key": "d_loss", "title": "Discriminator Loss", "ylabel": "D Loss", "log_scale": False},
    {"key": "bpd", "title": "Bits Per Dimension (BPD)", "ylabel": "BPD", "log_scale": False},
    {"key": "grad_norm", "title": "Gradient Norm", "ylabel": "Norm", "log_scale": True},
    {"key": "throughput", "title": "Throughput", "ylabel": "Samples/Sec", "log_scale": False},
    {"key": "secs", "title": "Time per Epoch", "ylabel": "Seconds", "log_scale": False},
    {"key": "gpu_mem_alloc_mb", "title": "GPU Memory Allocated", "ylabel": "MB", "log_scale": False},
    {"key": "fid", "title": "FID Score", "ylabel": "FID", "log_scale": False},
    {"key": "kid_mean", "title": "KID Mean", "ylabel": "KID", "log_scale": False},
]

def parse_log_file(filepath):
    """
    Parses a single train_log.jsonl file.
    Returns a dictionary of lists for time-series data and a float for final accuracy.
    """
    data = {
        "epochs": [],
        "loss": [],
        "g_loss": [],
        "d_loss": [],
        "bpd": [],
        "secs": [],
        "gpu_mem_alloc_mb": [],
        "grad_norm": [],
        "throughput": [],
        
        # Sparse metrics (e.g., recorded every 10 epochs)
        "fid_epochs": [], 
        "fid": [],
        "kid_epochs": [],
        "kid_mean": [],
        
        "accuracy": 0.0
    }
    
    if not os.path.exists(filepath):
        # Silent return if file doesn't exist to avoid cluttering stdout
        return None

    try:
        with open(filepath, 'r') as f:
            for line in f:
                line = line.strip()
                if not line: continue
                
                entry = json.loads(line)
                
                # Check if it's a standard training epoch
                if "epoch" in entry:
                    epoch = entry["epoch"]
                    data["epochs"].append(epoch)
                    
                    # --- Metric Parsing & Unification ---
                    raw_loss = entry.get("loss")
                    g_loss = entry.get("g_loss")
                    d_loss = entry.get("d_loss")
                    elbo = entry.get("elbo") # VAE specific
                    
                    # Unified Loss Logic:
                    # 1. Use 'elbo' if present (VAE) -> map to loss
                    # 2. Use average of GAN losses if available
                    # 3. Fallback to 'loss' if generic logging is used
                    if elbo is not None:
                        unified_loss = elbo
                    elif g_loss is not None and d_loss is not None:
                        unified_loss = (g_loss + d_loss) / 2.0
                    elif raw_loss is not None:
                        unified_loss = raw_loss
                    else:
                        unified_loss = None

                    data["loss"].append(unified_loss)
                    data["g_loss"].append(g_loss)
                    data["d_loss"].append(d_loss)
                    
                    # Other standard metrics
                    data["bpd"].append(entry.get("bpd"))
                    data["secs"].append(entry.get("secs"))
                    data["gpu_mem_alloc_mb"].append(entry.get("gpu_mem_alloc_mb"))
                    data["grad_norm"].append(entry.get("grad_norm"))
                    data["throughput"].append(entry.get("throughput"))
                    
                    # Check for sparse metrics (FID)
                    if "fid" in entry and entry["fid"] is not None:
                        data["fid_epochs"].append(epoch)
                        data["fid"].append(entry["fid"])
                        
                    # Check for sparse metrics (KID)
                    if "kid_mean" in entry and entry["kid_mean"] is not None:
                        data["kid_epochs"].append(epoch)
                        data["kid_mean"].append(entry["kid_mean"])
                
                # Check if it's the final summary line
                elif "gen_train_real_test_top1" in entry:
                    data["accuracy"] = entry["gen_train_real_test_top1"]
                    
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON in {filepath}: {e}")
        return None
        
    return data

def plot_single_metric(dataset_name, all_data, metric_config):
    """
    Creates and saves a single plot for a specific metric for a specific dataset,
    comparing all available models.
    """
    key = metric_config["key"]
    title = metric_config["title"]
    ylabel = metric_config["ylabel"]
    use_log = metric_config.get("log_scale", False)
    
    # Check if we have any data to plot for this metric across ANY model
    has_data = False
    for model in MODELS:
        model_data = all_data[model].get(dataset_name)
        if model_data:
            values = model_data.get(key, [])
            if values and any(v is not None for v in values):
                has_data = True
                break
    
    if not has_data:
        return # Skip generating empty plots

    plt.figure(figsize=(10, 6))
    
    for model in MODELS:
        # Access data: all_data[model][dataset]
        if dataset_name not in all_data[model] or all_data[model][dataset_name] is None:
            continue
            
        data = all_data[model][dataset_name]
        y = data.get(key, [])
        
        # Determine X axis
        if key == "fid":
            x = data["fid_epochs"]
            marker = 'o'
        elif key == "kid_mean":
            x = data["kid_epochs"]
            marker = 'o'
        else:
            x = data["epochs"]
            marker = None
        
        # Filter out None values just in case
        if x and y and len(x) == len(y):
            # Clean data (remove None points for plotting)
            clean_x, clean_y = [], []
            for cur_x, cur_y in zip(x, y):
                if cur_y is not None:
                    clean_x.append(cur_x)
                    clean_y.append(cur_y)
                    
            if clean_x:
                plt.plot(clean_x, clean_y, label=model, color=COLORS.get(model, 'black'), 
                         marker=marker, markersize=6, linewidth=2.5)

    if use_log:
        plt.yscale('log')
        title += " (Log Scale)"
        
        # Set specific Y-axis limits for Loss if needed
        if key == "loss":
             plt.ylim(1e-2, 1e2) 
             pass

    # Increased font sizes
    plt.title(f"{dataset_name.upper()}: {title}", fontsize=20, fontweight='bold')
    plt.xlabel("Epoch", fontsize=18)
    plt.ylabel(ylabel, fontsize=18)
    plt.xticks(fontsize=15)
    plt.yticks(fontsize=15)
    
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(fontsize=12)
    plt.tight_layout()
    
    output_filename = os.path.join(PLOTS_DIR, f"{dataset_name}_{key}.png")
    plt.savefig(output_filename)
    print(f"Saved plot: {output_filename}")
    plt.close()

def plot_accuracy_comparison(all_data):
    """
    Plots a grouped bar chart for accuracy.
    X-axis: Datasets
    Groups: Models
    """
    x = np.arange(len(DATASETS))
    width = 0.35  # width of the bars
    
    fig, ax = plt.subplots(figsize=(12, 6))
    
    for i, model in enumerate(MODELS):
        accuracies = []
        for dataset in DATASETS:
            # Safe retrieval
            if model in all_data and dataset in all_data[model] and all_data[model][dataset]:
                accuracies.append(all_data[model][dataset]["accuracy"])
            else:
                accuracies.append(0)
        
        # Calculate offset for grouped bars
        offset = x + (i * width) - (width * (len(MODELS) - 1) / 2)
        
        rects = ax.bar(offset, accuracies, width, label=model, color=COLORS.get(model, 'gray'))
        
        # Add value labels on top of bars
        ax.bar_label(rects, padding=3, fmt='%.4f', fontsize=9)

    # Increased font sizes
    ax.set_ylabel('Accuracy (Top 1)', fontsize=14)
    ax.set_title('Final Model Accuracy by Dataset and Model', fontsize=16, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels([d.upper() for d in DATASETS], fontsize=11, rotation=15)
    ax.tick_params(axis='y', labelsize=12)
    ax.legend(fontsize=12)
    ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    plt.tight_layout()
    output_filename = os.path.join(PLOTS_DIR, "accuracy_comparison.png")
    plt.savefig(output_filename)
    print(f"Saved plot: {output_filename}")
    plt.close()

def main():
    # 0. Setup
    os.makedirs(PLOTS_DIR, exist_ok=True)
    
    # 1. Load Data
    all_data = {model: {} for model in MODELS}
    
    print("Parsing log files...")
    for model in MODELS:
        for dataset in DATASETS:
            filepath = os.path.join(ROOT_DIR, dataset, model, "train_log.jsonl")
            parsed = parse_log_file(filepath)
            all_data[model][dataset] = parsed

    # 2. Generate Individual Plots (Grouped by Dataset, Comparing Models)
    print("\nGenerating individual plots...")
    for dataset in DATASETS:
        for metric_config in METRICS_CONFIG:
            plot_single_metric(dataset, all_data, metric_config)

    # 3. Generate Accuracy Plot
    print("\nGenerating accuracy comparison...")
    plot_accuracy_comparison(all_data)
    
    print(f"\nDone! Check the '{PLOTS_DIR}' directory.")

if __name__ == "__main__":
    main()