ComputerAssignments/CA10_knowledge_graphs/scripts/run_notebook_comprehensive.py

#!/usr/bin/env python3
"""
Comprehensive script to run all CA10 notebook cells and save outputs
"""

import json
import sys
import os
from io import StringIO
import traceback

# Import all required libraries
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
from typing import List, Dict, Tuple, Optional, Union, Set
from dataclasses import dataclass
from collections import defaultdict, deque
import json
import time
import random
from abc import ABC, abstractmethod
import pickle
from itertools import combinations, product
import warnings

warnings.filterwarnings("ignore")

# Set seeds
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def capture_output(func):
    """Decorator to capture stdout"""

    def wrapper(*args, **kwargs):
        old_stdout = sys.stdout
        sys.stdout = captured_output = StringIO()
        try:
            result = func(*args, **kwargs)
            output = captured_output.getvalue()
            return result, output
        finally:
            sys.stdout = old_stdout

    return wrapper


# Load the notebook
print("📓 Loading notebook...")
with open("notebooks/CA10.ipynb", "r") as f:
    notebook = json.load(f)

print(f"   Found {len(notebook['cells'])} cells")
print(f"   Starting execution...\n")

# Track execution
execution_summary = []

# Execute code cells
for cell_idx, cell in enumerate(notebook["cells"]):
    if cell["cell_type"] == "code":
        print(f"\n{'='*70}")
        print(f"🔄 Executing Cell {cell_idx}")
        print(f"{'='*70}")

        # Get the code
        code = "".join(cell["source"])

        # Capture output
        old_stdout = sys.stdout
        sys.stdout = captured = StringIO()

        try:
            # Execute the code
            exec(code, globals())

            # Get output
            output_text = captured.getvalue()

            # Save output to notebook
            cell["outputs"] = [
                {
                    "output_type": "stream",
                    "name": "stdout",
                    "text": output_text.split("\n"),
                }
            ]

            # Update execution count
            cell["execution_count"] = cell_idx + 1

            execution_summary.append(
                {
                    "cell": cell_idx,
                    "status": "SUCCESS",
                    "output_lines": len(output_text.split("\n")),
                }
            )

            print(output_text)
            print(f"✅ Cell {cell_idx} executed successfully")

        except Exception as e:
            error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
            print(f"❌ Error in cell {cell_idx}:", file=sys.stderr)
            print(error_msg, file=sys.stderr)

            # Save error to notebook
            cell["outputs"] = [
                {
                    "output_type": "error",
                    "ename": type(e).__name__,
                    "evalue": str(e),
                    "traceback": traceback.format_exc().split("\n"),
                }
            ]

            execution_summary.append(
                {"cell": cell_idx, "status": "ERROR", "error": str(e)}
            )

        finally:
            sys.stdout = old_stdout

# Save updated notebook
print(f"\n{'='*70}")
print("💾 Saving updated notebook...")
with open("notebooks/CA10.ipynb", "w") as f:
    json.dump(notebook, f, indent=2)
print("✅ Notebook saved successfully!")

# Print summary
print(f"\n{'='*70}")
print("📊 EXECUTION SUMMARY")
print(f"{'='*70}")
for item in execution_summary:
    status_icon = "✅" if item["status"] == "SUCCESS" else "❌"
    print(f"{status_icon} Cell {item['cell']}: {item['status']}")
    if "error" in item:
        print(f"   Error: {item['error']}")

success_count = sum(1 for item in execution_summary if item["status"] == "SUCCESS")
total_count = len(execution_summary)
print(
    f"\n📈 Success Rate: {success_count}/{total_count} ({100*success_count/total_count:.1f}%)"
)
print(f"{'='*70}\n")

print("🎉 Notebook execution complete!")