ComputerAssignments/CA20_distributed_memory_systems/src/visualization.py

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from typing import Dict, List, Any, Optional
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

class MemoryAnalyzer:

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


        plt.style.use("seaborn-v0_8")
        plt.rcParams["figure.figsize"] = (14, 8)
        plt.rcParams["font.size"] = 12

    def create_memory_hierarchy_plot(self, cache_results: Dict[str, float]) -> None:
        fig, ax = plt.subplots(figsize=(10, 6))

        access_patterns = ["Sequential", "Strided", "Random"]
        hit_rates = [
            cache_results["sequential"],
            cache_results["strided"],
            cache_results["random"],
        ]
        colors = ["green", "orange", "red"]

        bars = ax.bar(
            access_patterns, [h * 100 for h in hit_rates], color=colors, alpha=0.7
        )
        ax.set_ylabel("Cache Hit Rate (%)")
        ax.set_title("Cache Performance by Access Pattern")
        ax.set_ylim(0, 100)


        for bar, hit_rate in zip(bars, hit_rates):
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2.0,
                height + 1,
                f"{hit_rate*100:.1f}%",
                ha="center",
                va="bottom",
                fontweight="bold",
            )

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "memory_hierarchy_performance.png",
            dpi=300,
            bbox_inches="tight",
        )
        plt.show()

    def create_algorithm_performance_plot(
        self, matrix_results: List[Dict[str, Any]]
    ) -> None:
        if not matrix_results:
            logger.warning("No matrix results available for plotting")
            return

        fig, ax = plt.subplots(figsize=(12, 8))


        alg_performance = {}
        for result in matrix_results:
            if result["size"] == 256:
                alg_performance[result["algorithm"]] = result["gflops"]

        if not alg_performance:
            logger.warning("No performance data for size 256")
            return

        algorithms = list(alg_performance.keys())
        gflops = list(alg_performance.values())
        colors = plt.cm.viridis(np.linspace(0, 1, len(algorithms)))

        bars = ax.bar(algorithms, gflops, color=colors, alpha=0.7)
        ax.set_ylabel("Performance (GFLOPS)")
        ax.set_title("Matrix Algorithm Performance (256x256)")
        ax.tick_params(axis="x", rotation=45)
        ax.set_yscale("log")


        for bar, gflop in zip(bars, gflops):
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2.0,
                height * 1.1,
                f"{gflop:.1f}",
                ha="center",
                va="bottom",
                fontweight="bold",
            )

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "algorithm_performance.png", dpi=300, bbox_inches="tight"
        )
        plt.show()

    def create_model_memory_plot(self, model_analysis: List[Dict[str, Any]]) -> None:
        if not model_analysis:
            logger.warning("No model analysis data available")
            return

        fig, ax = plt.subplots(figsize=(12, 8))

        model_names = [r["model"].replace(" ", "\n") for r in model_analysis]
        param_memory = [r["param_memory_mb"] for r in model_analysis]
        training_memory = [r["training_memory_mb"] for r in model_analysis]

        x = np.arange(len(model_names))
        width = 0.35

        bars1 = ax.bar(
            x - width / 2,
            param_memory,
            width,
            label="Parameters",
            alpha=0.7,
            color="skyblue",
        )
        bars2 = ax.bar(
            x + width / 2,
            training_memory,
            width,
            label="Training (est.)",
            alpha=0.7,
            color="lightcoral",
        )

        ax.set_ylabel("Memory (MB)")
        ax.set_title("Model Memory Requirements")
        ax.set_xticks(x)
        ax.set_xticklabels(model_names, fontsize=10)
        ax.legend()
        ax.set_yscale("log")


        for bars in [bars1, bars2]:
            for bar in bars:
                height = bar.get_height()
                ax.text(
                    bar.get_x() + bar.get_width() / 2.0,
                    height * 1.1,
                    f"{height:.0f}",
                    ha="center",
                    va="bottom",
                    fontsize=8,
                )

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "model_memory_requirements.png",
            dpi=300,
            bbox_inches="tight",
        )
        plt.show()

    def create_optimization_impact_plot(
        self, optimization_results: List[Dict[str, Any]]
    ) -> None:
        if not optimization_results:
            logger.warning("No optimization results available")
            return

        fig, ax = plt.subplots(figsize=(12, 8))

        opt_names = [r["optimization"] for r in optimization_results]
        peak_memories = [r["peak_memory_mb"] for r in optimization_results]

        bars = ax.barh(
            opt_names,
            peak_memories,
            color=plt.cm.plasma(np.linspace(0, 1, len(opt_names))),
        )
        ax.set_xlabel("Peak Memory (MB)")
        ax.set_title("Memory Optimization Impact")


        for i, (bar, memory) in enumerate(zip(bars, peak_memories)):
            width = bar.get_width()
            ax.text(
                width + 5,
                bar.get_y() + bar.get_height() / 2,
                f"{memory:.0f}MB",
                ha="left",
                va="center",
                fontsize=10,
            )

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "optimization_impact.png", dpi=300, bbox_inches="tight"
        )
        plt.show()

    def create_roofline_model_plot(self) -> None:
        fig, ax = plt.subplots(figsize=(10, 8))

        arithmetic_intensity = np.logspace(-1, 2, 100)
        peak_performance = 1000
        peak_bandwidth = 100


        roofline = np.minimum(peak_performance, peak_bandwidth * arithmetic_intensity)

        ax.loglog(arithmetic_intensity, roofline, "k-", linewidth=3, label="Roofline")


        operations = {
            "DAXPY": (0.5, 4),
            "SpMV": (1, 8),
            "Dense MatMul": (64, 800),
            "FFT": (2.5, 40),
        }

        for op_name, (ai, perf) in operations.items():
            ax.plot(ai, perf, "o", markersize=10, label=op_name)

        ax.set_xlabel("Arithmetic Intensity (FLOPs/Byte)")
        ax.set_ylabel("Performance (GFLOPS)")
        ax.set_title("Roofline Performance Model")
        ax.legend()
        ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "roofline_model.png", dpi=300, bbox_inches="tight"
        )
        plt.show()

    def create_memory_scaling_plot(self) -> None:
        fig, ax = plt.subplots(figsize=(10, 8))

        model_sizes = [1, 7, 70, 175]
        memory_requirements = [4, 28, 280, 700]
        optimized_memory = [2, 14, 140, 350]

        ax.loglog(
            model_sizes,
            memory_requirements,
            "ro-",
            label="Baseline",
            linewidth=2,
            markersize=8,
        )
        ax.loglog(
            model_sizes,
            optimized_memory,
            "go-",
            label="Optimized",
            linewidth=2,
            markersize=8,
        )

        ax.set_xlabel("Model Size (B parameters)")
        ax.set_ylabel("Memory Requirements (GB)")
        ax.set_title("Memory Scaling with Model Size")
        ax.legend()
        ax.grid(True, alpha=0.3)


        model_names = ["GPT-1B", "GPT-7B", "GPT-70B", "GPT-175B"]
        for i, (size, mem, name) in enumerate(
            zip(model_sizes, memory_requirements, model_names)
        ):
            ax.annotate(
                name,
                (size, mem),
                xytext=(5, 5),
                textcoords="offset points",
                fontsize=10,
            )

        plt.tight_layout()
        plt.savefig(
            self.output_dir / "memory_scaling.png", dpi=300, bbox_inches="tight"
        )
        plt.show()

