{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "95cd5e56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 50, Loss: -0.2426\n",
      "Epoch 100, Loss: -0.3236\n",
      "Epoch 150, Loss: -0.3399\n",
      "Epoch 200, Loss: -0.3477\n",
      "Epoch 250, Loss: -0.3536\n",
      "Epoch 300, Loss: -0.3594\n",
      "RMSE (noisy): 0.1975\n",
      "RMSE (nash-fused): 0.0810\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 77416 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from torch import nn\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torchvision import datasets, transforms\n",
    "from skimage.restoration import denoise_tv_chambolle, denoise_nl_means, estimate_sigma\n",
    "from scipy.ndimage import gaussian_filter, median_filter\n",
    "import cvxpy as cp\n",
    "\n",
    "def build_graph_from_mnist_image(img, noise_level=0.3):\n",
    "    n = img.shape[0]\n",
    "    beta = img.flatten()\n",
    "    beta_hat = beta + noise_level * torch.randn_like(beta)\n",
    "    se = noise_level * torch.ones_like(beta)\n",
    "    grid_y, grid_x = torch.meshgrid(torch.arange(n), torch.arange(n), indexing='ij')\n",
    "    node_feat = torch.stack([beta , grid_x.flatten() / n, grid_y.flatten() / n], dim=1)\n",
    "\n",
    "    edge_index = []\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            idx = i * n + j\n",
    "            if i > 0: edge_index.append((idx, idx - n))\n",
    "            if i < n - 1: edge_index.append((idx, idx + n))\n",
    "            if j > 0: edge_index.append((idx, idx - 1))\n",
    "            if j < n - 1: edge_index.append((idx, idx + 1))\n",
    "    edge_index = torch.tensor(edge_index).t().long()\n",
    "    edge_attr = torch.ones(edge_index.shape[1])\n",
    "    return beta, beta_hat, se, node_feat, edge_index, edge_attr\n",
    "\n",
    "class ManualStyleGNN(MessagePassing):\n",
    "    def __init__(self, node_feat_dim, hidden_dim):\n",
    "        super().__init__(aggr='add')\n",
    "        self.fc1 = nn.Linear(node_feat_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc_out = nn.Linear(hidden_dim, 4)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        h = F.relu(self.fc1(x))\n",
    "        h = F.relu(self.fc2(h))\n",
    "        return self.propagate(edge_index, x=h)\n",
    "\n",
    "    def message(self, x_j):\n",
    "        return x_j\n",
    "\n",
    "    def update(self, aggr_out, x):\n",
    "        h = F.relu(x + aggr_out)\n",
    "        return self.fc_out(h)\n",
    "\n",
    "class EB_GNN_Denoiser(nn.Module):\n",
    "    def __init__(self, gnn_model):\n",
    "        super().__init__()\n",
    "        self.gnn = gnn_model\n",
    "\n",
    "    def marginal_log_likelihood(self, beta_hat, se, X, edge_index):\n",
    "        output = self.gnn(X, edge_index)\n",
    "        mu, logit_pi, log_lambda1, log_lambda2 = output[:, 0], output[:, 1], output[:, 2], output[:, 3]\n",
    "        pi = torch.sigmoid(logit_pi)\n",
    "        lambda1 = torch.exp(log_lambda1).clamp(min=1e-4)\n",
    "        lambda2 = torch.exp(log_lambda2).clamp(min=1e-4)\n",
    "        lambda_sum = lambda1 + lambda2\n",
    "\n",
    "        log_lap_prod = -torch.abs(beta_hat - mu) / lambda_sum - torch.log(2 * lambda_sum)\n",
    "        log_pointmass = torch.zeros_like(beta_hat)\n",
    "        prior_log_density = torch.logaddexp(torch.log(1 - pi + 1e-8) + log_pointmass,\n",
    "                                            torch.log(pi + 1e-8) + log_lap_prod)\n",
    "\n",
    "        log_lik = -0.5 * ((beta_hat - mu) / se)**2 - torch.log(se) - 0.5 * np.log(2 * np.pi)\n",
    "        return -torch.mean(prior_log_density + log_lik)\n",
    "\n",
    "    def posterior_mean(self, beta_hat, se, X, edge_index):\n",
    "        output = self.gnn(X, edge_index)\n",
    "        mu, logit_pi, log_lambda1, log_lambda2 = output[:, 0], output[:, 1], output[:, 2], output[:, 3]\n",
    "        pi = torch.sigmoid(logit_pi)\n",
    "        lambda1 = torch.exp(log_lambda1).clamp(min=1e-4)\n",
    "        lambda2 = torch.exp(log_lambda2).clamp(min=1e-4)\n",
    "        lambda_sum = lambda1 + lambda2\n",
    "\n",
    "        ll_pointmass = torch.exp(-0.5 * (beta_hat / se)**2) / (se * np.sqrt(2 * np.pi))\n",
    "        ll_lap = torch.exp(-torch.abs(beta_hat - mu) / lambda_sum) / (2 * lambda_sum)\n",
    "        w0 = (1 - pi) * ll_pointmass\n",
    "        w1 = pi * ll_lap\n",
    "        norm = w0 + w1 + 1e-8\n",
    "        post_weight = w1 / norm\n",
    "        return post_weight * mu, post_weight, pi, lambda1, lambda2\n",
    "\n",
    "def benchmark_baselines(beta_np, beta_hat_np, shape=(28, 28)):\n",
    "    noisy_img = beta_hat_np.reshape(shape)\n",
    "    tv = denoise_tv_chambolle(noisy_img, weight=0.1).reshape(-1)\n",
    "    rmse_tv = np.sqrt(np.mean((tv - beta_np)**2))\n",
    "\n",
    "    Z = cp.Variable(shape)\n",
    "    obj = cp.Minimize(0.5 * cp.sum_squares(Z - noisy_img) + 0.1 * cp.tv(Z))\n",
    "    cp.Problem(obj).solve(solver=cp.SCS)\n",
    "    fused = Z.value.reshape(-1)\n",
    "    rmse_fused = np.sqrt(np.mean((fused - beta_np)**2))\n",
    "\n",
    "    gauss = gaussian_filter(noisy_img, sigma=1).reshape(-1)\n",
    "    rmse_gauss = np.sqrt(np.mean((gauss - beta_np)**2))\n",
    "\n",
    "    median = median_filter(noisy_img, size=3).reshape(-1)\n",
    "    rmse_median = np.sqrt(np.mean((median - beta_np)**2))\n",
    "\n",
    "    sigma_est = np.mean(estimate_sigma(noisy_img, channel_axis=None))\n",
    "    nlm = denoise_nl_means(noisy_img, h=1.15 * sigma_est, fast_mode=True,\n",
    "                           patch_size=3, patch_distance=5, channel_axis=None).reshape(-1)\n",
    "    rmse_nlm = np.sqrt(np.mean((nlm - beta_np)**2))\n",
    "\n",
    "    return {\n",
    "        \"TV\": (tv, rmse_tv),\n",
    "        \"Fused Lasso\": (fused, rmse_fused),\n",
    "        \"Gaussian\": (gauss, rmse_gauss),\n",
    "        \"Median\": (median, rmse_median),\n",
    "        \"NLM\": (nlm, rmse_nlm)\n",
    "    }\n",
    "\n",
    "def run_pipeline():\n",
    "    transform = transforms.Compose([transforms.ToTensor()])\n",
    "    mnist = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "    img = mnist.data[np.random.randint(1, 10000)].float() / 255.0\n",
    "    img = img[:28, :28]\n",
    "\n",
    "    beta, beta_hat, se, node_feat, edge_index, _ = build_graph_from_mnist_image(img, noise_level=0.2)\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    beta_hat, se, node_feat, edge_index = beta_hat.to(device), se.to(device), node_feat.to(device), edge_index.to(device)\n",
    "\n",
    "    model = EB_GNN_Denoiser(ManualStyleGNN(node_feat_dim=3, hidden_dim=64).to(device))\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "    for epoch in range(300):\n",
    "        optimizer.zero_grad()\n",
    "        loss = model.marginal_log_likelihood(beta_hat, se, node_feat, edge_index)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if (epoch + 1) % 50 == 0:\n",
    "            print(f\"Epoch {epoch + 1}, Loss: {loss.item():.4f}\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        beta_denoised, pi_post, pi_prior, lambda1, lambda2 = model.posterior_mean(beta_hat, se, node_feat, edge_index)\n",
    "        beta_np = beta.cpu().numpy()\n",
    "        beta_hat_np = beta_hat.cpu().numpy()\n",
    "        beta_denoised_np = beta_denoised.cpu().numpy().reshape(28, 28)\n",
    "\n",
    "    rmse_noisy = np.sqrt(np.mean((beta_hat_np - beta_np) ** 2))\n",
    "    rmse_denoised = np.sqrt(np.mean((beta_denoised_np.flatten() - beta_np) ** 2))\n",
    "    print(f\"RMSE (noisy): {rmse_noisy:.4f}\")\n",
    "    print(f\"RMSE (nash-fused): {rmse_denoised:.4f}\")\n",
    "\n",
    "    baselines = benchmark_baselines(beta_np, beta_hat_np)\n",
    "\n",
    "    # --- Plot: Noisy, Original, Nash-fused, and baselines ---\n",
    "    fig, axes = plt.subplots(2, 4, figsize=(18, 10))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    font_title = {\"fontsize\": 15}\n",
    "    font_ticks = {\"fontsize\": 12}\n",
    "\n",
    "    # Noisy\n",
    "    axes[0].imshow(beta_hat_np.reshape(28, 28), cmap=\"viridis\")\n",
    "    axes[0].set_title(f\"Noisy\\nRMSE={rmse_noisy:.4f}\", **font_title)\n",
    "\n",
    "    # Original\n",
    "    axes[1].imshow(beta_np.reshape(28, 28), cmap=\"viridis\")\n",
    "    axes[1].set_title(\"Original\", **font_title)\n",
    "\n",
    "    # Nash-fused\n",
    "    axes[2].imshow(beta_denoised_np, cmap=\"viridis\")\n",
    "    axes[2].set_title(f\"NASH-Fused\\nRMSE={rmse_denoised:.4f}\", **font_title)\n",
    "\n",
    "    # Baselines\n",
    "    for i, (name, (img, rmse)) in enumerate(baselines.items(), start=3):\n",
    "        axes[i].imshow(img.reshape(28, 28), cmap=\"viridis\")\n",
    "        axes[i].set_title(f\"{name}\\nRMSE={rmse:.4f}\", **font_title)\n",
    "\n",
    "    for ax in axes:\n",
    "        ax.axis(\"off\")\n",
    "        for label in (ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "            label.set_fontsize(font_ticks[\"fontsize\"])\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Run the pipeline\n",
    "run_pipeline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3c6f38a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add at the beginning:\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Noise2Self CNN\n",
    "class DenoiserCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 1, 3, padding=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "def noise2self_loss(pred, target, mask):\n",
    "    return F.mse_loss(pred[mask], target[mask])\n",
    "\n",
    "def train_noise2self_model(model, loader, device, noise_std=0.2, epochs=3):\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        total_loss = 0\n",
    "        for x, _ in loader:\n",
    "            x = x.to(device)\n",
    "            noisy = x + noise_std * torch.randn_like(x)\n",
    "            mask = (torch.rand_like(x) > 0.9)\n",
    "            masked_noisy = noisy.clone()\n",
    "            masked_noisy[mask] = 0.0\n",
    "            output = model(masked_noisy)\n",
    "            loss = noise2self_loss(output, noisy, mask)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += loss.item()\n",
    "        print(f\"[Epoch {epoch + 1}] Noise2Self Loss: {total_loss / len(loader):.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "16b3ef07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_pipeline():\n",
    "    transform = transforms.Compose([transforms.ToTensor()])\n",
    "    mnist = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "    img = mnist.data[np.random.randint(1, 10000)].float() / 255.0\n",
    "    img = img[:28, :28]\n",
    "\n",
    "    beta, beta_hat, se, node_feat, edge_index, _ = build_graph_from_mnist_image(img, noise_level=0.2)\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    beta_hat, se, node_feat, edge_index = beta_hat.to(device), se.to(device), node_feat.to(device), edge_index.to(device)\n",
    "\n",
    "    # --- EB-GNN Training ---\n",
    "    model = EB_GNN_Denoiser(ManualStyleGNN(node_feat_dim=3, hidden_dim=64).to(device))\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "    for epoch in range(800):\n",
    "        optimizer.zero_grad()\n",
    "        loss = model.marginal_log_likelihood(beta_hat, se, node_feat, edge_index)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if (epoch + 1) % 50 == 0:\n",
    "            print(f\"Epoch {epoch + 1}, Loss: {loss.item():.4f}\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        beta_denoised, *_ = model.posterior_mean(beta_hat, se, node_feat, edge_index)\n",
    "        beta_np = beta.cpu().numpy()\n",
    "        beta_hat_np = beta_hat.cpu().numpy()\n",
    "        beta_denoised_np = beta_denoised.cpu().numpy().reshape(28, 28)\n",
    "\n",
    "    rmse_noisy = np.sqrt(np.mean((beta_hat_np - beta_np) ** 2))\n",
    "    rmse_denoised = np.sqrt(np.mean((beta_denoised_np.flatten() - beta_np) ** 2))\n",
    "    print(f\"RMSE (noisy): {rmse_noisy:.4f}\")\n",
    "    print(f\"RMSE (nash-fused): {rmse_denoised:.4f}\")\n",
    "\n",
    "    # --- Train and run Noise2Self ---\n",
    "    train_loader = DataLoader(mnist, batch_size=64, shuffle=True)\n",
    "    noise2self_model = DenoiserCNN().to(device)\n",
    "    train_noise2self_model(noise2self_model, train_loader, device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        n2s_input = torch.tensor(beta_hat_np.reshape(1, 1, 28, 28)).float().to(device)\n",
    "        n2s_out = noise2self_model(n2s_input).squeeze().cpu().numpy()\n",
    "        rmse_n2s = np.sqrt(np.mean((n2s_out - beta_np.reshape(28, 28))**2))\n",
    "\n",
    "    baselines = benchmark_baselines(beta_np, beta_hat_np)\n",
    "\n",
    "    vmin, vmax = beta_np.min(), beta_np.max()\n",
    "\n",
    "    # --- Plot ---\n",
    "    fig, axes = plt.subplots(2, 4, figsize=(18, 10))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    axes[0].imshow(beta_hat_np.reshape(28, 28), cmap=\"viridis\", vmin=vmin, vmax=vmax)\n",
    "    axes[0].set_title(f\"Noisy\\nRMSE={rmse_noisy:.4f}\")\n",
    "\n",
    "    axes[1].imshow(beta_np.reshape(28, 28), cmap=\"viridis\", vmin=vmin, vmax=vmax)\n",
    "    axes[1].set_title(\"Original\")\n",
    "\n",
    "    axes[2].imshow(beta_denoised_np, cmap=\"viridis\", vmin=vmin, vmax=vmax)\n",
    "    axes[2].set_title(f\"NASH-Fused\\nRMSE={rmse_denoised:.4f}\")\n",
    "\n",
    "    axes[3].imshow(n2s_out, cmap=\"viridis\", vmin=vmin, vmax=vmax)\n",
    "    axes[3].set_title(f\"Noise2Self\\nRMSE={rmse_n2s:.4f}\")\n",
    "\n",
    "    # Plot only selected baselines (remove \"Median\")\n",
    "    filtered_baselines = {k: v for k, v in baselines.items() if k != \"Median\"}\n",
    "    for i, (name, (img, rmse)) in enumerate(filtered_baselines.items(), start=4):\n",
    "        axes[i].imshow(img.reshape(28, 28), cmap=\"viridis\", vmin=vmin, vmax=vmax)\n",
    "        axes[i].set_title(f\"{name}\\nRMSE={rmse:.4f}\", fontsize=14)\n",
    "\n",
    "    for ax in axes:\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "af776ae0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 50, Loss: -0.2029\n",
      "Epoch 100, Loss: -0.2687\n",
      "Epoch 150, Loss: -0.2892\n",
      "Epoch 200, Loss: -0.2992\n",
      "Epoch 250, Loss: -0.3067\n",
      "Epoch 300, Loss: -0.3146\n",
      "Epoch 350, Loss: -0.3228\n",
      "Epoch 400, Loss: -0.3329\n",
      "Epoch 450, Loss: -0.3447\n",
      "Epoch 500, Loss: -0.3569\n",
      "Epoch 550, Loss: -0.3673\n",
      "Epoch 600, Loss: -0.3695\n",
      "Epoch 650, Loss: -0.3182\n",
      "Epoch 700, Loss: -0.3798\n",
      "Epoch 750, Loss: -0.3659\n",
      "Epoch 800, Loss: -0.3884\n",
      "RMSE (noisy): 0.2003\n",
      "RMSE (nash-fused): 0.0771\n",
      "[Epoch 1] Noise2Self Loss: 0.0516\n",
      "[Epoch 2] Noise2Self Loss: 0.0487\n",
      "[Epoch 3] Noise2Self Loss: 0.0483\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 67800 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run_pipeline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "85cc2c8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 50, Loss: -0.1005\n",
      "Epoch 100, Loss: -0.1727\n",
      "Epoch 150, Loss: -0.1872\n",
      "Epoch 200, Loss: -0.1937\n",
      "Epoch 250, Loss: -0.1986\n",
      "Epoch 300, Loss: -0.2040\n",
      "Epoch 350, Loss: -0.2090\n",
      "Epoch 400, Loss: -0.2126\n",
      "Epoch 450, Loss: -0.2180\n",
      "Epoch 500, Loss: -0.2244\n",
      "Epoch 550, Loss: -0.2140\n",
      "Epoch 600, Loss: -0.2393\n",
      "Epoch 650, Loss: -0.2196\n",
      "Epoch 700, Loss: -0.2468\n",
      "Epoch 750, Loss: -0.2510\n",
      "Epoch 800, Loss: -0.2574\n",
      "RMSE (noisy): 0.2099\n",
      "RMSE (nash-fused): 0.0924\n",
      "[Epoch 1] Noise2Self Loss: 0.0512\n",
      "[Epoch 2] Noise2Self Loss: 0.0485\n",
      "[Epoch 3] Noise2Self Loss: 0.0483\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 66188 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run_pipeline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8486cee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_efficient_benchmark(num_images=100, noise_level=0.2, ebgnn_epochs=300):\n",
    "    from tqdm import tqdm\n",
    "    import random\n",
    "\n",
    "    transform = transforms.ToTensor()\n",
    "    mnist = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    # Train Noise2Self model once\n",
    "    noise2self_model = DenoiserCNN().to(device)\n",
    "    train_loader = DataLoader(mnist, batch_size=64, shuffle=True)\n",
    "    train_noise2self_model(noise2self_model, train_loader, device, noise_std=noise_level, epochs=3)\n",
    "\n",
    "    methods = [\"Noisy\", \"NASH-Fused\", \"Noise2Self\", \"TV\", \"Fused Lasso\", \"Gaussian\", \"NLM\"]\n",
    "    results = {m: [] for m in methods}\n",
    "    examples_to_plot = []\n",
    "\n",
    "    for idx in tqdm(range(num_images), desc=\"Running Benchmark\"):\n",
    "        img = mnist.data[np.random.randint(0, len(mnist))].float() / 255.0\n",
    "        img = img[:28, :28]\n",
    "\n",
    "        beta, beta_hat, se, node_feat, edge_index, _ = build_graph_from_mnist_image(img, noise_level=noise_level)\n",
    "        beta_hat, se, node_feat, edge_index = beta_hat.to(device), se.to(device), node_feat.to(device), edge_index.to(device)\n",
    "\n",
    "        # Train EB-GNN model\n",
    "        model = EB_GNN_Denoiser(ManualStyleGNN(node_feat_dim=3, hidden_dim=64).to(device))\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "        for epoch in range(ebgnn_epochs):\n",
    "            optimizer.zero_grad()\n",
    "            loss = model.marginal_log_likelihood(beta_hat, se, node_feat, edge_index)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            beta_denoised, *_ = model.posterior_mean(beta_hat, se, node_feat, edge_index)\n",
    "            beta_np = beta.cpu().numpy()\n",
    "            beta_hat_np = beta_hat.cpu().numpy()\n",
    "            beta_denoised_np = beta_denoised.cpu().numpy().reshape(28, 28)\n",
    "\n",
    "        # Noise2Self output\n",
    "        with torch.no_grad():\n",
    "            n2s_input = torch.tensor(beta_hat_np.reshape(1, 1, 28, 28)).float().to(device)\n",
    "            n2s_out = noise2self_model(n2s_input).squeeze().cpu().numpy()\n",
    "\n",
    "        # Classical baselines\n",
    "        baselines = benchmark_baselines(beta_np, beta_hat_np)\n",
    "        baselines = {k: v for k, v in baselines.items() if k != \"Median\"}\n",
    "\n",
    "        # Compute RMSE\n",
    "        rmse_noisy = np.sqrt(np.mean((beta_hat_np - beta_np)**2))\n",
    "        rmse_denoised = np.sqrt(np.mean((beta_denoised_np.flatten() - beta_np)**2))\n",
    "        rmse_n2s = np.sqrt(np.mean((n2s_out - beta_np.reshape(28, 28))**2))\n",
    "        rmse_baselines = {k: np.sqrt(np.mean((v[0] - beta_np)**2)) for k, v in baselines.items()}\n",
    "\n",
    "        results[\"Noisy\"].append(rmse_noisy)\n",
    "        results[\"NASH-Fused\"].append(rmse_denoised)\n",
    "        results[\"Noise2Self\"].append(rmse_n2s)\n",
    "        for k, val in rmse_baselines.items():\n",
    "            results[k].append(val)\n",
    "\n",
    "        # Store 5 random examples for later plotting\n",
    "        if len(examples_to_plot) < 5 and random.random() < 0.15:\n",
    "            vmin, vmax = beta_np.min(), beta_np.max()\n",
    "            panel = [\n",
    "                (\"Original\", beta_np.reshape(28, 28)),\n",
    "                (\"Noisy\", beta_hat_np.reshape(28, 28)),\n",
    "                (\"NASH-Fused\", beta_denoised_np),\n",
    "                (\"Noise2Self\", n2s_out),\n",
    "            ] + [(k, v[0].reshape(28, 28)) for k, v in baselines.items()]\n",
    "            examples_to_plot.append((panel, vmin, vmax))\n",
    "\n",
    "    return results, examples_to_plot\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a20fc6e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch 1] Noise2Self Loss: 0.0512\n",
      "[Epoch 2] Noise2Self Loss: 0.0489\n",
      "[Epoch 3] Noise2Self Loss: 0.0485\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running Benchmark: 100%|██████████| 100/100 [01:17<00:00,  1.29it/s]\n"
     ]
    }
   ],
   "source": [
    "results, examples = run_efficient_benchmark(num_images=100)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fa70cf5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Flatten the results into a list of dicts\n",
    "data = []\n",
    "for method, rmse_list in results.items():\n",
    "    for rmse in rmse_list:\n",
    "        data.append({\"Method\": method, \"RMSE\": rmse})\n",
    "\n",
    "# Create DataFrame\n",
    "df = pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "04186d7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53404 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "sns.boxplot(data=df, x=\"Method\", y=\"RMSE\")\n",
    "plt.xticks(rotation=30)\n",
    "plt.title(\"RMSE Comparison Across Methods (100 MNIST digits)\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
