sheet11.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": [
    "# 2) Event Generation with Diffusion Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(c)\n",
    "\n",
    "The following code provides a minimal implementation of a Denoising Diffusion Probabilistic Model (DDPM). It uses the same notation as in the lecture notes and https://arxiv.org/pdf/2305.10475.pdf. We recommend reading it alongside Figures 1, 2 of the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DDPM(nn.Module):\n",
    "    \n",
    "    def __init__(self, \n",
    "                data_dim, T=100, \n",
    "                 n_layers=3, hidden_dim=64):\n",
    "        super().__init__()\n",
    "        # initialize diffusion parameters\n",
    "        self.T = T\n",
    "        self.data_dim = data_dim\n",
    "        self.beta = torch.linspace(1e-4, 2e-2, self.T)\n",
    "        self.beta_bar = 1-torch.cumprod(1-self.beta, dim=0)\n",
    "        self.sigma2 = torch.ones_like(self.beta)\n",
    "        self.sigma2[1:] = self.beta_bar[:-1]/self.beta_bar[1:] * self.beta[1:]\n",
    "        \n",
    "        # construct simple MLP\n",
    "        layers = []\n",
    "        # network gets xt (data_dim-dimensional) and T (1-dimensional) as inputs\n",
    "        layers.append(nn.Linear(self.data_dim+1, hidden_dim))\n",
    "        for _ in range(n_layers-1):\n",
    "            layers.append(nn.Linear(hidden_dim, hidden_dim))\n",
    "            layers.append(nn.ReLU())\n",
    "        layers.append(nn.Linear(hidden_dim, self.data_dim))\n",
    "        self.net = nn.Sequential(*layers)\n",
    "        \n",
    "    def batch_loss(self, x0):\n",
    "        t = torch.randint(low=1, high=self.T, size=(x0.shape[0], 1))\n",
    "        epsilon = torch.randn_like(x0)\n",
    "        \n",
    "        # TODO: Compute x_t(x_0, epsilon)\n",
    "        xt = torch.sqrt(1-self.beta_bar[t])*x0 + torch.sqrt(self.beta_bar[t])*epsilon\n",
    "        \n",
    "        epsilon_theta = self.net(torch.cat((xt, t), dim=1))\n",
    "        \n",
    "        # loss function (eqn 17 in the paper)\n",
    "        prefactor = 1/(2*self.sigma2[t]) * self.beta[t]**2 / (1-self.beta[t]) / self.beta_bar[t]\n",
    "        loss = prefactor.unsqueeze(-1) * (epsilon - epsilon_theta)**2 \n",
    "        return loss.mean()\n",
    "    \n",
    "    def sample(self, n_samples):\n",
    "        # x_t starts at x_T and is diffused over x_{T-1}, x_{T-2} etc to x_0\n",
    "        xt = torch.randn(n_samples, self.data_dim)\n",
    "        for t in reversed(range(self.T)):\n",
    "            z = torch.randn(n_samples, self.data_dim) if t>0 else 0.\n",
    "            epsilon_theta = self.net(torch.cat((xt, torch.full((xt.shape[0], 1),\n",
    "                                                             t)), dim=1))\n",
    "            \n",
    "            # TODO: Compute x_t (x_{t+1}, epsilon_theta, z)\n",
    "            xt = (xt - epsilon_theta * self.beta[t] / torch.sqrt(self.beta_bar[t])) / torch.sqrt(1-self.beta[t]) + torch.sqrt(self.sigma2[t]) * z\n",
    "        return xt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard imports\n",
    "from sklearn.datasets import make_moons\n",
    "\n",
    "BATCHSIZE = 100\n",
    "N_DIM = 2\n",
    "    \n",
    "model = DDPM(data_dim=N_DIM)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# a very basic training loop\n",
    "losses = []\n",
    "for i in range(10000):\n",
    "    optimizer.zero_grad()\n",
    "    # sample data from the moons distribution\n",
    "    data, label = make_moons(n_samples=BATCHSIZE, noise=0.05)\n",
    "    x0 = torch.Tensor(data)\n",
    "    # compute loss of DDPM\n",
    "    loss = model.batch_loss(x0)\n",
    "    # backpropagate and update the weights\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    losses.append(loss.item())\n",
    "\n",
    "# sample from the DDPM\n",
    "with torch.no_grad():\n",
    "    samples = model.sample(BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(losses)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    samples = model.sample(10000)\n",
    "truth, _ = make_moons(n_samples=10000, noise=0.05)\n",
    "plt.plot(truth[:,0], truth[:,1], \"k,\")\n",
    "plt.plot(samples[:,0], samples[:,1], \"r,\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download data to folder data09\n",
    "# this might take some time (50MB)\n",
    "# you can also do this manually (download + unpack zip)\n",
    "import os, sys\n",
    "import wget\n",
    "from zipfile import ZipFile\n",
    "from tqdm import tqdm\n",
    "\n",
    "DESTINATION = \"data09\"\n",
    "url = \"https://www.thphys.uni-heidelberg.de/~plehn/pics/\"\n",
    "filename = \"tutorial-11-data.zip\"\n",
    "url = url + filename\n",
    "\n",
    "os.makedirs(DESTINATION, exist_ok=True)\n",
    "os.chdir(DESTINATION)\n",
    "wget.download(url, filename)\n",
    "with ZipFile(filename, \"r\") as zip_ref:\n",
    "    for file in tqdm(iterable=zip_ref.namelist(), total=len(zip_ref.namelist())):\n",
    "        zip_ref.extract(member=file)\n",
    "os.chdir(\"..\")\n",
    "%ls data09"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 264,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1384081, 8) (296588, 8) (296588, 8)\n"
     ]
    }
   ],
   "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": 265,
   "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": 266,
   "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": 267,
   "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": [
    "# TODO: Construct and train DDPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Sample from DDPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plotting code\n",
    "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: Plot distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(e)\n",
    "\n",
    "This is a minimal implementation of Conditional Flow Matching (CFM). Again, it uses the same notation as in the lecture notes and https://arxiv.org/pdf/2305.10475.pdf. We recommend reading it alongside Figures 3 of the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use ode solver that can operate on gpu\n",
    "from torchdiffeq import odeint\n",
    "\n",
    "class CFM(nn.Module):\n",
    "    \n",
    "    def __init__(self, data_dim, n_layers=3, hidden_dim=64):\n",
    "        super().__init__()\n",
    "        self.data_dim = data_dim\n",
    "        \n",
    "        # construct simple MLP\n",
    "        layers = []\n",
    "        # network gets x_t (data_dim-dimensional) and t (1-dimensional) as inputs\n",
    "        layers.append(nn.Linear(self.data_dim+1, hidden_dim))\n",
    "        for _ in range(n_layers-1):\n",
    "            layers.append(nn.Linear(hidden_dim, hidden_dim))\n",
    "            layers.append(nn.ReLU())\n",
    "        layers.append(nn.Linear(hidden_dim, self.data_dim))\n",
    "        self.net = nn.Sequential(*layers)\n",
    "        \n",
    "    def batch_loss(self, x0):\n",
    "        t = torch.rand(x0.shape[0], 1)\n",
    "        epsilon = torch.randn_like(x0)\n",
    "        \n",
    "        # TODO: Define linear trajectory\n",
    "        x_t = (1-t)*x0 + t*epsilon\n",
    "        v_t = -x0 + epsilon\n",
    "        \n",
    "        # Predict v_theta with neural network\n",
    "        v_theta = self.net(torch.cat((x_t, t), dim=1))\n",
    "        \n",
    "        # MSE loss\n",
    "        loss = torch.mean( (v_theta - v_t)**2)\n",
    "        return loss\n",
    "    \n",
    "    def sample(self, n_samples):\n",
    "        # velocity function to be called by the ODE solver\n",
    "        def velocity(t, x_t):\n",
    "            t = t * torch.ones(x_t.shape[0], 1).float()\n",
    "            v_t = self.net(torch.cat((x_t, t), dim=1))\n",
    "            return v_t\n",
    "\n",
    "        # apply ODE solver starting from N(0,1) noise\n",
    "        epsilon = torch.randn(n_samples, self.data_dim)\n",
    "        x_t = odeint(velocity, epsilon, torch.tensor([1.,0.]))[-1]\n",
    "        return x_t         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard imports\n",
    "from sklearn.datasets import make_moons\n",
    "\n",
    "BATCHSIZE = 1000\n",
    "N_DIM = 2\n",
    "    \n",
    "model = CFM(data_dim=N_DIM)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# a very basic training loop\n",
    "losses = []\n",
    "for i in range(1000):\n",
    "    optimizer.zero_grad()\n",
    "    # sample data from the moons distribution\n",
    "    data, label = make_moons(n_samples=BATCHSIZE, noise=0.05)\n",
    "    x0 = torch.Tensor(data)\n",
    "    # compute loss of DDPM\n",
    "    loss = model.batch_loss(x0)\n",
    "    # backpropagate and update the weights\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    losses.append(loss.item())\n",
    "\n",
    "# sample from the CFM\n",
    "with torch.no_grad():\n",
    "    samples = model.sample(BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(losses)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    samples = model.sample(10000)\n",
    "truth, _ = make_moons(n_samples=10000, noise=0.05)\n",
    "plt.plot(truth[:,0], truth[:,1], \"k,\")\n",
    "plt.plot(samples[:,0], samples[:,1], \"r,\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(f) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Construct and train CFM "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Sample from CFM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Plot distributions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3) Unfolding with Normalizing Flows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this exercise you we have to extend the INN of sheet 10, exercise 2 to allow for extra, conditional inputs. The following code demonstrates how to implement a 1-dimensional random condition for INNs within the FrEIA framework. "
   ]
  },
  {
   "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",
    "              cond=0, cond_shape=(1,))\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",
    "    c = torch.randn(BATCHSIZE, 1)\n",
    "    # pass to INN and get transformed variable z and log Jacobian determinant\n",
    "    z, log_jac_det = inn(x, c=(c,))\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",
    "c = torch.randn(BATCHSIZE, 1)\n",
    "samples, _ = inn(z, c=(c,), rev=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the dataset (same as exercise 1)"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def toy_detector_sim(data):\n",
    "    split_frac = np.random.uniform(0,1,(len(data),1))\n",
    "    q1 = split_frac * data[:,:4]\n",
    "    q2 = (1.-split_frac) * data[:,:4]\n",
    "    q3 = data[:,4:]\n",
    "    q_all = np.concatenate((q1, q2, q3), axis=1)\n",
    "    q_noise = q_all * np.random.normal(0.9, 0.05, q_all.shape)\n",
    "    return q_noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Apply toy detector simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Preprocess detector-level events (standardization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Create dataloaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Construct and train conditional INN for unfolding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Unfold all events once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize the distributions using the plot(truth, model) function "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Unfold the first 10 events in the test set 1000 times each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Visualize the distribution of unfolded events of the energy of the first muon\n",
    "# For each of the 10 plots, show the distribution of unfolded events as well as the true event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}