def create_comprehensive_analysis(
    results: Dict[str, Any], output_dir: str = "visualizations"
) -> None:
    logger.info("Creating comprehensive analysis visualization")

    analyzer = MemoryAnalyzer(output_dir)


    if "cache_simulation" in results:
        analyzer.create_memory_hierarchy_plot(results["cache_simulation"])

    if "matrix_results" in results:
        analyzer.create_algorithm_performance_plot(results["matrix_results"])

    if "model_analysis" in results:
        analyzer.create_model_memory_plot(results["model_analysis"])

    if "optimization_results" in results:
        analyzer.create_optimization_impact_plot(results["optimization_results"])


    analyzer.create_roofline_model_plot()
    analyzer.create_memory_scaling_plot()


    create_comprehensive_dashboard(results, output_dir)

    logger.info(f"All visualizations saved to {output_dir}")

def create_comprehensive_dashboard(results: Dict[str, Any], output_dir: str) -> None:
    logger.info("Creating comprehensive dashboard")

    fig = plt.figure(figsize=(20, 16))
    gs = fig.add_gridspec(4, 4, height_ratios=[1, 1, 1, 1], width_ratios=[1, 1, 1, 1])

    fig.suptitle(
        "CA20: Advanced Memory Systems in AI - Comprehensive Analysis",
        fontsize=16,
        fontweight="bold",
        y=0.98,
    )


    ax1 = fig.add_subplot(gs[0, 0])
    if "cache_simulation" in results:
        access_patterns = ["Sequential", "Strided", "Random"]
        hit_rates = [
            results["cache_simulation"]["sequential"],
            results["cache_simulation"]["strided"],
            results["cache_simulation"]["random"],
        ]
        colors = ["green", "orange", "red"]

        bars = ax1.bar(
            access_patterns, [h * 100 for h in hit_rates], color=colors, alpha=0.7
        )
        ax1.set_ylabel("Cache Hit Rate (%)")
        ax1.set_title("Cache Performance by Access Pattern")
        ax1.set_ylim(0, 100)


    ax2 = fig.add_subplot(gs[0, 1])
    if "matrix_results" in results:
        alg_performance = {}
        for result in results["matrix_results"]:
            if result["size"] == 256:
                alg_performance[result["algorithm"]] = result["gflops"]

        if alg_performance:
            algorithms = list(alg_performance.keys())
            gflops = list(alg_performance.values())
            colors = plt.cm.viridis(np.linspace(0, 1, len(algorithms)))

            bars = ax2.bar(algorithms, gflops, color=colors, alpha=0.7)
            ax2.set_ylabel("Performance (GFLOPS)")
            ax2.set_title("Matrix Algorithm Performance")
            ax2.tick_params(axis="x", rotation=45)
            ax2.set_yscale("log")


    ax3 = fig.add_subplot(gs[0, 2])
    if "model_analysis" in results:
        model_names = [r["model"].replace(" ", "\n") for r in results["model_analysis"]]
        param_memory = [r["param_memory_mb"] for r in results["model_analysis"]]
        training_memory = [r["training_memory_mb"] for r in results["model_analysis"]]

        x = np.arange(len(model_names))
        width = 0.35

        bars1 = ax3.bar(
            x - width / 2, param_memory, width, label="Parameters", alpha=0.7
        )
        bars2 = ax3.bar(
            x + width / 2, training_memory, width, label="Training (est.)", alpha=0.7
        )

        ax3.set_ylabel("Memory (MB)")
        ax3.set_title("Model Memory Requirements")
        ax3.set_xticks(x)
        ax3.set_xticklabels(model_names, fontsize=8)
        ax3.legend()
        ax3.set_yscale("log")


    ax4 = fig.add_subplot(gs[0, 3])
    if "optimization_results" in results:
        opt_names = [r["optimization"] for r in results["optimization_results"]]
        peak_memories = [r["peak_memory_mb"] for r in results["optimization_results"]]

        bars = ax4.barh(
            opt_names,
            peak_memories,
            color=plt.cm.plasma(np.linspace(0, 1, len(opt_names))),
        )
        ax4.set_xlabel("Peak Memory (MB)")
        ax4.set_title("Memory Optimization Impact")


    ax5 = fig.add_subplot(gs[3, :2])
    ax5.axis("off")

