sheet10.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1) Hands on Diffusion model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fShRSlQfDMoC"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "<!--   <td>\n",
    "    <a target=\"_blank\" href=\"https://www.infocusp.com/home\"><img src=\"https://www.infocusp.in/static/media/logo2.2b25740fa600fe779f23c6cf86e678b1.svg\" width=100 />Infocusp</a>\n",
    "  </td> -->\n",
    "<!--   <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/drive/1e8fQCnPVcEmT9Rdr2hu87qz120VqnJ97?authuser=1#scrollTo=C0EVLOt8MwG1\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td> -->\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/InFoCusp/diffusion_models\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View original source on GitHub</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8nSwjTSmApII"
   },
   "source": [
    "Author: Falak Shah, Infocusp Innovations Private Limited\n",
    "\n",
    "Translated from TensorFlow to PyTorch and modified by Roman Remme"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PGUQcdqYOZ_t"
   },
   "source": [
    "# Hands on Diffusion model\n",
    "\n",
    "[Diffusion models](https://arxiv.org/pdf/2006.11239.pdf) are a family of models that have shown amazing capability of generating photorealistic images with/ without text prompt. They have two flows as shown in the figure below - \n",
    "1. Deterministic forward flow (from image to noise) and \n",
    "2. Generative reverse flow (recreating image from noise).\n",
    "\n",
    "Diffusion models get their name from the forward flow where they follow a markov chain of diffusion steps, each of which adds a small amount of random noise to the data. Then they learn the model to reverse the diffusion process and construct desired data samples from noise. \n",
    "\n",
    "<figure>\n",
    "<p style=\"text-align:center;\"  align = \"center\"><img src=\"https://developer-blogs.nvidia.com/wp-content/uploads/2022/04/Fixed_Forward_Diffusion_Process.png\" alt=\"Trulli\" style=\"width:100%\"  align = \"center\"></p>\n",
    "<figcaption align = \"center\">Forward and reverse process <a href=\"https://developer.nvidia.com/blog/improving-diffusion-models-as-an-alternative-to-gans-part-1/\">Ref: Nvidia blog</a> </figcaption>\n",
    "</figure>\n",
    " \n",
    "\n",
    "\n",
    "Since they map noise to data, these models can be said to be capable of learning the distributions that generate data of any particular domain.\n",
    "\n",
    "This notebook showcases a minimal example of the forward diffusion process and its reverse mapping using a dense network. It is meant to give the reader side by side code snippets to match the equations in the paper and visual examples of the complete process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "choHKvGkOUbd"
   },
   "source": [
    "### Imports and utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "C0EVLOt8MwG1",
    "outputId": "63c9cd2d-dffd-4f43-e9c4-e6b22f0fe0c2"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "! pip install celluloid\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from celluloid import Camera\n",
    "from IPython.display import display\n",
    "import functools\n",
    "import sklearn.datasets\n",
    "from tqdm.auto import tqdm\n",
    "# For plotting\n",
    "from IPython.display import HTML\n",
    "from base64 import b64encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6cMVVPNCoiMC"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "# Utility function for displaying video inline\n",
    "\n",
    "def show_video(vname):\n",
    "    mp4 = open(vname,'rb').read()\n",
    "    data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "    display(HTML(\"\"\"\n",
    "    <video width=400 controls>\n",
    "        <source src=\"%s\" type=\"video/mp4\">\n",
    "    </video>\n",
    "    \"\"\" % data_url))\n",
    "\n",
    "def save_animation(vname, interval=30):\n",
    "    anim = camera.animate(blit=True, interval=interval)\n",
    "    anim.save(vname)\n",
    "    plt.close()\n",
    "\n",
    "# Utility function for random noise\n",
    "def noise_like(shape):\n",
    "    return torch.randn_like(shape=shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vm7lVgUvNiPB"
   },
   "source": [
    "## Data distribution\n",
    "\n",
    "Images can be thought of as points sampled from $height \\times width$ dimensional space. \n",
    "\n",
    "Consider an image of dimension $height \\times width$. Then the total number of pixels are $height \\times width$. Each pixel has a value ranging from 0 to 255.\n",
    "Now, consider a vector space, where we flatten this image and represent the intensity of each pixel along one dimension of the vector space. For example, an image with $height = 2$ and $ width =3$ (2px x 3px image) becomes a single vector of length 6 where each component of this vector will have a value between 0 to 255.\n",
    "\n",
    "So, in this image vector space, there are small clusters of valid (photorealistic) images sparsely distributed over the space. Rest of the vector space is made up of invalid (not real looking) images.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SDlU_wb8T-T5"
   },
   "source": [
    "For the example in this notebook, we consider a **hypothetical** simplified version of the above representation. We consider images made of just 2 pixels, each of which can have values between [-5, 5]. This is to allow visualization of each dimension of the data as it moves through the forward and reverse process (and additionally faster training 😅). \n",
    "\n",
    "The same code can be extended to the original image dimensions with just updated data dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JbdNUSgTNOjX",
    "outputId": "877f13d7-518d-4ff9-f7ef-b394baabb609"
   },
   "outputs": [],
   "source": [
    "# Evaluate to generate and use gaussian mixture dataset\n",
    "# Generate original points which are around [0.5, 0.5] in all quadrants and \n",
    "# 4 corners ([0,1], [1,0], [0,-1], [-1, 0])\n",
    "# Some region around these points indicates valid images region (true data distribution)\n",
    "\n",
    "num_samples_per_center = 1000\n",
    "stddev = 0.1\n",
    "mean = 0\n",
    "\n",
    "centers = torch.tensor([[0,1], [1,0], [0,-1], [-1, 0],\n",
    "                    [0.5, 0.5], [0.5, -0.5], [-0.5, -0.5], \n",
    "                    [-0.5,0.5]]) * 4\n",
    "\n",
    "all_data = []\n",
    "# Data for all clusters\n",
    "for center in centers:\n",
    "    center_data = torch.randn(size=(num_samples_per_center, 2)) * stddev + center\n",
    "    all_data.append(center_data)\n",
    "\n",
    "train_data = torch.cat(all_data, dim=0)\n",
    "train_data -= train_data.mean(0, keepdims=True)\n",
    "train_data /= train_data.std(0, keepdims=True)\n",
    "print(f'{train_data.shape[0]} samples of {train_data.shape[1]} dimensions in training data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CvTi5iiHohsy",
    "outputId": "b82a84fa-ed8f-452f-94ea-29c0e929003d"
   },
   "outputs": [],
   "source": [
    "# Evaluate to generate and use V-Shaped dataset\n",
    "\n",
    "num_samples_per_center = 4000\n",
    "stddev = 0.1\n",
    "\n",
    "def sample_along_line(p1, p2, n):\n",
    "    t = torch.rand(n, dtype=torch.float32)[:, None]\n",
    "    return p1[None] * t + p2[None] * (1-t)\n",
    "\n",
    "lines = torch.tensor([[[0, 0], [2, 3]], \n",
    "                      [[0, 0], [-2, 3]]])\n",
    "train_data = torch.cat([sample_along_line(*line, num_samples_per_center) \n",
    "                        for line in lines])\n",
    "train_data += torch.randn_like(train_data) * stddev\n",
    "train_data -= train_data.mean(0, keepdims=True)\n",
    "train_data /= train_data.std(0, keepdims=True)\n",
    "\n",
    "print(f'{train_data.shape[0]} samples of {train_data.shape[1]} dimensions in training data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IRKBnOsKFXQD",
    "outputId": "e6629659-767a-4d43-d8c8-d2c7b44be791"
   },
   "outputs": [],
   "source": [
    "# Evaluate to generate and use two moons dataset\n",
    "\n",
    "num_samples_per_center = 4000\n",
    "train_data, labels = sklearn.datasets.make_moons(8000)\n",
    "train_data = train_data[np.argsort(labels)] # sort by moon for plotting\n",
    "train_data = torch.from_numpy(train_data).float()\n",
    "train_data -= train_data.mean(0, keepdims=True)\n",
    "train_data /= train_data.std(0, keepdims=True)\n",
    "\n",
    "print(f'{train_data.shape[0]} samples of {train_data.shape[1]} dimensions in training data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kwjTwgC6Wq1h"
   },
   "source": [
    "The x's  in the plot below can be thought of as valid images in 2d space with the rest of the white region representing the rest of the invalid images. The blue clusters around the x's are also valid images (corresponding to minor pixel perturbations in original images)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 336
    },
    "id": "LJrXAdsLO6t8",
    "outputId": "829e433a-3610-422e-f94a-6034262d41bc"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "# Visualize the data\n",
    "plt.figure(figsize=[5,5])\n",
    "plt.scatter(train_data[:,0], train_data[:,1], alpha=0.1)\n",
    "plt.title('Original data distribution')\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O-RDRbzZsD7H"
   },
   "source": [
    "## Beta schedule\n",
    "\n",
    "Now that we have the original (non noisy) data, let's start now with the actual diffusion implementation. The first thing is to add noise to the input images following a fixed variance schedule (also known as beta schedule). The original paper uses a linear schedule. And 1000 timesteps to move forward and back. We use smaller number of timesteps (250) as the data is simpler in our case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xa-cIhVdWkii"
   },
   "outputs": [],
   "source": [
    "num_diffusion_timesteps=250\n",
    "beta_start = 0.0001\n",
    "beta_end = 0.02\n",
    "# schedule_type='linear'\n",
    "\n",
    "def get_beta_schedule(schedule_type, beta_start, beta_end, num_diffusion_timesteps):\n",
    "    if schedule_type == 'quadratic':\n",
    "        betas = np.linspace(beta_start ** 0.5, beta_end ** 0.5, num_diffusion_timesteps, dtype=np.float32) ** 2\n",
    "    elif schedule_type == 'linear':\n",
    "        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float32)\n",
    "    elif schedule_type == 'cosine':\n",
    "        # Implements eq. (17) and below from https://arxiv.org/pdf/2102.09672.pdf\n",
    "        t = torch.arange(num_diffusion_timesteps)\n",
    "        s = 0.008\n",
    "        f = torch.cos((t / num_diffusion_timesteps + s) / (1 + s) * np.pi/2) ** 2\n",
    "        alphas_cumprod = f / f[0]\n",
    "        betas = 1 - alphas_cumprod / torch.cat([torch.ones(1), alphas_cumprod[:-1]])\n",
    "        betas = betas.clip(max=0.999).numpy()\n",
    "    return betas\n",
    "\n",
    "betas_linear = get_beta_schedule('linear', beta_start, beta_end, num_diffusion_timesteps)\n",
    "betas_quad = get_beta_schedule('quadratic', beta_start, beta_end, num_diffusion_timesteps)\n",
    "betas_cosine = get_beta_schedule('cosine', beta_start, beta_end, num_diffusion_timesteps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sDIz3B0mBEiG"
   },
   "source": [
    "### Visualize beta schedules\n",
    "\n",
    "The below plot shows that the variance of noise is low at the start and increases as we move forward in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 312
    },
    "id": "60_TTZzFA_uq",
    "outputId": "b0f2cf30-a284-4dc2-febe-6f2ca7019d18"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "plt.plot(betas_linear, label = 'linear')\n",
    "plt.plot(betas_quad, label='quad')\n",
    "plt.plot(betas_cosine, label='cosine')\n",
    "plt.title('Beta schedule')\n",
    "plt.ylabel('Beta value')\n",
    "plt.xlabel('Timestep')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qimwKYzhNM6w"
   },
   "source": [
    "### Beta derivatives\n",
    "\n",
    "Next, let's compute all the derivatives from beta that are used repeatedly in the forward and reverse process of diffusion. Since the variance schedule ($\\beta_t$) is fixed, the derivatives of $\\beta_t$ are also fixed. We precompute these to save time/ compute.\n",
    "\n",
    "We'll see the use cases of these variables in the respective sections below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4WsolSX3NNS7"
   },
   "outputs": [],
   "source": [
    "class BetaDerivatives():\n",
    "    def __init__(self, betas, dtype=torch.float32):\n",
    "        \"\"\"Take in betas and pre-compute the dependent values to use in forward/ backward pass.\n",
    "\n",
    "        Values are precomputed for all timesteps so that they can be used as and\n",
    "        when required.\n",
    "        \"\"\"\n",
    "        self.np_betas = betas\n",
    "        timesteps, = betas.shape\n",
    "        self.num_timesteps = int(timesteps)\n",
    "\n",
    "        self.betas = torch.from_numpy(betas).to(dtype=dtype)\n",
    "        self.alphas = 1 - self.betas\n",
    "        self.alphas_cumprod = self.alphas.cumprod(dim=0)\n",
    "        self.alphas_cumprod_prev = torch.cat([torch.ones(1), self.alphas_cumprod[:-1]], dim=0)\n",
    "\n",
    "        # calculations required for diffusion q(x_t | x_{t-1}) and others\n",
    "        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)\n",
    "        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)\n",
    "        self.log_one_minus_alphas_cumprod = torch.log(1. - self.alphas_cumprod)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XbuYzr3dOCQK"
   },
   "outputs": [],
   "source": [
    "gdb = BetaDerivatives(betas_linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6HQ5X_w-AT-j"
   },
   "source": [
    "### Visualize beta derivatives over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 276
    },
    "id": "fSp_uO7TN88L",
    "outputId": "95bb859f-5aa2-44f4-a183-58ef84e83e54"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "# Visualizing betas and other variables\n",
    "plt.figure(figsize=(16, 9))\n",
    "\n",
    "plt.subplot(2,4,1)\n",
    "plt.plot(gdb.betas)\n",
    "plt.title('Betas')\n",
    "plt.subplot(2,4,2)\n",
    "plt.plot(gdb.alphas)\n",
    "plt.title('Alphas')\n",
    "\n",
    "plt.subplot(2,4,3)\n",
    "plt.plot(gdb.alphas_cumprod, label='alphas_cumprod')\n",
    "plt.plot(gdb.sqrt_alphas_cumprod, label='sqrt_alphas_cumprod')\n",
    "plt.legend();\n",
    "plt.subplot(2,4,4)\n",
    "plt.plot(1-gdb.alphas_cumprod, label='one_minus_alphas_cumprod')\n",
    "plt.plot(gdb.sqrt_one_minus_alphas_cumprod, label='sqrt_one_minus_alphas_cumprod')\n",
    "plt.plot(gdb.log_one_minus_alphas_cumprod, label='log_one_minus_alphas_cumprod')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PadYxkN7_QPs"
   },
   "source": [
    "## Forward pass of diffusion model\n",
    "\n",
    "In the forward pass, the diffused input at timestep t can be computed directly using the closed form equation (For derivation of how we arrive at this, refer to the paper).\n",
    "\n",
    "$q(x_t| x_0) = N(\\sqrt{\\bar{\\alpha_t}}x_o, 1-\\bar{\\alpha_t}I)$\n",
    "\n",
    "This is done in the q_sample function below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZpJM9cUKQ1XU"
   },
   "outputs": [],
   "source": [
    "class DiffusionForward(BetaDerivatives):\n",
    "    \"\"\"\n",
    "    Forward pass of the diffusion model.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, betas):\n",
    "        super().__init__(betas)\n",
    "\n",
    "\n",
    "    def q_sample(self, x_start, t, noise=None):\n",
    "        \"\"\"\n",
    "        Forward pass - sample of diffused data at time t.\n",
    "        \"\"\"\n",
    "        if noise is None:\n",
    "          noise = torch.randn(x_start.shape)\n",
    "        p1 = self.sqrt_alphas_cumprod[t] * x_start \n",
    "        p2 = self.sqrt_one_minus_alphas_cumprod[t] * noise \n",
    "        return (p1 + p2)\n",
    "\n",
    "    def q_sample_step(self, x_t, t, noise=None):\n",
    "        \"\"\"\n",
    "        Single step of forward diffusion, samples from p(x_(t+1)|x_t)\n",
    "        \"\"\"\n",
    "        noise = torch.randn(x_t.shape) if noise is None else noise\n",
    "        p1 = torch.sqrt(1 - self.betas[t]) * x_t\n",
    "        p2 = torch.sqrt(self.betas[t]) * noise\n",
    "        return p1 + p2\n",
    "\n",
    "diff_forward = DiffusionForward(betas_linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PFxZ2UhF_beI"
   },
   "source": [
    "### Visualize the forward diffusion of the entire data over time\n",
    "\n",
    "We start with original data distribution and move it through the forward diffusion process 10 steps at a time. We can see that the original data distribution information is lost till it resembles gaussian after num_diffusion_steps. \n",
    "\n",
    "Also, the slow perturbations at the start and large ones towards the end as per the beta schedule are evident from the video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 319,
     "referenced_widgets": [
      "6829284938ab4781b7e48f8de97ac8f1",
      "3026258fdfec4144ae6f08030f6bc04b",
      "0917bebe7cfc4f5d98d12881e3b3933f",
      "b7e519a141a14cd9a943462521fcf9c1",
      "4e6b045df999460284dce897ba448baa",
      "cdbd9bffdd7f41359d8147ff40a8100f",
      "1635da91dc084070908e20d026ab7650",
      "32e047d4fde9402aa84b7e88c1227a1e",
      "35ea6f0085e744b5a35dc9142acd41fb",
      "8070a6c587cf4fcd9ef4a2980e2e2bcc",
      "3be102fe006848c7855c08275b633a12"
     ]
    },
    "id": "b7SVU1wNfl15",
    "outputId": "4024a1c6-1ec0-45ce-b1bb-84b4cb37dff8"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "camera = Camera(plt.figure())\n",
    "stepsize = 10\n",
    "x0 = train_data[:]\n",
    "for timestep in tqdm(range(0, num_diffusion_timesteps, stepsize)): \n",
    "    tstep = torch.ones_like(x0[0], dtype=torch.long) * timestep\n",
    "    shifted = diff_forward.q_sample(x0, tstep)\n",
    "    order = np.arange(len(x0)).astype(np.int32)\n",
    "    np.random.shuffle(order)\n",
    "    plt.scatter(shifted[order,0], shifted[order,1], \n",
    "              marker='.',\n",
    "              c=np.arange(x0.shape[0])[order]//num_samples_per_center, \n",
    "              alpha=0.1)\n",
    "    plt.gca().annotate(\n",
    "        f'Time step: {timestep}', xy=(5, 5), xycoords='axes points',\n",
    "        size=14, ha='left', va='bottom'\n",
    "    )\n",
    "    camera.snap()\n",
    "\n",
    "save_animation('scatter.mp4', 30 * stepsize)\n",
    "show_video('scatter.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 319,
     "referenced_widgets": [
      "c0c5422d3bb94330bb576e8222ec1c3c",
      "03094c896e204a7b959db6d86ef9f8e8",
      "582182f289e942f8b3859264cbacaf80",
      "27bf3cff2edd424591efc0cfee071b8b",
      "1adbd2702cce4232a6f34465f01139cb",
      "a327559aad554b5eaa0fa57bd9f3424e",
      "f519d668c597464ca31fa089f6c4b8eb",
      "5a222302b33247ff944d3a4da44eea9d",
      "a9b82d9d151e499b81a1a9a35b605271",
      "7925e99529cc42f18147af540d30848a",
      "3add974db2c3404aaddc7e7b87f413cb"
     ]
    },
    "id": "DG58oGE2yY2h",
    "outputId": "b98b5dc7-db9a-484d-967c-71d3e4827134"
   },
   "outputs": [],
   "source": [
    "# same as above, but with single diffusion steps p(x_t|x_t-1), \n",
    "# instead of sampling from p(x_t|x_0) for every frame\n",
    "\n",
    "camera = Camera(plt.figure())\n",
    "\n",
    "xt = train_data[:]\n",
    "\n",
    "stepsize = 10\n",
    "\n",
    "# order for plotting the points\n",
    "order = np.arange(len(x0)).astype(np.int32)\n",
    "np.random.shuffle(order)\n",
    "\n",
    "for timestep in tqdm(range(0, num_diffusion_timesteps+1)): \n",
    "    if timestep > 0:\n",
    "        tstep = torch.ones_like(x0[0], dtype=torch.long) * (timestep - 1)\n",
    "        xt = diff_forward.q_sample_step(xt, tstep)\n",
    "    if timestep % 5 == 0:\n",
    "        plt.scatter(xt[order, 0], xt[order, 1], \n",
    "                    marker='.',\n",
    "                    c=np.arange(x0.shape[0])[order] // num_samples_per_center, \n",
    "                    alpha=0.1)\n",
    "        plt.gca().annotate(\n",
    "            f'Time step: {timestep}', xy=(5, 5), xycoords='axes points',\n",
    "                  size=14, ha='left', va='bottom')\n",
    "        camera.snap()\n",
    "\n",
    "save_animation('scatter.mp4', 150)\n",
    "show_video('scatter.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NJdCHKvW_BiZ"
   },
   "source": [
    "### Visualize the forward pass of single point\n",
    "\n",
    "We perform the forward diffusion of a single point over time. At every timestep, we generate 500 possible diffused samples of the same input point. We observe the distribution of these points over time. \n",
    "\n",
    "These too are closer to the original point at the start and move towards gaussian as the forward process reaches num_timesteps same as the example above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eL_o1p_ko2AZ"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "camera = Camera(plt.figure())\n",
    "\n",
    "x0 = torch.tensor([[0.0,4.0]]).repeat((500, 1))\n",
    "for timestep in range(0, num_diffusion_timesteps, 10): \n",
    "    tstep = torch.ones_like(x0[0], dtype=torch.long) * timestep\n",
    "    shifted = diff_forward.q_sample(x0, tstep)\n",
    "    plt.scatter(shifted[:, 0], shifted[:, 1], c='b')\n",
    "    plt.scatter(x0[0,0], x0[0,1], marker='x', c='r')\n",
    "    camera.snap()\n",
    "\n",
    "save_animation('pointshifting.mp4', 300)\n",
    "show_video('pointshifting.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NsSaG3wu28F4"
   },
   "source": [
    "## Model building\n",
    "\n",
    "With the data taken care of, let's build a model that can fit the data. We use a DNN with few layers since we're just using data with 2 features that we wish to reconstruct. Would be replaced with unet with similar loss function for the case of image data.\n",
    "\n",
    "The model takes in 2 inputs:\n",
    "* Timestep embedding of $t$\n",
    "* $x_t$\n",
    "\n",
    "And predicts \n",
    "* The noise $n$ that lead from $x_0$ to $x_t$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iThLbMiT-8Q7"
   },
   "source": [
    "### Timestep embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iqW5wDLJSJKr",
    "outputId": "9c1dfef1-41de-4d3e-9d68-ca1ee22dbe70"
   },
   "outputs": [],
   "source": [
    "# We create a 128 dimensional embedding for the timestep input to the model. \n",
    "# Fixed embeddings similar to positional embeddings in transformer are used - \n",
    "# could be replaced by trainable embeddings later\n",
    "emb_size = 128\n",
    "\n",
    "def get_timestep_embedding(timesteps, embedding_dim: int):\n",
    "    half_dim = embedding_dim // 2\n",
    "    emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)\n",
    "    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)\n",
    "    emb = timesteps[:, None] * emb[None, :]\n",
    "    emb = torch.cat([torch.sin(emb), torch.cos(emb)], axis=1)\n",
    "    if embedding_dim % 2 == 1:  # zero pad\n",
    "        emb = torch.pad(emb, [[0, 0], [0, 1]])\n",
    "    return emb\n",
    "\n",
    "# test on an example\n",
    "temb = get_timestep_embedding(torch.tensor([2, 3]), emb_size)\n",
    "print(temb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2nDm0rvOQ0VD"
   },
   "outputs": [],
   "source": [
    "# Actual model that takes in x_t and t and outputs n_{t-1}\n",
    "# Experiments showed that prediction of n_{t-1} worked better compared to\n",
    "# prediction of x_{t-1}\n",
    "from torchvision.ops import MLP\n",
    "from torch.optim import Adam\n",
    "\n",
    "class DiffusionModel(nn.Module):\n",
    "    def __init__(self, in_channels):\n",
    "        super().__init__()\n",
    "        self.x_emb = torch.nn.Linear(in_channels, emb_size)\n",
    "        self.mlp = MLP(\n",
    "            in_channels=emb_size * 2, \n",
    "            hidden_channels=[128, 64, 32, 16, 2], \n",
    "            activation_layer=torch.nn.ELU,\n",
    "        )\n",
    "\n",
    "    def forward(self, x, temb):\n",
    "        # inpute embedding\n",
    "        x = self.x_emb(x)\n",
    "        # concatenate timestep embedding\n",
    "        x = torch.cat([temb, x], dim=1)\n",
    "        # main MLP\n",
    "        x = self.mlp(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NlpJ_xTmUUFh"
   },
   "source": [
    "### Data generation for diffusion model\n",
    "\n",
    "Next, let's generate the data for the model to train. We generate $x_t$ given the input $x_0$ using the deterministic forward process equation described above. This $x_t$ and timestep embedding of \n",
    "$t$ are input to the model that is tasked with predicting the noise $n$.\n",
    "\n",
    "$t$ is picked uniformly between [0, num_diffusion_timesteps]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ssu-fnyeUXXm"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "batch_size = 128\n",
    "\n",
    "class DiffusionDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data.shape[0]\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        # randomly sample timestep and noise\n",
    "        tstep = torch.randint(0, num_diffusion_timesteps, (1,), dtype=torch.long)\n",
    "        noise = torch.randn(self.data.shape[1])\n",
    "        noisy_out = diff_forward.q_sample(self.data[i], tstep, noise)\n",
    "        return ((noisy_out, get_timestep_embedding(tstep, emb_size)[0]), noise) \n",
    "    \n",
    "\n",
    "class DiffusionDataloader(torch.utils.data.Dataset):\n",
    "    def __init__(self, data, batch_size):\n",
    "        self.data = data\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data.shape[0] // self.batch_size\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        # randomly sample timestep and noise\n",
    "        tstep = torch.randint(0, num_diffusion_timesteps, \n",
    "                              (self.batch_size,), dtype=torch.long)\n",
    "        noise = torch.randn(self.batch_size, self.data.shape[1])\n",
    "        # randomly sample subset of the data: \n",
    "        # this could be improved to not use duplicates during an epoch\n",
    "        ind = torch.randint(0, self.data.shape[0], \n",
    "                              (self.batch_size,), dtype=torch.long)\n",
    "        noisy_out = diff_forward.q_sample(self.data[ind], tstep[:, None], noise)\n",
    "        return ((noisy_out, get_timestep_embedding(tstep, emb_size)), noise) \n",
    "\n",
    "    def __iter__(self):\n",
    "        for i in range(len(self)):\n",
    "            yield self[i]\n",
    "\n",
    "            \n",
    "# dataset = DiffusionDataset(train_data)\n",
    "# dataloader = torch.utils.data.DataLoader(\n",
    "#     dataset, \n",
    "#     shuffle=True, \n",
    "#     batch_size=batch_size,\n",
    "#     drop_last=True,\n",
    "# )\n",
    "\n",
    "# Use the custom dataloader: it's much faster\n",
    "dataloader = DiffusionDataloader(train_data, batch_size=batch_size)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "V-lNH2zsUakB",
    "outputId": "f6cebb7d-c855-4942-dc4c-e3a0c5b476ba"
   },
   "outputs": [],
   "source": [
    "# Let's test the data generator\n",
    "(xx,tt),yy = next(iter(dataloader))\n",
    "print(xx.shape, tt.shape, yy.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ddr-DrYA-x2P"
   },
   "source": [
    "### Train and evaluate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 661,
     "referenced_widgets": [
      "c50a72362c804a6487caea81462509c7",
      "183c9b356fdc49e2816179eee1c6a145",
      "d1cc4c4c36574836bd4a478cf8287a65",
      "e0d4abd7fe5f48c280a50fd67f0079e6",
      "528eb55bf39d463483afc5100c25eb76",
      "d44add7a78eb402e946f78ad8497866a",
      "784ae90efe5a46c2971dcfc4c2452a1a",
      "3e699bdf928c4829b90360c2241fe09d",
      "5a7108db66cb439dbf3ce3af9e123921",
      "464f36ffb1c54c35ac9bfe6be49b9678",
      "c4b3e0b38d1343c39f76af76fdd90fdd"
     ]
    },
    "id": "w7v01phcpr7M",
    "outputId": "a5660091-6ff8-4b09-d177-ef4d4ceaf7ff"
   },
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "# device = torch.device('cuda:0')\n",
    "device = torch.device('cpu')\n",
    "\n",
    "model = DiffusionModel(in_channels=2)\n",
    "print(model)\n",
    "model = model.to(device)\n",
    "optimizer = Adam(model.parameters(), lr=1e-3)\n",
    "criterion = torch.nn.MSELoss().to(device)\n",
    "\n",
    "num_epochs = 500\n",
    "losses = []\n",
    "for epoch in tqdm(range(num_epochs), desc=f'Training for {num_epochs} epochs..'):\n",
    "    for iteration, ((x, temb), target) in enumerate(dataloader):\n",
    "        optimizer.zero_grad()\n",
    "        target = target.to(device)\n",
    "        x, temb = x.to(device), temb.to(device)\n",
    "        pred = model(x, temb)\n",
    "        loss = criterion(pred, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        losses.append(loss.item())\n",
    "\n",
    "    if epoch % 10 == 0:\n",
    "        print(f'Epoch {epoch}/{num_epochs}. loss: {np.mean(losses[-len(dataloader) * 10:]):.2e}')\n",
    "plt.plot(losses)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4pUhljaJTvpx"
   },
   "source": [
    "### Scatter plots of reconstructed values v/s target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 441
    },
    "id": "v8vKw7K4puSL",
    "outputId": "ad3f353d-62cb-46c4-be84-53b7ee272d39"
   },
   "outputs": [],
   "source": [
    "# TODO: Create scatterplots between predicted and ground truth noise. Color the points by the diffusion timestep.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E85QVROF-i87"
   },
   "source": [
    "## Reverse process of diffusion\n",
    "\n",
    "The model provides a decent estimate of the noise given the data and t.  Now comes the tricky part: given the data at timestep t $x_t$, and the noise estimate from the model, reconstructing original data distribution.\n",
    "\n",
    "\n",
    "There are 4 parts in the reverse process: \n",
    "1. Pass $x_t$ and $t$ (converted to time embedding) into the model that predicts the noise $ϵ$ \n",
    "2. Using the noise estimate $ϵ$ and $x_t$, compute $x_0$ using equation : $\\frac{1}{\\sqrt{\\bar{\\alpha}_t}}x_t - (\\sqrt{\\frac{1}{\\bar{\\alpha}_t}-1}) \\epsilon$\n",
    "\n",
    "\n",
    "3. Compute mean and variance using the equations: \n",
    "\n",
    "$\\tilde{\\mu}(x_t, x_0) = \\frac{\\sqrt{\\bar{\\alpha}_{t-1}} \\beta_t}{1-\\bar{\\alpha_t}}x_0 + \\frac{\\sqrt{\\bar{\\alpha}_t}(1-\\bar{\\alpha}_{t-1})}{1-\\bar{\\alpha}_t}x_t$ and\n",
    "variance $\\tilde{\\beta}_t = \\frac{(1-\\bar{\\alpha}_{t-1})}{1-\\bar{\\alpha}_t}\\beta_t$\n",
    "\n",
    "4. Sample using this mean and variance\n",
    "$q(x_{t−1}|x_t, x_0)=N(x_{t−1}; \\tilde{\\mu}(x_t, x_0), \\tilde{\\beta}_tI)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E8dw9P9P5q-M"
   },
   "outputs": [],
   "source": [
    "class DiffusionReconstruct(BetaDerivatives):\n",
    "  \n",
    "    def __init__(self, betas):\n",
    "        super().__init__(betas)\n",
    "\n",
    "        self.sqrt_recip_alphas_cumprod = torch.sqrt(1. / self.alphas_cumprod)\n",
    "        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1. / self.alphas_cumprod - 1)\n",
    "\n",
    "        # calculations required for posterior q(x_{t-1} | x_t, x_0)\n",
    "        # Variance choice corresponds to 2nd choice mentioned in the paper\n",
    "        self.posterior_variance = self.betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)  \n",
    "\n",
    "\n",
    "        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain\n",
    "        self.posterior_log_variance_clipped = np.log(np.maximum(self.posterior_variance, 1e-20))\n",
    "        self.posterior_mean_coef1 = self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)\n",
    "        self.posterior_mean_coef2 = (1. - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1. - self.alphas_cumprod)\n",
    "\n",
    "    def predict_start_from_noise(self, x_t, t, noise):\n",
    "        \"\"\"\n",
    "        Reconstruct x_0 using x_t, t and noise. Uses deterministic process\n",
    "        \"\"\"\n",
    "        return (\n",
    "            self.sqrt_recip_alphas_cumprod[t, None] * x_t -\n",
    "            self.sqrt_recipm1_alphas_cumprod[t, None] * noise\n",
    "        )\n",
    "\n",
    "    def q_posterior(self, x_start, x_t, t):\n",
    "        \"\"\"\n",
    "        Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0)\n",
    "        \"\"\"\n",
    "        posterior_mean = (\n",
    "            self.posterior_mean_coef1[t, None] * x_start +\n",
    "            self.posterior_mean_coef2[t, None] * x_t\n",
    "        )\n",
    "        posterior_log_variance_clipped = self.posterior_log_variance_clipped[t]\n",
    "        return posterior_mean, posterior_log_variance_clipped\n",
    "\n",
    "    def p_sample(self, model, x_t, t):\n",
    "        \"\"\"\n",
    "        Sample from the model. This does 4 things\n",
    "        * Predict the noise from the model using x_t and t\n",
    "        * Create estimate of x_0 using x_t and noise (reconstruction)\n",
    "        * Estimate of model mean and log_variance of x_{t-1} using x_0, x_t and t\n",
    "        * Sample data (for x_{t-1}) using the mean and variance values\n",
    "        \"\"\"\n",
    "        noise_pred = model(x_t, get_timestep_embedding(t, emb_size)) # Step 1\n",
    "        x_recon = self.predict_start_from_noise(x_t, t=t, noise=noise_pred) # Step 2\n",
    "        model_mean, model_log_variance = self.q_posterior(x_start=x_recon, x_t=x_t, t=t) # Step 3\n",
    "        noise = torch.randn(x_t.shape)\n",
    "        nonzero_mask = (t > 0).reshape(x_t.shape[0], 1) \n",
    "        return model_mean + torch.exp(0.5 * model_log_variance)[:, None] * noise * nonzero_mask # Step 4\n",
    "\n",
    "    \n",
    "    # TODO: implement for task (c)\n",
    "    def p_sample_samenoise(self, model, x_t, t):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    # TODO: implement for task (d)\n",
    "    def p_sample_DDIM(self, model, x_t, t, eta=0):\n",
    "        raise NotImplementedError\n",
    "\n",
    "\n",
    "    def p_sample_loop_trajectory(self, model, shape=None, latents=None, method=\"p_sample\"):\n",
    "        \"\"\"\n",
    "        Generate the visualization of intermediate steps of the reverse of diffusion\n",
    "        process.\n",
    "        \"\"\"\n",
    "        method = getattr(self, method)\n",
    "\n",
    "        t = self.num_timesteps - 1 \n",
    "        if latents is None:\n",
    "            latents = torch.randn(shape)\n",
    "        if shape is None:\n",
    "            shape = latents.shape\n",
    "        assert latents.shape == shape, f'Shape mismatch: {latents.shape}!={shape}'\n",
    "\n",
    "        imgs = latents[None]\n",
    "\n",
    "        with torch.no_grad():\n",
    "            while t >= 0:\n",
    "                imgs = torch.cat(\n",
    "                    [imgs, \n",
    "                     method(\n",
    "                         model=model, x_t=imgs[-1],\n",
    "                         t=torch.ones(shape[0], dtype=torch.long) * t)[None]\n",
    "                     ]\n",
    "                )\n",
    "                t -= 1\n",
    "\n",
    "        return imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: use rec_diff.p_sample_loop_trajectory with different methods to visualize the reverse process on\n",
    "#       points sampled from a gaussian and the linear interpolations as described on the sheet\n",
    "\n",
    "rec_diff = DiffusionReconstruct(betas_linear)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LxItzOc7FmyF"
   },
   "source": [
    "References:\n",
    "1. [Diffusion models repo](https://github.com/hojonathanho/diffusion)\n",
    "2. [Diffusion models paper](https://arxiv.org/pdf/2006.11239.pdf)\n",
    "3. [Improved Denoising Diffusion Probabilistic Models paper](https://arxiv.org/pdf/2102.09672.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2) Event Generation with a Normalizing Flow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implementing normalizing flows is tedious. In this exercise, we use the FrEIA package (https://github.com/vislearn/FrEIA) for this task. It has been developed in Heidelberg by the group of Prof. Köthe. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install FrEIA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code snippet demonstrates how you can train an INN on a toy dataset with FrEIA. For more information, check out the tutorial (https://vislearn.github.io/FrEIA/_build/html/tutorial/tutorial.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from sklearn.datasets import make_moons\n",
    "\n",
    "# FrEIA imports\n",
    "import FrEIA.framework as Ff\n",
    "import FrEIA.modules as Fm\n",
    "\n",
    "BATCHSIZE = 100\n",
    "N_DIM = 2\n",
    "\n",
    "# we define a subnet for use inside an affine coupling block\n",
    "# for more detailed information see the full tutorial\n",
    "def subnet_fc(dims_in, dims_out):\n",
    "    return nn.Sequential(nn.Linear(dims_in, 512), nn.ReLU(),\n",
    "                         nn.Linear(512,  dims_out))\n",
    "\n",
    "# a simple chain of operations is collected by ReversibleSequential\n",
    "inn = Ff.SequenceINN(N_DIM)\n",
    "for k in range(8):\n",
    "    inn.append(Fm.AllInOneBlock, subnet_constructor=subnet_fc, permute_soft=True)\n",
    "\n",
    "optimizer = torch.optim.Adam(inn.parameters(), lr=0.001)\n",
    "\n",
    "# a very basic training loop\n",
    "for i in range(500):\n",
    "    optimizer.zero_grad()\n",
    "    # sample data from the moons distribution\n",
    "    data, label = make_moons(n_samples=BATCHSIZE, noise=0.05)\n",
    "    x = torch.Tensor(data)\n",
    "    # pass to INN and get transformed variable z and log Jacobian determinant\n",
    "    z, log_jac_det = inn(x)\n",
    "    # calculate the negative log-likelihood of the model with a standard normal prior\n",
    "    loss = 0.5*torch.sum(z**2, 1) - log_jac_det\n",
    "    loss = loss.mean() / N_DIM\n",
    "    # backpropagate and update the weights\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# sample from the INN by sampling from a standard normal and transforming\n",
    "# it in the reverse direction\n",
    "z = torch.randn(BATCHSIZE, N_DIM)\n",
    "samples, _ = inn(z, rev=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_trn = np.load(\"data09/tutorial-11-data/dy_trn_data.npy\")\n",
    "data_tst = np.load(\"data09/tutorial-11-data/dy_tst_data.npy\")\n",
    "data_val = np.load(\"data09/tutorial-11-data/dy_val_data.npy\")\n",
    "print(data_trn.shape, data_tst.shape, data_val.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mass(particle):\n",
    "    return np.sqrt(np.clip(particle[:,0]**2 - np.sum(particle[:,1:]**2, axis=-1), 0, None))\n",
    "\n",
    "def get_pt(particle):\n",
    "    return np.sqrt(particle[:,1]**2 + particle[:,2]**2)\n",
    "\n",
    "def get_eta(particle):\n",
    "    p_absolute = np.sqrt(np.sum(particle[:,1:]**2, axis=-1))\n",
    "    return np.arctanh(particle[:,3] / p_absolute)\n",
    "\n",
    "def get_phi(particle):\n",
    "    return np.arctan2(particle[:,2], particle[:,1])\n",
    "\n",
    "def get_pt_phi_eta_mass(particle):\n",
    "    pt = get_pt(particle)\n",
    "    phi = get_phi(particle)\n",
    "    eta = get_eta(particle)\n",
    "    mass = get_mass(particle)\n",
    "    return np.stack((pt, phi, eta, mass), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cuts in the dataset\n",
    "get_pt(data_tst[:,:4]).min(), get_pt(data_tst[:,4:]).min() # pt cuts\n",
    "get_eta(data_tst[:,:4]).min(), get_eta(data_tst[:,:4]).max()\n",
    "\n",
    "eta_cut = 2.5 + 1e-5 # should be 2.5 but one event is weird\n",
    "pt_cut = 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(events, mean=None, std=None):\n",
    "    particle1, particle2 = events[:,:4], events[:,4:]\n",
    "    events_jetcoordinates = np.stack((get_pt(particle1), get_phi(particle1), get_eta(particle1), get_mass(particle1), get_pt(particle2), get_phi(particle2), get_eta(particle2), get_mass(particle2)), axis=-1)\n",
    "    events_reduced = events_jetcoordinates[:,[0,2,6]]\n",
    "    \n",
    "    events_reduced[:,0] = np.log(events_reduced[:,0] - pt_cut)\n",
    "    events_reduced[:,1:] = np.arctanh(events_reduced[:,1:] / eta_cut)\n",
    "    \n",
    "    if mean is None or std is None:\n",
    "        mean = events_reduced.mean(axis=0)\n",
    "        std = events_reduced.std(axis=0)\n",
    "    events_reduced = (events_reduced - mean) / std\n",
    "    \n",
    "    assert np.isfinite(events_reduced).all()\n",
    "    events_reduced = torch.tensor(events_reduced).float()\n",
    "    return events_reduced, mean, std\n",
    "    \n",
    "def undo_preprocess(events_reduced, mean, std):\n",
    "    events_reduced = events_reduced * std + mean\n",
    "    \n",
    "    events_reduced[:,0] = np.exp(events_reduced[:,0]) + pt_cut\n",
    "    events_reduced[:,[1,2]] = np.tanh(events_reduced[:,[1,2]]) * eta_cut\n",
    "    \n",
    "    pt1, eta1, eta2 = events_reduced.T\n",
    "    phi1 = np.random.uniform(0, 2*np.pi, events_reduced.shape[0])\n",
    "    mass1, mass2 = np.ones((2, events_reduced.shape[0])) * 0.105\n",
    "    px1 = pt1 * np.cos(phi1)\n",
    "    py1 = pt1 * np.sin(phi1)\n",
    "    pz1 = pt1 * np.sinh(eta1)\n",
    "    e1 = np.sqrt(mass1**2 + px1**2 + py1**2 + pz1**2)\n",
    "    px2 = -px1\n",
    "    py2 = -py1\n",
    "    pt2 = pt1\n",
    "    pz2 = pt2 * np.sinh(eta2)\n",
    "    e2 = np.sqrt(mass2**2 + px2**2 + py2**2 + pz2**2)\n",
    "    return np.stack((e1, px1, py1, pz1, e2, px2, py2, pz2), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "data_trn_prepd, mean, std = preprocess(data_trn)\n",
    "data_tst_prepd, _, _ = preprocess(data_tst, mean, std)\n",
    "data_val_prepd, _, _ = preprocess(data_val, mean, std)\n",
    "\n",
    "dataset_trn = TensorDataset((data_trn_prepd))\n",
    "dataset_tst = TensorDataset((data_tst_prepd))\n",
    "dataset_val = TensorDataset((data_val_prepd))\n",
    "\n",
    "batchsize = 1024\n",
    "dataloader_trn = DataLoader(dataset_trn, batch_size=batchsize, shuffle=True)\n",
    "dataloader_tst = DataLoader(dataset_tst, batch_size=batchsize, shuffle=False)\n",
    "dataloader_val = DataLoader(dataset_val, batch_size=batchsize, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Contruct event generation INN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Train event generation INN\n",
    "# Hint: INNs train much slower compared to simple MLPs due to the more sophisticated architecture. \n",
    "# Training for 1-2 epochs should be enough"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Generate events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "components_Eppp = [\"E [GeV]\", \"px [GeV]\", \"py [GeV]\", \"pz [GeV]\"]\n",
    "components_jetcoordinates = [\"pt [GeV]\", \"phi\", \"eta\", \"mass [GeV]\"]\n",
    "def plot(truth, generated, bins=50):\n",
    "    fig, axs = plt.subplots(4,4, figsize=(15,15))\n",
    "    \n",
    "    # plot (E, px, py, pz) for both particles\n",
    "    for iparticle in range(2):\n",
    "        for icomponent in range(4):\n",
    "            ax = axs[iparticle, icomponent]\n",
    "            xlabel = f\"{components_Eppp[icomponent]} of particle {iparticle+1}\"\n",
    "            i = iparticle*4 + icomponent\n",
    "            \n",
    "            bins_local = bins\n",
    "            _, bins_local, _ = ax.hist(truth[:,i], bins=bins_local, alpha=.5, label=\"truth\", density=True)\n",
    "            ax.hist(generated[:,i], bins=bins_local, alpha=.5, label=\"model\", density=True)\n",
    "            ax.legend()\n",
    "            ax.set_xlabel(xlabel)\n",
    "    \n",
    "    # plot (pt phi eta mass) for both particles\n",
    "    truth = np.concatenate((get_pt_phi_eta_mass(truth[:,:4]), \n",
    "                            get_pt_phi_eta_mass(truth[:,4:])), axis=-1)\n",
    "    generated = np.concatenate((get_pt_phi_eta_mass(generated[:,:4]), \n",
    "                                get_pt_phi_eta_mass(generated[:,4:])), axis=-1)\n",
    "    for iparticle in range(2):\n",
    "        for icomponent in range(4):\n",
    "            ax = axs[2+iparticle, icomponent]\n",
    "            xlabel = f\"{components_jetcoordinates[icomponent]} of particle {iparticle+1}\"\n",
    "            i = iparticle*4 + icomponent\n",
    "            \n",
    "            bins_local = bins\n",
    "            _, bins_local, _ = ax.hist(truth[:,i], bins=bins_local, alpha=.5, label=\"truth\", density=True)\n",
    "            ax.hist(generated[:,i], bins=bins_local, alpha=.5, label=\"model\", density=True)\n",
    "            ax.legend()\n",
    "            ax.set_xlabel(xlabel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize distributions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Compute likelihoods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Plot distributions of the 10% events with the largest likelihood"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "00125859113d4d3fa8c76f31bb1000a6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_decc2aa61b644bab8e2c34fe3d42a63a",
       "IPY_MODEL_ba4b9067f20f48d6b3bdd181197bd164",
       "IPY_MODEL_bd65c7a0f71b452e8d6c71f62f748c82"
      ],
      "layout": "IPY_MODEL_33ade5f8103546349b49915d56451271"
     }
    },
    "03094c896e204a7b959db6d86ef9f8e8": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a327559aad554b5eaa0fa57bd9f3424e",
      "placeholder": "​",
      "style": "IPY_MODEL_f519d668c597464ca31fa089f6c4b8eb",
      "value": "100%"
     }
    },
    "0917bebe7cfc4f5d98d12881e3b3933f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_32e047d4fde9402aa84b7e88c1227a1e",
      "max": 25,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_35ea6f0085e744b5a35dc9142acd41fb",
      "value": 25
     }
    },
    "1635da91dc084070908e20d026ab7650": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "183c9b356fdc49e2816179eee1c6a145": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d44add7a78eb402e946f78ad8497866a",
      "placeholder": "​",
      "style": "IPY_MODEL_784ae90efe5a46c2971dcfc4c2452a1a",
      "value": "Training for 20 epochs..: 100%"
     }
    },
    "1adbd2702cce4232a6f34465f01139cb": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "27bf3cff2edd424591efc0cfee071b8b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7925e99529cc42f18147af540d30848a",
      "placeholder": "​",
      "style": "IPY_MODEL_3add974db2c3404aaddc7e7b87f413cb",
      "value": " 251/251 [00:03&lt;00:00, 73.92it/s]"
     }
    },
    "3026258fdfec4144ae6f08030f6bc04b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_cdbd9bffdd7f41359d8147ff40a8100f",
      "placeholder": "​",
      "style": "IPY_MODEL_1635da91dc084070908e20d026ab7650",
      "value": "100%"
     }
    },
    "32e047d4fde9402aa84b7e88c1227a1e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "33ade5f8103546349b49915d56451271": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "35ea6f0085e744b5a35dc9142acd41fb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "3add974db2c3404aaddc7e7b87f413cb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "3be102fe006848c7855c08275b633a12": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "3e699bdf928c4829b90360c2241fe09d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "464f36ffb1c54c35ac9bfe6be49b9678": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4e6b045df999460284dce897ba448baa": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "528eb55bf39d463483afc5100c25eb76": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "56c09020b61c43c58a3f2538faf3c07a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "582182f289e942f8b3859264cbacaf80": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5a222302b33247ff944d3a4da44eea9d",
      "max": 251,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_a9b82d9d151e499b81a1a9a35b605271",
      "value": 251
     }
    },
    "5a222302b33247ff944d3a4da44eea9d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5a7108db66cb439dbf3ce3af9e123921": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "5eaa58477f3f47adb5c22127a8040484": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6829284938ab4781b7e48f8de97ac8f1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_3026258fdfec4144ae6f08030f6bc04b",
       "IPY_MODEL_0917bebe7cfc4f5d98d12881e3b3933f",
       "IPY_MODEL_b7e519a141a14cd9a943462521fcf9c1"
      ],
      "layout": "IPY_MODEL_4e6b045df999460284dce897ba448baa"
     }
    },
    "784ae90efe5a46c2971dcfc4c2452a1a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "7925e99529cc42f18147af540d30848a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8070a6c587cf4fcd9ef4a2980e2e2bcc": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a327559aad554b5eaa0fa57bd9f3424e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a8a30dd2f61241af89081bc89f536092": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a9b82d9d151e499b81a1a9a35b605271": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "b7e519a141a14cd9a943462521fcf9c1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_8070a6c587cf4fcd9ef4a2980e2e2bcc",
      "placeholder": "​",
      "style": "IPY_MODEL_3be102fe006848c7855c08275b633a12",
      "value": " 25/25 [00:01&lt;00:00, 11.54it/s]"
     }
    },
    "ba4b9067f20f48d6b3bdd181197bd164": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5eaa58477f3f47adb5c22127a8040484",
      "max": 250,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_f360e679ff5c4f60bf24a1796334693c",
      "value": 250
     }
    },
    "bd65c7a0f71b452e8d6c71f62f748c82": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d7fca449e4e04f638ba0df804c970ad1",
      "placeholder": "​",
      "style": "IPY_MODEL_56c09020b61c43c58a3f2538faf3c07a",
      "value": " 250/250 [00:19&lt;00:00, 14.40it/s]"
     }
    },
    "c0c5422d3bb94330bb576e8222ec1c3c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_03094c896e204a7b959db6d86ef9f8e8",
       "IPY_MODEL_582182f289e942f8b3859264cbacaf80",
       "IPY_MODEL_27bf3cff2edd424591efc0cfee071b8b"
      ],
      "layout": "IPY_MODEL_1adbd2702cce4232a6f34465f01139cb"
     }
    },
    "c4b3e0b38d1343c39f76af76fdd90fdd": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "c50a72362c804a6487caea81462509c7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_183c9b356fdc49e2816179eee1c6a145",
       "IPY_MODEL_d1cc4c4c36574836bd4a478cf8287a65",
       "IPY_MODEL_e0d4abd7fe5f48c280a50fd67f0079e6"
      ],
      "layout": "IPY_MODEL_528eb55bf39d463483afc5100c25eb76"
     }
    },
    "cdbd9bffdd7f41359d8147ff40a8100f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d1cc4c4c36574836bd4a478cf8287a65": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3e699bdf928c4829b90360c2241fe09d",
      "max": 20,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_5a7108db66cb439dbf3ce3af9e123921",
      "value": 20
     }
    },
    "d44add7a78eb402e946f78ad8497866a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d7fca449e4e04f638ba0df804c970ad1": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "decc2aa61b644bab8e2c34fe3d42a63a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a8a30dd2f61241af89081bc89f536092",
      "placeholder": "​",
      "style": "IPY_MODEL_f7973434b2704187ba214d80bf7b141a",
      "value": "plotting frames..: 100%"
     }
    },
    "e0d4abd7fe5f48c280a50fd67f0079e6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_464f36ffb1c54c35ac9bfe6be49b9678",
      "placeholder": "​",
      "style": "IPY_MODEL_c4b3e0b38d1343c39f76af76fdd90fdd",
      "value": " 20/20 [00:04&lt;00:00,  5.21it/s]"
     }
    },
    "f360e679ff5c4f60bf24a1796334693c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "f519d668c597464ca31fa089f6c4b8eb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "f7973434b2704187ba214d80bf7b141a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}