{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ddmms/camml-tutorials/blob/main/notebooks/05-generative/autoencoders.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autoencoders, Variational Autoencoders and Conditional Varaiational Autoencoders\n",
    "\n",
    "Autoencoders (AEs) are neural networks designed for unsupervised learning, primarily used for dimensionality reduction and feature extraction. They consist of an encoder, which compresses input data into a lower-dimensional latent space, and a decoder, which reconstructs the original input from this compressed representation. Standard autoencoders can learn meaningful representations but often struggle with generating diverse, high-quality samples.\n",
    "\n",
    "Variational Autoencoders (VAEs) improve upon AEs by introducing a probabilistic framework, ensuring that the learned latent space follows a structured distribution (typically Gaussian). Instead of mapping inputs to fixed points in the latent space, VAEs learn distributions, allowing for smooth interpolation and controlled sampling. This makes them powerful for generative tasks, as they can generate novel data points that resemble the training distribution.\n",
    "\n",
    "Conditional Variational Autoencoders (cVAEs) extend VAEs by incorporating additional information, such as labels or descriptors, into the encoding and decoding process. This enables the generation of data conditioned on specific attributes, making cVAEs particularly useful in chemistry for designing molecules with desired properties.\n",
    "\n",
    "These models form the backbone of modern generative approaches in machine learning, with applications in molecular design, materials discovery, and reaction prediction.\n",
    "\n",
    "## Further Reading on Autoencoders, VAEs, and cVAEs\n",
    "\n",
    "### Autoencoders\n",
    "- [Autoencoder - Wikipedia](https://en.wikipedia.org/wiki/Autoencoder)\n",
    "\n",
    "### Variational Autoencoders (VAEs)\n",
    "- [Variational Autoencoder - Wikipedia](https://en.wikipedia.org/wiki/Variational_autoencoder)\n",
    "- [Tutorial on Variational Autoencoders - arXiv](https://arxiv.org/pdf/1606.05908)\n",
    "- [Variational Autoencoders: How They Work and Why They Matter](https://www.datacamp.com/tutorial/variational-autoencoders)\n",
    "\n",
    "### Conditional Variational Autoencoders (cVAEs)\n",
    "\n",
    "- [Conditional Variational Auto-encoder - Pyro Tutorials](https://pyro.ai/examples/cvae.html)\n",
    "- [Hands-on with Conditional Variational Autoencoders (CVAE) - YouTube](https://www.youtube.com/watch?v=e4izCK7l4Kc)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Dataset\n",
    "\n",
    "In this notebook, we will use the `mnist-digits` dataset. It is simpler than the `mnist-fashion` dataset, allowing us to use only two latent features so that we can conveniently visualise and examine the encodings distribution in the latent space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load MNIST dataset\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Lambda(lambda x: x.view(-1))])\n",
    "mnist_train = datasets.MNIST(root=\"./data\", train=True, transform=transform, download=True)\n",
    "mnist_test = datasets.MNIST(root=\"./data\", train=False, transform=transform, download=True)\n",
    "train_loader = DataLoader(mnist_train, batch_size=128, shuffle=True)\n",
    "test_loader = DataLoader(mnist_test, batch_size=128, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Take a look at the size of the data. \n",
    "\n",
    "We have 60k (28, 28) iamges over 10 separate classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print dataset info\n",
    "print(f'Training dataset size: {len(mnist_train)}')\n",
    "print(f'Image size (pixels per sample): {mnist_train.data.shape[1] * mnist_train.data.shape[2]}')\n",
    "print(f'Number of classes: {len(mnist_train.classes)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset visualisation\n",
    "### Look at some of the samples from the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_images(images, title=\"Images\"):\n",
    "    fig, axes = plt.subplots(4, 4, figsize=(5, 5))\n",
    "    for i, ax in enumerate(axes.flat):\n",
    "        ax.imshow(images[i].squeeze().reshape(28, 28), cmap='gray')\n",
    "        ax.axis('off')\n",
    "    plt.suptitle(title)\n",
    "    plt.show()\n",
    "\n",
    "# Visualize some training images\n",
    "images, _ = next(iter(train_loader))\n",
    "plot_images(images[:16], \"Training Data Samples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autoencoder\n",
    "Why we need a VAE? To answer this question, let us start with an ordinary AE and see what is unsatisfactory when we use it to generate new images. \n",
    "\n",
    "## Model Structure\n",
    "The autoencoder consists of two main components: an **encoder** and a **decoder**.\n",
    "\n",
    "### Encoder\n",
    "- The encoder compresses the input image into a **latent representation**.\n",
    "- It consists of:\n",
    "  - A fully connected layer (`fc1`) that transforms the input into a hidden representation.\n",
    "  - A final fully connected layer (`fc_latent`) that reduces the hidden representation into a **2D latent space**.\n",
    "\n",
    "### Decoder\n",
    "- The decoder reconstructs the original image from the latent representation.\n",
    "- It consists of:\n",
    "  - A fully connected layer (`fc_decode`) that expands the latent representation back to the hidden dimension.\n",
    "  - A final fully connected layer (`fc_out`) that reconstructs the original image using a **sigmoid activation**.\n",
    "\n",
    "## Loss Function\n",
    "The autoencoder is trained using the **binary cross-entropy (BCE) loss**:\n",
    "\n",
    "$ \\mathcal{L}(x, \\hat{x}) = - \\sum x \\log(\\hat{x}) + (1 - x) \\log(1 - \\hat{x}) $\n",
    "\n",
    "where:\n",
    "-  $x$ is the original input image.\n",
    "- $\\hat{x}$ is the reconstructed image.\n",
    "- The loss measures how well the reconstructed image matches the original.\n",
    "\n",
    "This approach helps the autoencoder learn meaningful features from the data and reconstruct images accurately.\n",
    "\n",
    "## Training Process\n",
    "- The model is trained using the **Adam optimizer** with a learning rate of $10^{-3}$.\n",
    "- The loss is minimized over multiple epochs using a batch size of 128.\n",
    "- The encoder's learned representations can be visualized in a **2D latent space** to analyze digit clustering.\n",
    "\n",
    "## Visualization\n",
    "1. **Reconstruction**: Displays original and reconstructed images to assess model performance.\n",
    "2. **Latent Space**: Plots the encoded representations to understand digit distributions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the Autoencoder\n",
    "class Autoencoder(nn.Module):\n",
    "    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=2):\n",
    "        super(Autoencoder, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.fc_latent = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.fc_decode = nn.Linear(latent_dim, hidden_dim)\n",
    "        self.fc_out = nn.Linear(hidden_dim, input_dim)\n",
    "\n",
    "    def encode(self, x):\n",
    "        h = F.relu(self.fc1(x))\n",
    "        return self.fc_latent(h)\n",
    "\n",
    "    def decode(self, z):\n",
    "        h = F.relu(self.fc_decode(z))\n",
    "        return torch.sigmoid(self.fc_out(h))\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.encode(x)\n",
    "        return self.decode(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "### Set up the training loop\n",
    "\n",
    "Using this code template, set up the code to train for 20 epochs with a latent dimensionality of 2. Use a learning rate of 0.001\n",
    "\n",
    "```python\n",
    "# Initialize the model, optimizer, and train\n",
    "epochs = \n",
    "latent_dim = \n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "autoencoder = Autoencoder(latent_dim=latent_dim).to(device)\n",
    "optimizer = optim.Adam(autoencoder.parameters(), lr=)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize the model, optimizer, and train\n",
    "epochs = 20\n",
    "latent_dim = 2\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "autoencoder = Autoencoder(latent_dim=latent_dim).to(device)\n",
    "optimizer = optim.Adam(autoencoder.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loss function\n",
    "def loss_function(recon_x, x):\n",
    "    return F.binary_cross_entropy(recon_x, x, reduction='sum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    autoencoder.train()\n",
    "    train_loss = 0\n",
    "    for x, _ in train_loader:\n",
    "        x = x.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        recon_x = autoencoder(x)\n",
    "        loss = loss_function(recon_x, x)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    print(f\"Epoch {epoch+1}, Loss: {train_loss/len(train_loader.dataset):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visulaise the results\n",
    "\n",
    "### Reconstructed outputs\n",
    "\n",
    "We first plot a series of input images and their outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Visualize original and reconstructed images\n",
    "def visualize_reconstructions():\n",
    "    autoencoder.eval()\n",
    "    with torch.no_grad():\n",
    "        for x, _ in test_loader:\n",
    "            x = x.to(device)\n",
    "            recon_x = autoencoder(x)\n",
    "            break\n",
    "    x = x.cpu().view(-1, 28, 28)[:10]\n",
    "    recon_x = recon_x.cpu().view(-1, 28, 28)[:10]\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 10, figsize=(10, 2))\n",
    "    for i in range(10):\n",
    "        axes[0, i].imshow(x[i], cmap=\"gray\")\n",
    "        axes[1, i].imshow(recon_x[i], cmap=\"gray\")\n",
    "        axes[0, i].axis(\"off\")\n",
    "        axes[1, i].axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "visualize_reconstructions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualise the latent space\n",
    "\n",
    "We can look at the latent space and see how the different digits are distributed. Note that although this was technically an unsupervised learning exercise, we do have access to labels, but they are never used in training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Visualize latent space\n",
    "def visualize_latent_space():\n",
    "    autoencoder.eval()\n",
    "    embeddings, labels = [], []\n",
    "    with torch.no_grad():\n",
    "        for x, y in test_loader:\n",
    "            x = x.to(device)\n",
    "            z = autoencoder.encode(x)\n",
    "            embeddings.append(z.cpu())\n",
    "            labels.append(y.cpu())\n",
    "    \n",
    "    embeddings = torch.cat(embeddings).numpy()\n",
    "    labels = torch.cat(labels).numpy()\n",
    "    \n",
    "    plt.figure(figsize=(8, 6))\n",
    "    scatter = plt.scatter(embeddings[:, 0], embeddings[:, 1], c=labels, cmap=\"tab10\", alpha=0.5)\n",
    "    plt.colorbar(scatter, label=\"Digit Label\")\n",
    "    plt.title(\"Latent Space Visualization\")\n",
    "    plt.xlabel(\"z1\")\n",
    "    plt.ylabel(\"z2\")\n",
    "    plt.show()\n",
    "\n",
    "visualize_latent_space()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualise Digit distributions\n",
    "\n",
    "For a differnt perspective we can plot the distribtion of the latent dimensions of the different digits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# KDE Plot for each digit in subplots\n",
    "def visualize_kde_plots():\n",
    "    autoencoder.eval()\n",
    "    embeddings, labels = [], []\n",
    "    with torch.no_grad():\n",
    "        for x, y in test_loader:\n",
    "            x = x.to(device)\n",
    "            z = autoencoder.encode(x)\n",
    "            embeddings.append(z.cpu())\n",
    "            labels.append(y.cpu())\n",
    "    \n",
    "    embeddings = torch.cat(embeddings).numpy()\n",
    "    labels = torch.cat(labels).numpy()\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 5, figsize=(15, 6))\n",
    "    axes = axes.flatten()\n",
    "    for digit in range(10):\n",
    "        idx = labels == digit\n",
    "        sns.kdeplot(x=embeddings[idx, 0], y=embeddings[idx, 1], fill=True, ax=axes[digit])\n",
    "        axes[digit].set_title(f\"Digit {digit}\")\n",
    "        axes[digit].set_xlabel(\"z1\")\n",
    "        axes[digit].set_ylabel(\"z2\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "visualize_kde_plots()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Regularity of the latent space\n",
    "\n",
    "Both the scatter plot and the histogram plots show that the data distributions in the latent space are rather *irregular*. Some of the digits have very wide distributions (such as 1 and 7) and some very narrow distributions (such as 2 and 3). \n",
    "\n",
    "Remember that our goal of training this AE is neither dimensionality reduction nor denoising but to generate new images out of the original dataset. Image generation is done by the decoder, taking the latent representation (`X` and `Y` in the plots) as the input. An irregular latent space makes image generation less controllable and robust. Taking our case for example, two shortcomings are likely to emerge:\n",
    "\n",
    "1. **Controllability**: sampling the entire latent space, we will generate much more of the widely distributed digits than the narrowly distributed ones; instead, if we limit the range of the latent space, we will loss some characteristics of the widely distributed ones;\n",
    "\n",
    "2. **Robustness**: images that do not resemble any of the digits will be generated by the gaps between the distributions of the digits; such gaps increase with the range of the latent space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational Autoencoder\n",
    "\n",
    "Overfitting is the essential reason behind an irregular latent space of a naive AE, that is, the neural networks for encoding and decoding try their best to fit the data from end to end without caring about how the latent space is organised with respect to the original data. A VAE can regularise the latent space by imposing additional distributional properties on the latent space.\n",
    "\n",
    "The following figure summarises **the two extensions** from an AE to a VAE:\n",
    "\n",
    "1. Unlike a naive AE that encodes an input data $x$ as a single point $z$ in the latent space, a VAE encodes it as a normal distribution $\\mathcal{N}(\\mu, \\sigma)$, and the latent representation $z$ is sampled from this distribution and then passed to the decoder;\n",
    "\n",
    "2. An AE only minimises the reconstruction error $\\lVert x-x'\\rVert^2$ to fit the data, whereas a VAE minimises the sum of the reconstruction error and the KL divergence (Kullback–Leibler divergence) between the latent distribution $\\mathcal{N}(\\mu, \\sigma)$ and the standard normal distribution $\\mathcal{N}(0, 1)$.\n",
    "\n",
    "How does a VAE regularise the latent space? The loss function provides a straightforward answer: in addition to fitting the data by minimising the reconstruction error, it also drags the latent distribution to a standard normal distribution. The final model is a trade-off between the two effects. Also, because each input image is encoded as a Gaussian blob instead of a single point, the gaps in the latent space can be filled by such blurring so that meaningless decodings can be largely avoided.\n",
    "\n",
    "![ae-vae.png](https://github.com/stfc-sciml/sciml-workshop/blob/master/course_3.0_with_solutions/markdown_pic/ae-vae.png?raw=1)\n",
    "\n",
    "\n",
    "**📚 Theory for VAEs** \n",
    "\n",
    "<details> <summary>Reveal / Hide</summary> \n",
    "<p>\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "### Derivation of ELBO\n",
    "\n",
    "Posterior Probability $p(z|x)$, which can be expressed as:\n",
    "\n",
    "\\begin{eqnarray}\n",
    "p(z|x)&=&\\frac{p(x|z)p(z)}{p(x)}\\nonumber\\\\\n",
    "&=& \\frac{p(x|z)p(z)}{\\int p(x|z)p(z)}\n",
    "\\end{eqnarray}\n",
    "\n",
    "where $\\int p(x|z)p(z)$, which is the marginal, can be intractable and cannot be computed directly. One way  to compute the overall solution $p(z|x)$ is using Monte Carlo methods (such as sampling). The method used in this notebook (and the underlying VAE paper) is  ***variational inference***. \n",
    "\n",
    "The idea is to identify another proxy distribution $q(z|x)$ that reasonably approximates $p(z|x)=p(x|z)p(z)$. i.e. if the KL-divergence between two pdfs, $q(x)$ and $p(z|x)$ is denoted by\n",
    "\n",
    "$$\\mathrm{KL}(q(x)||p(z|x))$$\n",
    "\n",
    "\n",
    "it can be minimized by selecting an alternative pdf $q(z|x)$, which is a good proxy for $p(z|x)$. But \n",
    "\n",
    "\\begin{eqnarray}\n",
    "\\mathrm{KL}(q(z|x)||p(z|x)) &=& -\\int q(z|x)\\log\\frac{p(z|x)}{q(z|x)} dz\\nonumber\\\\\n",
    "                          &=& -\\int q(z|x)\\log\\frac{p(x|z)p(z)}{p(x)q(z|x)} dz\\nonumber\\\\\n",
    "                          &=& -\\int q(z|x)\\log\\frac{p(x|z)p(z)}{q(z|x)}dz + \\int_{z} q(z|x)\\log p(x)dz \\nonumber\\\\\n",
    "                          &=& -\\int q(z|x)\\log\\frac{p(x|z)p(z)}{q(z|x)} +  \\log p(x)\\int_{z} q(z|x)dz\\nonumber\\\\   \n",
    "                          &=& -\\int q(z|x)\\log\\frac{p(x|z)p(z)}{q(z|x)}dz +  \\log p(x)\\nonumber\\\\\n",
    "                          &=& -\\int q(z|x)\\log\\frac{p(z)}{q(z|x)}dz -\\int q(z|x)\\log{p(x|z)}dz + \\log p(x) \n",
    "\\end{eqnarray}\n",
    "\n",
    "\n",
    "Given that $\\mathrm{KL}\\left(q(z|x)||p(z|x)\\right)\\geq 0$, \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\\begin{eqnarray}\n",
    "-\\int q(z|x)\\log\\frac{p(z)}{q(z|x)}dz -\\int q(z|x)\\log{p(x|z)}dz + \\log p(x) &\\geq& 0 \\\\\n",
    "\\log p(x) &\\geq& \\int q(z|x)\\log\\frac{p(z)}{q(z|x)}dz + \\int q(z|x)\\log{p(x|z)}dz\\\\\n",
    "\\log p(x) &\\geq& - \\mathrm{KL}(q(z|x)||p(z)) + \\int q(z|x)\\log p(x|z)dz \\nonumber\\\\\n",
    "\\log p(x) &\\geq& - \\mathrm{KL}(q(z|x)||p(z)) + \\mathbb{E}_{q(z|x)}\\left[\\log p(x|z)\\right] \\nonumber\\\\\n",
    "\\end{eqnarray}\n",
    "\n",
    "This is the *variational lower-bound*, or the evidence of lower bound (ELBO).  This remains as the objective function for the VAE. However, frameworks like TensorFlow or PyTorch need a loss function to be minimized. Maximising the log likelihood of the model evidence $p(x)$ is same as minimizing the $-\\log p(x)$. The first term of the ELBO, namely,  $\\mathrm{KL}(q(z|x)||p(z))$ is the *regularising* term and constrains the posterior distribution. The second term of the ELBO models the reconstruction loss. \n",
    "\n",
    "Now, this leaves fair bit of freedom on the choice of the prior $p(z)$. Let's assume:\n",
    "\n",
    "\n",
    "$$\n",
    "p(z)={\\cal N}(\\mu_p, \\sigma_p^2)\n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "q(z|x)={\\cal N}(\\mu_q, \\sigma_q^2)\n",
    "$$\n",
    "\n",
    "Thus, \n",
    "\n",
    "$$\n",
    "p(z)=\\frac{1}{\\sqrt{2\\pi\\sigma_p^2}}\\exp\\left(\\frac{(x-\\mu_p)^2}{2\\sigma_p^2}\\right)\n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "q(z|x)=\\frac{1}{\\sqrt{2\\pi\\sigma_q^2}}\\exp\\left(\\frac{(x-\\mu_q)^2}{2\\sigma_q^2}\\right)\n",
    "$$\n",
    "\n",
    "The direct derivation of $\\mathrm{KL}(q(z|x)||p(z))$ will give (with some simplifications)\n",
    "\n",
    "\n",
    "$$\n",
    "-\\mathrm{KL}(q(z|x)||p(z)) = \\log\\frac{\\sigma_q}{\\sigma_p} - \\frac{\\left(\\log\\sigma_q^2-(\\mu_p-\\mu_q)^2\\right)}{2\\sigma_p^2} +\\frac{1}{2} \n",
    "$$\n",
    "\n",
    "By fixing the prior distribution $p(z)={\\cal N}(0,1^2)$, \n",
    "\n",
    "$$\n",
    "-\\mathrm{KL}(q(z|x)||p(z)) = \\frac{1}{2}\\left[ 1 + \\log\\sigma_q^2 - \\sigma_q^2 -\\mu_q^2\\right]\n",
    "$$\n",
    "\n",
    "Hence, the new ELBO is\n",
    "\n",
    "\n",
    "$$\n",
    "\\frac{1}{2}\\left[ 1 + \\log\\sigma_q^2 - \\sigma_q^2 -\\mu_q^2\\right] + \\mathbb{E}_{q(z|x)}\\left[\\log p(x|z)\\right] \n",
    "$$\n",
    "\n",
    "\n",
    "Let $J, B$ and $\\cal{L}$ be the dimension of the latent space, and the batch size over which the sampling is done. The loss function we need to minimise (from the point of  implementation) is\n",
    "\n",
    "$$\n",
    "{\\cal L} = - \\sum_{j=1}^J \\frac{1}{2}\\Bigl[ 1 + \\log\\sigma_j^2 - \\sigma_j^2 -\\mu_j^2\\Bigr] - \\frac{1}{B}\\sum_{l}\\mathbb{E}_{q(z|x_i)}\\left[\\log p(x_i|z^{(i,l)})\\right] \n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "This can be observed in the code implementation below (see function implementation ``loss_function`` below)\n",
    "\n",
    "### Reparameterisation\n",
    "\n",
    "A valid reparameterization would be \n",
    "\n",
    "$$\n",
    "z = \\mu+\\sigma\\epsilon\n",
    "$$\n",
    "\n",
    "\n",
    "where $\\epsilon$ is an auxiliary noise variable $\\epsilon\\sim{\\cal{N}}(0, 1)$, which actually enables the reparameterization technique. Although it is possible to use $\\sigma$ or more specifically $\\sigma^2$, working on log scales improves the stability. i.e. \n",
    "\n",
    "\\begin{eqnarray}\n",
    "p &=& \\log(\\sigma^2)\\\\\n",
    "&=& 2  \\log(\\sigma)\n",
    "\\end{eqnarray}\n",
    "\n",
    "To get the log standard deviation, $\\log(\\sigma)$, \n",
    "\\begin{eqnarray}\n",
    "\\log(\\sigma) &=&  p/2 \\\\\n",
    "\\label{eqn:log_sigma}\n",
    "\\end{eqnarray} \n",
    "\n",
    "and hence\n",
    "\n",
    "$$\n",
    "\\sigma = \\exp^{p/2}\n",
    "$$\n",
    "\n",
    "The resulting estimator (or the loss function) becomes (see Page 5 of [Auto-Encoding Variational Bayes Paper](https://arxiv.org/abs/1312.6114)),\n",
    "\n",
    "$$\n",
    "-\\text{KLD} = \\frac{1}{2}\\sum_{j=1}^{J}(1+\\log(\\sigma_j^i)^2 - (\\mu_j^i)^2 -(\\sigma_j^i)^2)\n",
    "$$\n",
    "\n",
    "\n",
    "It is important to see that the KL divergence can be computed and differentiated without estimation. This is a very remarkable thing (no esimtation!).\n",
    "\n",
    "The $\\boldsymbol{\\epsilon}$ must be sampled from a zero-mean, unit-variance Gaussian distribution, and should be of the same size as $\\boldsymbol{\\sigma}$. \n",
    "\n",
    "\n",
    "<br>\n",
    "\n",
    "\n",
    "    \n",
    "</p>\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model structure\n",
    "\n",
    "### Encoder\n",
    "\n",
    "Maps the input data to a latent space by computing the mean (`μ`) and log-variance (`log(σ²)`) using fully connected layers.\n",
    "  - A **fully connected layer** (`fc1`) maps the input (`input_dim = 784`) to a hidden dimension (`hidden_dim = 400`).\n",
    "  - Two separate **fully connected layers** (`fc_mu` and `fc_logvar`) transform the hidden representation into the mean (`μ`) and log-variance (`logvar`) of the latent space (`latent_dim = 2`).\n",
    "\n",
    "### Decoder\n",
    "\n",
    "Reconstructs the input from the latent representation `z`.  \n",
    "  - A **fully connected layer** (`fc_decode`) maps `z` back to the hidden dimension (`hidden_dim = 400`).\n",
    "  - A second **fully connected layer** (`fc_out`) transforms the hidden representation back to the original input dimension (`input_dim = 784`).\n",
    "  - The final output is passed through a **sigmoid activation** to constrain pixel values between 0 and 1.\n",
    "\n",
    "## Loss function\n",
    "\n",
    "The loss function is the VAE loss discussed above.\n",
    "\n",
    "## Training Process\n",
    "- The model is trained using the **Adam optimizer** with a learning rate of $10^{-3}$.\n",
    "- The loss is minimized over multiple epochs using a batch size of 128.\n",
    "- The encoder's learned representations can be visualized in a **2D latent space** to analyze digit clustering.\n",
    "\n",
    "## Visualization\n",
    "1. **Reconstruction**: Displays original and reconstructed images to assess model performance.\n",
    "2. **Latent Space**: Plots the encoded representations to understand digit distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the Variational Autoencoder\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=2):\n",
    "        super(VAE, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.fc_mu = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.fc_decode = nn.Linear(latent_dim, hidden_dim)\n",
    "        self.fc_out = nn.Linear(hidden_dim, input_dim)\n",
    "\n",
    "    def encode(self, x):\n",
    "        h = F.relu(self.fc1(x))\n",
    "        mu = self.fc_mu(h)\n",
    "        logvar = self.fc_logvar(h)\n",
    "        return mu, logvar\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps * std\n",
    "\n",
    "    def decode(self, z):\n",
    "        h = F.relu(self.fc_decode(z))\n",
    "        return torch.sigmoid(self.fc_out(h))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar\n",
    "    \n",
    "def loss_function(recon_x, x, mu, logvar):\n",
    "    recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')\n",
    "    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "    return recon_loss + kl_div"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "The setup for the VAE is very similar to the AE that we did earlier.\n",
    "\n",
    "### Set up the training loop\n",
    "\n",
    "Based on the code that you used earlier for the AE - set up the trainig parameters for the VAE.\n",
    "\n",
    "<details>\n",
    "  <summary>Click to expand code</summary>\n",
    "\n",
    "  ```python\n",
    "# Initialize the model, optimizer, and train \n",
    "epochs = 20 \n",
    "latent_dim = 2 \n",
    "vae = VAE(latent_dim=latent_dim).to(device)\n",
    "optimizer = optim.Adam(vae.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the model, optimizer, and train\n",
    "epochs = 20\n",
    "latent_dim = 2\n",
    "vae = VAE(latent_dim=latent_dim).to(device)\n",
    "optimizer = optim.Adam(vae.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    vae.train()\n",
    "    train_loss = 0\n",
    "    for x, _ in train_loader:\n",
    "        x = x.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        recon_x, mu, logvar = vae(x)\n",
    "        loss = loss_function(recon_x, x, mu, logvar)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    print(f\"Epoch {epoch+1}, Loss: {train_loss/len(train_loader.dataset):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise the results\n",
    "\n",
    "### Reconstructed plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Visualize original and reconstructed images\n",
    "def visualize_reconstructions():\n",
    "    vae.eval()\n",
    "    with torch.no_grad():\n",
    "        for x, _ in test_loader:\n",
    "            x = x.to(device)\n",
    "            recon_x, _, _ = vae(x)\n",
    "            break\n",
    "    x = x.cpu().view(-1, 28, 28)[:10]\n",
    "    recon_x = recon_x.cpu().view(-1, 28, 28)[:10]\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 10, figsize=(10, 2))\n",
    "    for i in range(10):\n",
    "        axes[0, i].imshow(x[i], cmap=\"gray\")\n",
    "        axes[1, i].imshow(recon_x[i], cmap=\"gray\")\n",
    "        axes[0, i].axis(\"off\")\n",
    "        axes[1, i].axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "visualize_reconstructions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Visualise the latent space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# KDE Plot for each digit in subplots\n",
    "def visualize_kde_plots():\n",
    "    vae.eval()\n",
    "    embeddings, labels = [], []\n",
    "    with torch.no_grad():\n",
    "        for x, y in test_loader:\n",
    "            x = x.to(device)\n",
    "            mu, _ = vae.encode(x)\n",
    "            embeddings.append(mu.cpu())\n",
    "            labels.append(y.cpu())\n",
    "    \n",
    "    embeddings = torch.cat(embeddings).numpy()\n",
    "    labels = torch.cat(labels).numpy()\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 5, figsize=(15, 6))\n",
    "    axes = axes.flatten()\n",
    "    for digit in range(10):\n",
    "        idx = labels == digit\n",
    "        sns.kdeplot(x=embeddings[idx, 0], y=embeddings[idx, 1], fill=True, ax=axes[digit])\n",
    "        axes[digit].set_title(f\"Digit {digit}\")\n",
    "        axes[digit].set_xlabel(\"z1\")\n",
    "        axes[digit].set_ylabel(\"z2\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Run visualization\n",
    "visualize_kde_plots()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conditional Variational Autoencoder (CVAE) Explanation\n",
    "\n",
    "## Overview\n",
    "A **Conditional Variational Autoencoder (CVAE)** is a type of generative model that extends the **Variational Autoencoder (VAE)** by incorporating **labels** during training and generation. This allows the model to generate specific outputs based on given labels.\n",
    "\n",
    "In our implementation, the CVAE is trained on the **MNIST dataset**, using digit labels to guide the latent space representation.\n",
    "\n",
    "## Model Structure\n",
    "The CVAE consists of three main components:\n",
    "\n",
    "### 1. Encoder\n",
    "The encoder maps an input image $x$ and its corresponding **one-hot encoded label** $ y $ into a latent space:\n",
    "- The input image $x $ (flattened 28×28 pixels) is concatenated with the label $y$.\n",
    "- This combined input is passed through a fully connected **encoder network**.\n",
    "- The network produces **mean** $\\mu$ and **log variance** $\\log \\\\sigma^2$for the latent representation.\n",
    "\n",
    "### 2. Latent Space Sampling (Reparameterization Trick)\n",
    "To enable backpropagation, the model uses the **reparameterization trick**:\n",
    "$\n",
    "z = \\mu + \\sigma \\cdot \\epsilon, \\quad \\epsilon \\sim \\mathcal{N}(0, I)\n",
    "$\n",
    "where $z$ is a sampled latent vector that represents the compressed information.\n",
    "\n",
    "### 3. Decoder\n",
    "The decoder reconstructs the image using the sampled latent vector $z$ and the label $y$:\n",
    "- The latent vector $z$ is concatenated with the label $y$.\n",
    "- This combined input is passed through a **decoder network**.\n",
    "- The network outputs a **reconstructed image**.\n",
    "\n",
    "## Loss Function\n",
    "The CVAE is trained using a combination of two loss terms:\n",
    "\n",
    "1. **Reconstruction Loss**: Measures how accurately the reconstructed image resembles the original.\n",
    "$\n",
    "   \\mathcal{L}_{\\text{recon}} = \\sum x \\log(\\hat{x}) + (1 - x) \\log(1 - \\hat{x})\n",
    "$\n",
    "\n",
    "2. **KL Divergence Loss**: Encourages the learned latent distribution to be close to a normal distribution.\n",
    "$\n",
    "   \\mathcal{L}_{\\text{KL}} = -\\frac{1}{2} \\sum (1 + \\log \\sigma^2 - \\mu^2 - \\sigma^2)\n",
    "$\n",
    "\n",
    "The total loss is:\n",
    "$\n",
    "\\mathcal{L} = \\mathcal{L}_{\\text{recon}} + \\mathcal{L}_{\\text{KL}}\n",
    "$\n",
    "\n",
    "## Code Explanation\n",
    "\n",
    "### 1. Model Definition (`CVAE`)\n",
    "- `encode(x, y)`: Encodes the input and label into latent space (produces $\\mu$ and $\\log \\sigma^2$).\n",
    "- `reparameterize(mu, logvar)`: Samples from the latent space.\n",
    "- `decode(z, y)`: Reconstructs the image from the latent vector and label.\n",
    "- `forward(x, y)`: Full forward pass through encoder, reparameterization, and decoder.\n",
    "\n",
    "### 2. Training (`train_cvae`)\n",
    "- Uses **Adam optimizer** to minimize the total loss.\n",
    "- Runs for multiple epochs, updating weights via backpropagation.\n",
    "\n",
    "### 3. Conditional Image Generation (`generate_from_prompt`)\n",
    "- Allows users to generate images by specifying a digit.\n",
    "- Uses a random latent vector combined with a one-hot encoded label.\n",
    "\n",
    "## Summary\n",
    "- The **CVAE** can generate **conditional outputs**, making it more powerful than a standard VAE.\n",
    "- The **latent space visualization** helps understand how different digits are encoded.\n",
    "- The **generation function** allows creating specific digits on demand.\n",
    "\n",
    "This implementation provides a flexible way to explore generative modeling using deep learning. 🚀\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " \n",
    "class CVAE(nn.Module):\n",
    "    def __init__(self, feature_size, latent_size, class_size):\n",
    "        super(CVAE, self).__init__()\n",
    "        self.feature_size = feature_size\n",
    "        self.class_size = class_size\n",
    "\n",
    "        # encode\n",
    "        self.fc1  = nn.Linear(feature_size + class_size, 400)\n",
    "        self.fc21 = nn.Linear(400, latent_size)\n",
    "        self.fc22 = nn.Linear(400, latent_size)\n",
    "\n",
    "        # decode\n",
    "        self.fc3 = nn.Linear(latent_size + class_size, 400)\n",
    "        self.fc4 = nn.Linear(400, feature_size)\n",
    "\n",
    "        self.elu = nn.ELU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def encode(self, x, c): # Q(z|x, c)\n",
    "        '''\n",
    "        x: (bs, feature_size)\n",
    "        c: (bs, class_size)\n",
    "        '''\n",
    "        inputs = torch.cat([x, c], 1) # (bs, feature_size+class_size)\n",
    "        h1 = self.elu(self.fc1(inputs))\n",
    "        z_mu = self.fc21(h1)\n",
    "        z_var = self.fc22(h1)\n",
    "        return z_mu, z_var\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5*logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "\n",
    "    def decode(self, z, c): # P(x|z, c)\n",
    "        '''\n",
    "        z: (bs, latent_size)\n",
    "        c: (bs, class_size)\n",
    "        '''\n",
    "        inputs = torch.cat([z, c], 1) # (bs, latent_size+class_size)\n",
    "        h3 = self.elu(self.fc3(inputs))\n",
    "        return self.sigmoid(self.fc4(h3))\n",
    "\n",
    "    def forward(self, x, c):\n",
    "        mu, logvar = self.encode(x.view(-1, 28*28), c)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z, c), mu, logvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert labels to one-hot encoding\n",
    "def one_hot_encoding(labels, num_classes=10):\n",
    "    return torch.eye(num_classes, device=device)[labels]\n",
    "\n",
    "# Loss function\n",
    "def loss_function(recon_x, x, mu, logvar):\n",
    "    recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')\n",
    "    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "    return recon_loss + kl_div"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the model, optimizer, and train\n",
    "epochs = 20\n",
    "latent_dim = 20\n",
    "cvae = CVAE(feature_size=28*28, latent_size=latent_dim, class_size=10).to(device)\n",
    "optimizer = optim.Adam(cvae.parameters(), lr=5e-4)\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    cvae.train()\n",
    "    train_loss = 0\n",
    "    for x, y in train_loader:\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        y_one_hot = one_hot_encoding(y).to(device)\n",
    "        optimizer.zero_grad()\n",
    "        recon_x, mu, logvar = cvae(x, y_one_hot)\n",
    "        loss = loss_function(recon_x, x, mu, logvar)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    print(f\"Epoch {epoch+1}, Loss: {train_loss/len(train_loader.dataset):.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Conditional generation of new outputs based on prompts\n",
    "def generate_from_prompt(digit, num_samples=5):\n",
    "    cvae.eval()\n",
    "    fig, axes = plt.subplots(1, num_samples, figsize=(10, 2))\n",
    "    with torch.no_grad():\n",
    "        for i in range(num_samples):\n",
    "            y = torch.zeros(1, 10).to(device)\n",
    "            y[0, digit] = 1  # One-hot encode the desired digit\n",
    "            z = torch.randn(1, 20).to(device)  # Random latent vector\n",
    "            generated = cvae.decode(z, y).cpu().view(28, 28)\n",
    "            axes[i].imshow(generated, cmap=\"gray\")\n",
    "            axes[i].axis(\"off\")\n",
    "    plt.suptitle(f\"Generated Samples for Digit {digit}\")\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_from_prompt(9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can always try to improve the qulaity of the generated data by training for longer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "alignn-2",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}