summary_text =

    ax5.text(
        0.05,
        0.95,
        summary_text,
        transform=ax5.transAxes,
        verticalalignment="top",
        fontsize=12,
        fontfamily="monospace",
        bbox=dict(boxstyle="round", facecolor="lightblue", alpha=0.7),
    )


    ax6 = fig.add_subplot(gs[3, 2:])
    ax6.axis("off")

recommendations_text =

    ax6.text(
        0.05,
        0.95,
        recommendations_text,
        transform=ax6.transAxes,
        verticalalignment="top",
        fontsize=12,
        fontfamily="monospace",
        bbox=dict(boxstyle="round", facecolor="lightgreen", alpha=0.7),
    )

    plt.tight_layout()
    plt.savefig(
        Path(output_dir) / "comprehensive_dashboard.png", dpi=300, bbox_inches="tight"
    )
    plt.show()

def generate_final_report(
    results: Dict[str, Any], output_dir: str = "visualizations"
) -> None:
    logger.info("Generating final report")

    report_path = Path(output_dir) / "final_report.txt"

    with open(report_path, "w", encoding="utf-8") as f:
        f.write("=" * 80 + "\n")
        f.write("📊 CA20: ADVANCED MEMORY SYSTEMS IN AI - FINAL REPORT\n")
        f.write("=" * 80 + "\n")


        f.write("\n🏛️ MEMORY HIERARCHY ANALYSIS:\n")
        f.write("-" * 50 + "\n")

        if "spatial_results" in results and "row_major" in results["spatial_results"]:
            row_time = results["spatial_results"]["row_major"]["mean_time"]
            col_time = results["spatial_results"]["column_major"]["mean_time"]
            speedup = col_time / row_time
            f.write(
                f"✅ Spatial Locality Impact: {speedup:.2f}x performance difference\n"
            )
            f.write(f"   Row-major access: {row_time*1000:.2f} ms\n")
            f.write(f"   Column-major access: {col_time*1000:.2f} ms\n")

        if "cache_simulation" in results:
            seq_hit = results["cache_simulation"]["sequential"] * 100
            rand_hit = results["cache_simulation"]["random"] * 100
            f.write(f"✅ Cache Performance Analysis:\n")
            f.write(f"   Sequential access: {seq_hit:.1f}% hit rate\n")
            f.write(f"   Random access: {rand_hit:.1f}% hit rate\n")


        f.write("\n⚡ ALGORITHM PERFORMANCE ANALYSIS:\n")
        f.write("-" * 50 + "\n")

        if "matrix_results" in results:
            best_perf = max(results["matrix_results"], key=lambda x: x["gflops"])
            f.write(f"✅ Best Algorithm: {best_perf['algorithm']}\n")
            f.write(f"   Performance: {best_perf['gflops']:.1f} GFLOPS\n")
            f.write(f"   Matrix Size: {best_perf['size']}x{best_perf['size']}\n")


        f.write("\n💾 MEMORY EFFICIENCY ANALYSIS:\n")
        f.write("-" * 50 + "\n")

        if "model_analysis" in results:
            most_efficient = min(
                results["model_analysis"], key=lambda x: x["param_memory_mb"]
            )
            f.write(f"✅ Most Memory Efficient: {most_efficient['model']}\n")
            f.write(f"   Parameters: {most_efficient['parameters_M']:.2f}M\n")
            f.write(f"   Memory: {most_efficient['param_memory_mb']:.1f} MB\n")


        f.write("\n💡 KEY INSIGHTS AND RECOMMENDATIONS:\n")
        f.write("-" * 50 + "\n")
        f.write(
            "1. ✅ Memory access patterns significantly impact performance (2-4x difference)\n"
        )
        f.write(
            "2. ✅ Cache-aware algorithms provide substantial improvements for large datasets\n"
        )
        f.write(
            "3. ✅ Efficient architectures can reduce memory by 3-5x with minimal accuracy loss\n"
        )
        f.write("4. ✅ Mixed precision training offers ~50% memory reduction\n")
        f.write(
            "5. ✅ Gradient checkpointing enables training larger models on limited hardware\n"
        )
        f.write(
            "6. ✅ Memory optimization is crucial for scaling AI to larger models\n"
        )

        f.write("\n🎯 PRACTICAL RECOMMENDATIONS:\n")
        f.write("-" * 50 + "\n")
        f.write("• Always profile memory usage before optimization\n")
        f.write("• Use blocked/tiled algorithms for large matrix operations\n")
        f.write("• Enable mixed precision training when hardware supports it\n")
        f.write("• Implement gradient checkpointing for memory-constrained training\n")
        f.write("• Consider efficient architectures for deployment scenarios\n")
        f.write("• Optimize data loading and preprocessing pipelines\n")

        f.write("\n" + "=" * 80 + "\n")
        f.write("🎉 CA20 ANALYSIS COMPLETE!\n")
        f.write("=" * 80 + "\n")

    logger.info(f"Final report saved to {report_path}")