ComputerAssignments/CA8_neural_program_synthesis/scripts/visualizations.py

#!/usr/bin/env python3

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import json
import os
from typing import Dict, List, Tuple, Optional
from pathlib import Path
import torch


plt.style.use("seaborn-v0_8")
sns.set_palette("husl")

class SynthesisVisualizer:

    def __init__(self, output_dir: str = "./visualizations"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)

    def plot_training_curves(
        self,
        train_losses: List[float],
        val_losses: List[float],
        train_accs: Optional[List[float]] = None,
        val_accs: Optional[List[float]] = None,
        save_name: str = "training_curves.png",
    ):
        fig, axes = plt.subplots(1, 2, figsize=(15, 5))


        axes[0].plot(train_losses, label="Training Loss", linewidth=2)
        axes[0].plot(val_losses, label="Validation Loss", linewidth=2)
        axes[0].set_xlabel("Epoch")
        axes[0].set_ylabel("Loss")
        axes[0].set_title("Training and Validation Loss")
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)


        if train_accs and val_accs:
            axes[1].plot(train_accs, label="Training Accuracy", linewidth=2)
            axes[1].plot(val_accs, label="Validation Accuracy", linewidth=2)
            axes[1].set_xlabel("Epoch")
            axes[1].set_ylabel("Accuracy")
            axes[1].set_title("Training and Validation Accuracy")
            axes[1].legend()
            axes[1].grid(True, alpha=0.3)
        else:
            axes[1].text(
                0.5,
                0.5,
                "Accuracy data not available",
                ha="center",
                va="center",
                transform=axes[1].transAxes,
            )
            axes[1].set_title("Accuracy Curves")

        plt.tight_layout()
        plt.savefig(self.output_dir / save_name, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"Training curves saved to {self.output_dir / save_name}")

    def plot_evaluation_metrics(
        self, metrics: Dict[str, float], save_name: str = "evaluation_metrics.png"
    ):
        fig, ax = plt.subplots(figsize=(12, 6))

        metric_names = list(metrics.keys())
        metric_values = list(metrics.values())

        bars = ax.bar(metric_names, metric_values, alpha=0.7)


        for bar, value in zip(bars, metric_values):
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2.0,
                height + 0.01,
                f"{value:.3f}",
                ha="center",
                va="bottom",
            )

        ax.set_ylabel("Score")
        ax.set_title("Model Evaluation Metrics")
        ax.set_ylim(0, max(metric_values) * 1.1)


        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        plt.savefig(self.output_dir / save_name, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"Evaluation metrics saved to {self.output_dir / save_name}")

    def plot_success_rates(
        self,
        task_results: Dict[str, Dict[str, float]],
        save_name: str = "success_rates.png",
    ):
        fig, ax = plt.subplots(figsize=(12, 8))

        tasks = list(task_results.keys())
        metrics = [
            "syntax_valid_rate",
            "functionally_correct_rate",
            "overall_success_rate",
        ]

        x = np.arange(len(tasks))
        width = 0.25

        for i, metric in enumerate(metrics):
            values = [task_results[task].get(metric, 0) for task in tasks]
            ax.bar(x + i * width, values, width, label=metric.replace("_", " ").title())

        ax.set_xlabel("Task Type")
        ax.set_ylabel("Success Rate")
        ax.set_title("Success Rates by Task Type")
        ax.set_xticks(x + width)
        ax.set_xticklabels(tasks, rotation=45, ha="right")
        ax.legend()
        ax.set_ylim(0, 1)

        plt.tight_layout()
        plt.savefig(self.output_dir / save_name, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"Success rates saved to {self.output_dir / save_name}")

    def plot_model_comparison(
        self,
        model_results: Dict[str, Dict[str, float]],
        save_name: str = "model_comparison.png",
    ):
        fig, ax = plt.subplots(figsize=(12, 6))

        models = list(model_results.keys())
        metrics = [
            "overall_success_rate",
            "syntax_valid_rate",
            "functionally_correct_rate",
        ]

        x = np.arange(len(models))
        width = 0.25

        for i, metric in enumerate(metrics):
            values = [model_results[model].get(metric, 0) for model in models]
            ax.bar(x + i * width, values, width, label=metric.replace("_", " ").title())

        ax.set_xlabel("Model Architecture")
        ax.set_ylabel("Performance")
        ax.set_title("Model Architecture Comparison")
        ax.set_xticks(x + width)
        ax.set_xticklabels(models)
        ax.legend()
        ax.set_ylim(0, 1)

        plt.tight_layout()
        plt.savefig(self.output_dir / save_name, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"Model comparison saved to {self.output_dir / save_name}")

    def plot_generated_programs_analysis(
        self, program_data: List[Dict], save_name: str = "program_analysis.png"
    ):
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))


        lengths = [len(prog.get("program", "").split()) for prog in program_data]
        axes[0, 0].hist(lengths, bins=20, alpha=0.7, edgecolor="black")
        axes[0, 0].set_xlabel("Program Length (tokens)")
        axes[0, 0].set_ylabel("Frequency")
        axes[0, 0].set_title("Generated Program Length Distribution")


        success_by_length = {}
        for prog in program_data:
            length = len(prog.get("program", "").split())
            success = prog.get("functionally_correct", False)
            if length not in success_by_length:
                success_by_length[length] = {"total": 0, "success": 0}
            success_by_length[length]["total"] += 1
            if success:
                success_by_length[length]["success"] += 1

        lengths_sorted = sorted(success_by_length.keys())
        success_rates = [
            success_by_length[l]["success"] / success_by_length[l]["total"]
            for l in lengths_sorted
        ]

        axes[0, 1].plot(lengths_sorted, success_rates, "o-", linewidth=2)
        axes[0, 1].set_xlabel("Program Length (tokens)")
        axes[0, 1].set_ylabel("Success Rate")
        axes[0, 1].set_title("Success Rate vs Program Length")
        axes[0, 1].grid(True, alpha=0.3)


        syntax_correct = [prog.get("syntax_valid", False) for prog in program_data]
        functional_correct = [
            prog.get("functionally_correct", False) for prog in program_data
        ]

        confusion_matrix = np.array(
            [
                [
                    sum(
                        not s and not f
                        for s, f in zip(syntax_correct, functional_correct)
                    ),
                    sum(
                        not s and f for s, f in zip(syntax_correct, functional_correct)
                    ),
                ],
                [
                    sum(
                        s and not f for s, f in zip(syntax_correct, functional_correct)
                    ),
                    sum(s and f for s, f in zip(syntax_correct, functional_correct)),
                ],
            ]
        )

        sns.heatmap(
            confusion_matrix,
            annot=True,
            fmt="d",
            xticklabels=["Functionally Incorrect", "Functionally Correct"],
            yticklabels=["Syntax Invalid", "Syntax Valid"],
            ax=axes[1, 0],
        )
        axes[1, 0].set_title("Syntax vs Functional Correctness")


        task_performance = {}
        for prog in program_data:
            task = prog.get("task_type", "unknown")
            success = prog.get("functionally_correct", False)
            if task not in task_performance:
                task_performance[task] = {"total": 0, "success": 0}
            task_performance[task]["total"] += 1
            if success:
                task_performance[task]["success"] += 1

        tasks = list(task_performance.keys())
        success_rates = [
            task_performance[t]["success"] / task_performance[t]["total"] for t in tasks
        ]

        axes[1, 1].bar(tasks, success_rates, alpha=0.7)
        axes[1, 1].set_xlabel("Task Type")
        axes[1, 1].set_ylabel("Success Rate")
        axes[1, 1].set_title("Performance by Task Type")
        axes[1, 1].set_ylim(0, 1)

        plt.tight_layout()
        plt.savefig(self.output_dir / save_name, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"Program analysis saved to {self.output_dir / save_name}")

    def create_summary_report(
        self, results: Dict, save_name: str = "summary_report.html"
    ):
        """Create an HTML summary report"""
        html_content = """
        <!DOCTYPE html>
        <html>
        <head>
            <title>Synthesis Results Summary</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                table { border-collapse: collapse; width: 100%; }
                th, td { border: 1px solid #ddd; padding: 12px; text-align: left; }
                th { background-color: #4CAF50; color: white; }
                tr:nth-child(even) { background-color: #f2f2f2; }
            </style>
        </head>
        <body>
            <h1>Neural Program Synthesis Results</h1>
            <table>
                <tr><th>Metric</th><th>Value</th></tr>
        """

        for metric, value in results.items():
            if isinstance(value, (int, float)):
                html_content += f"<tr><td>{metric.replace('_', ' ').title()}</td><td>{value:.3f}</td></tr>\n"

        html_content += """
            </table>
        </body>
        </html>
        """

        with open(self.output_dir / save_name, "w", encoding="utf-8") as f:
            f.write(html_content)

        print(f"Summary report saved to {self.output_dir / save_name}")

    def save_results_json(self, results: Dict, save_name: str = "results.json"):
        with open(self.output_dir / save_name, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        print(f"Results saved to {self.output_dir / save_name}")

def create_visualizations_from_results(results_dir: str = "./results"):
    visualizer = SynthesisVisualizer()


    results_files = list(Path(results_dir).glob("*.json"))
    if not results_files:
        print(f"No results files found in {results_dir}")
        return

    all_results = {}
    for file_path in results_files:
        with open(file_path, "r") as f:
            results = json.load(f)
            all_results[file_path.stem] = results


    if "training_results" in all_results:
        training_results = all_results["training_results"]
        if "train_losses" in training_results and "val_losses" in training_results:
            visualizer.plot_training_curves(
                training_results["train_losses"],
                training_results["val_losses"],
                training_results.get("train_accs"),
                training_results.get("val_accs"),
            )

    if "evaluation_results" in all_results:
        eval_results = all_results["evaluation_results"]
        visualizer.plot_evaluation_metrics(eval_results)


    combined_results = {}
    for results in all_results.values():
        combined_results.update(results)

    visualizer.create_summary_report(combined_results)
    visualizer.save_results_json(combined_results)

    print("All visualizations created successfully!")

if __name__ == "__main__":

    visualizer = SynthesisVisualizer()


    train_losses = [2.5, 2.1, 1.8, 1.5, 1.2, 0.9, 0.7, 0.6, 0.5, 0.4]
    val_losses = [2.6, 2.2, 1.9, 1.7, 1.4, 1.1, 0.9, 0.8, 0.7, 0.6]

    visualizer.plot_training_curves(train_losses, val_losses)


    metrics = {
        "syntax_valid_rate": 0.85,
        "functionally_correct_rate": 0.72,
        "overall_success_rate": 0.68,
        "execution_success_rate": 0.78,
    }

    visualizer.plot_evaluation_metrics(metrics)

    print("Example visualizations created!")