notebooks/alanine_dipeptide_basics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "disciplinary-trade",
   "metadata": {},
   "source": [
    "# Training a Boltzmann Generator for Alanine Dipeptide\n",
    "\n",
    "This notebook introduces basic concepts behind `bgflow`. \n",
    "\n",
    "It shows how to build an train a Boltzmann generator for a small peptide. The most important aspects it will cover are\n",
    "\n",
    "- retrieval of molecular training data\n",
    "- defining a internal coordinate transform\n",
    "- defining normalizing flow classes\n",
    "- combining different normalizing flows\n",
    "- training a Boltzmann generator via NLL and KLL\n",
    "\n",
    "The main purpose of this tutorial is to introduce the implementation. The network design is optimized for educational purposes rather than good performance. In the conlusions, we will discuss some aspects of the generator that are not ideal and outline improvements.\n",
    "\n",
    "## Some Preliminaries\n",
    "\n",
    "We instruct jupyter to reload any imports automatically and define the device and datatype, on which we want to perform the computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "revolutionary-counter",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "exact-fault",
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "continued-transformation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "dtype = torch.float32\n",
    "\n",
    "# a context tensor to send data to the right device and dtype via '.to(ctx)'\n",
    "ctx = torch.zeros([], device=device, dtype=dtype)\n",
    "\n",
    "# a brief check if this module is the main executable (or imported)\n",
    "main = (__name__ == \"__main__\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fabulous-roulette",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## Load the Data and the Molecular System\n",
    "\n",
    "Molecular trajectories and their corresponding potential energy functions are available from the `bgmol` repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "pending-coating",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using downloaded and verified file: /tmp/alanine-dipeptide-nowater.pdb\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from bgmol.datasets import Ala2TSF300\n",
    "\n",
    "is_data_here = os.path.isfile(\"Ala2TSF300.npy\")\n",
    "dataset = Ala2TSF300(download=(not is_data_here), read=True)\n",
    "system = dataset.system\n",
    "coordinates = dataset.coordinates\n",
    "temperature = dataset.temperature\n",
    "dim = dataset.dim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "likely-picture",
   "metadata": {},
   "source": [
    "The energy model is a `bgflow.Energy` that wraps around OpenMM. The `n_workers` argument determines the number of openmm contexts that are used for energy evaluations. In notebooks, we set `n_workers=1` to avoid hickups. In production, we can omit this argument so that `n_workers` is automatically set to the number of CPU cores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "occasional-seller",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_energy = dataset.get_energy_model(n_workers=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "explicit-cargo",
   "metadata": {},
   "source": [
    "### Visualize Data: Ramachandran Plot for the Backbone Angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "regulation-authority",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mdtraj as md \n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "def plot_phi_psi(ax, trajectory, system):\n",
    "    if not isinstance(trajectory, md.Trajectory):\n",
    "        trajectory = md.Trajectory(\n",
    "            xyz=trajectory.cpu().detach().numpy().reshape(-1, 22, 3), \n",
    "            topology=system.mdtraj_topology\n",
    "        )\n",
    "    phi, psi = system.compute_phi_psi(trajectory)\n",
    "    \n",
    "    ax.hist2d(phi, psi, 50, norm=LogNorm())\n",
    "    ax.set_xlim(-np.pi, np.pi)\n",
    "    ax.set_ylim(-np.pi, np.pi)\n",
    "    ax.set_xlabel(\"$\\phi$\")\n",
    "    _ = ax.set_ylabel(\"$\\psi$\")\n",
    "    \n",
    "    return trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "amateur-integration",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAANsAAADUCAYAAAAV1VaVAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAej0lEQVR4nO2da4wlR3XH/+f2fc971vt+2DHYa4zjGFgbsrEw2Bgvr0BAKAFCSEBCfEACBSkJMQpBiUXAygMSlMhKSIhEgBBwSAADhiQQTEz8iDFebMxi1t5d73pfMzt35s599smHO9vnVO303bu7c7vvnTk/aaXqqeru6r5bXf86deoUMTMMw+g/mbQrYBhrBWtshpEQ1tgMIyGssRlGQlhjM4yEsMZmGAmRTbsCRFQE8B0ABXTq8y/M/MFu5wQjI5ybnO6kyy0nbzRXj9IjgaSzCGOvF4KidIvd708LgZRjKUdwIZIplEDdyy+nJ1rqnHPyFlr5KN2oy08TVN2r5GYbUfqyq7cjbX6y95Bz3BqVuof6f5j3Mkj9dBmVvuLSjStYu/7z+GOHo3Rl4enjzLx+uXKpNzYAdQA3MvM8EeUAfJeI7mLme+NOyE1O4+J3/TYAYHzXMSfv+s1PROlrR38WpdcHc0453cAWwkKUnm2POOWOt8aidDWUxpCjtlsndTwa1KJ0ntyPQYPlle+vXeTk/e/xi6P0z/ZviNLTD7iNctO/74/Sd93/MaTNnufe6hwf2y3PtajajfcdQ/G4pMvH5AN1z+fft6L16zc3X39blP7mPR94Mq5c6o2NO7Pq80uHuaV/NtNurDpSb2wAQEQBgAcAPBvAJ5j5+8uUeSeAdwJAMDWFxlTnS7hhZN4p96zi0Sh9RV669+lM0ykXKEkzG4pUnG1XnHLrsnL9Wig9TNsb7o5kRLKOZRbl79RAHFtzM85xOZCyTVWnwxVXVo09f1vsNZPi5swbY/Nyz1sXpWttedHqkTrH6n/f6H55z35P+bW9t2HQ0M8/9+u/KBn3xJ8zEAYSZm4z8zUAtgG4joiuWqbMHcy8i5l3BaMjZ1zDMAadgWhsp2HmWQD/BWBPujUxjJUndRlJROsBNJl5lohKAF4G4CNdT8owwmJHRo7mXJk2nV2I0pNKOk5nXCNDQCJvisqIEWDRKycD95q6RpvjZeSkIyldK2igDDOTmaNOHsqSPDI9HqWf2TLmFDt6jfTsenCO7z3klLs7/Dz6Rbdraxm4OC3GEs+mBGUrQuXS0Sj9vc8NvoEku3VLlA7qvZkYUm9sADYD+NTSuC0D4J+Z+csp18kwVpzUGxszPwzgeWnXwzD6zUCN2QxjNZN6z3beZDo6WXt1+LTP8N9Qp8d8Z3LE3rGMuZpqoW2O3KkEPe7T1yiSe5+C8kjJBe4gZgvPRumd5Wei9A+ntjjlDm0oRemFrcUonZaN1p8GCJ5zeZSe/Km8p9lnuePm4qy8p5FDNQwri+t667OsZzOMhLDGZhgJMZwyMiRkFjpybK5ZcLKOKV/Gk23Jy8GVKYGSgTUlRSuh+0oqym+yphyHA9+xWR1q6VljVypm1JRD24v/koeUHQtkCmK6WHXKHcrJ9RcvkvqWXvJ8pIE/DaBN/+WHxUm5uvFip1x2UT2/N20x6Nx1QHxSX3pL95mq01jPZhgJYY3NMBJiKGUktYDC8c534vDcuJP3xLgsJVqfFafiBp9wyjmeIUoeVsKSU262LW4dvteIRss+fb1msOCUG/GW3Dj3UvfWy36KgXdOUeSmlpGtkufpmxLtRx+P0trTonTCldRBTX4DbcEcNtpFs0YaxkBhjc0wEmIoZaRmbsadyt03ITKyFKgJ1ULZKVdU1sgmi/w66a3UbirrZNhlklyvzj4ViPWwEhadctph2edEWyyphxuTUXq2UfRKSj30Mr38KXeiXU8299Mp2Z/U1vfSeYWLNzjlsk+KIzaPj2JYKT1dPXshWM9mGIlhjc0wEsIam2EkxFCO2YIGMHag433QmHKdW/dPTEfpvHL0Pdlwx2IjWRk7aWfmhbbrkbLYluu3VBCNfMY1x+vrTeZkGmAi6y78LGdksWvoTSWcaovp/8DiVJQ+Ou+NZ6rKxK+GkY0J913kkQzdxoPapN/2vEScN+hGwxsqwgce6amc9WyGkRDW2AwjIVKXkUS0HcA/AtiEjjvvHczcNfJo0GCM7+/ItoXNrll8flxM/PtI4l/MjLieIWN5kX3zzXjBVWvllv17wfPqKGTlWMdFmcq7ZuGSstXXPafnU02p489OiRyuVNy6ZyvyjSypMCbFo278FFISbqWnAbpdLy7MnfYmAYDWoacvuB7DROqNDR3p/j5mfpCIxgA8QER3M/OP0q6YYawkqctIZj7MzA8upSsAHgWwNd1aGcbKMwg9WwQRXYJO8J+uEZGLwRjyPzkCACjuvMQppy1ylYxIylrdlYPZrFgqs4E4xLba7veHlaVSrz7LeOETCnmRkfWiyDlfomaVA3TDCxFcqYskPj4jFsjMAVdGFlQg5eKs2sSjGh99uVfp2Kvc7Jan5aJe9+XLy7hyw4Z+F0TdQnEMCEQ0CuALAN7LzHN+vo6InM+UzryAYQw4A9HYlnav+QKATzPzF9Ouj2H0g9RlJHX63b8D8Cgz/1kv51x29fZoq6Sbbvywk1efECkW5kU6NkK3e28GSgbqrPht3JxPE3mRjhtKpjaaIg8nym44Bi1LKzV3Ar06L8d0XNIFdykeikpGlg/L9fU6MqD3NWLdNsmIKxfnbNzrOedy39XCIPRsvwTgrQBuJKKHlv69Mu1KGcZKk3rPxszfxZkbdBrGqmMQejbDWBOk3rOdD48/8ESk9zO7r3HyCqck3S5Kh9nwH1X1pWFWxm/kjdm0rzCr2B/wPEvCEfEMabVlzDYz7y5azaixXm3RnRbgilyzcFIqeNEjXvRlz1MkurbvoeGN4U7TbbHneeH9BsGMeM3QnGxy+Irt73HK+R4lqx3r2QwjIayxGUZCDKWMpFwO2Y0dCdLy1kjRlbujtIouh6Dg2mBYS8dWF/uMzqrJ62qPuWHZWJn7uSB5HLieJq2WlAtr7uvPzUre6fV6ANAuuN/EzJGTcj3lzPv1LnJwJRyRtezTkjC8bJNbcEqkc6YsUtn3cGGVl1S8lDSxns0wEsIam2EkxFDKSO1Bcsu1H3LyclWRX40x0YBn7OccKH3I2hoZLzfDgpKEbbccjYgjst4vw3dYbihLpbevBnISwBmjh8QCWTjouorGrQPzrX3auXclHJF1uLnjL1oXpbNedL62MrLmFuQhszXX1FveN3vO9RtmrGczjISwxmYYCTGUMlLz9fs+6By/+LW3R+n5zerxPINjmItzRPa0nUJPeIdlVxJlc6JTczqd9fSroj3jOiJrqdscEbnpllp5yRXrYHz9bU65cKPISC0d5zfHW3PrE5JXOuGWK5k10jCMfmCNzTASwhqbYSTE0I/Z/MWj4aQ484ZdQgKHeRmbZdRYib3BXTiivEEyaoogEz+2KxfFU2Ik73pNzNVkcWs173k9kxqnzYrp/2t73bHTShO3ENR3FJ65QfbErm5QTt5TTjE0x9Qmh1X9PXffbWGLeJp8577VOU7TWM9mGAlhjc0wEmIgZCQRfRLAqwEcZearLuRa1Y3LbzrRHPckm5JwbS0PPc8Q5FSoOJUVFN2IyOOjEgtER1vOBa7pv5RTGzTmvDqpw0yjWzCU/tEttkizdEmUrokDCeqb3PV2GeWI3ZqUl1YldxJj/Mm1tUB/UHq2fwCwJ+1KGEY/GYjGxszfAXDyrAUNY4gZCBnZCzoi8o4dO6K/hzn3e1GfkHRjQjkRl1xZlikpGaiiIDt/B5DJKsuaSvsOxloeahmZ92RkK1T3Ctw6qe3ZkKlLPfY891anXL+tk6fxQ+Fl2vLM2pG7MOF6Ik+UZSFhU0V9nvUk+szlYi7Wz5jU860UvYbkG4ierRd0ROT169ef/QTDGDCGprEZxrAzNDIyjuyCK/uyVZEmC9vVxLUnDwsltV4s71rTNDqCsd6AQ8tGAJgsiHSazEs6m4mXkTnPotlWW821xuQ58keSG852k0SFnS+M0mFenmPEe3/bxk9hOfZ5m5ZU16tNUK6c9osPDUO1sQYRfQbA/wDYSUQHiegdadfJMFaagejZmPlNadfBMPrNQPRshrEWGIie7UKY3+Huqd1Qpn8dwXhszN1NZqQgDsKlXPwmgnGMeueM5sT8XQrUeDDjjstqWRmnZD3T/2Jheedmf6PAfi607OZBkltQG0guuBs5akbVytKxnLz3+XHXM/zxyRHJ2yzvpVsslWHGejbDSAhrbIaREEMvI0efcuXhyeeoaLzF+FggYwU5T0vCrLezhjbd67xC4MrDUkakYymQ6/meJouhyKWiZzJfKMr1F9cv710BJBejw7+Prkdh5qIoXfU2CAmVB/h4Vt7zplLFKbd/XLyZWyW1maQfYXmVYD2bYSSENTbDSIihl5HaYfcMlILzHYLHlbPweE48PnzrYS5GOpYyrjWyHMj1AnXjthcKYKEla7r8kAkn1CYcla1i7fve5wbDMVfv2Z3bLTKyNeNahI9vFCvj5qJEc9ZWWgAol+Sd1Ytyju9cvlpYnU9lGAOINTbDSIihl5G1DSX3D0o6uk7ErozUlkUtb6ayVaeclpVaKpY9GTmWESkaKAtkm10Z2QzllU8V3O16D+jNOQIp528eEj7wSJROMnqwvtcNr/polJ6bcbc8Pl4VSXiiLNbhwLPMamm/IKec4VyuJ7mHeYLbejbDSAhrbIaRENbYDCMhhn7MVtrvLlSk68QkDTVearTcR82oMVtObR/je3zocdpkIOO5sYzruTIZLETpIsmYwzf9h+r79nRpwsn7yYjUvTku5nS9kBQAvjUAu7yUHz4UpUuXX+Lkzc7JOO3IyLick3VN/6G/tdASrRH3txr6/6RLWM9mGAlhjc0wEmIgemgi2gPgYwACAH/LzH/S87lz885x/pRIsZpac7XYcM3T1ZZIs8WspLXjLOB6g2i5ORa4ZvtpJTEnM/Hr40ZI8k6WR528H49viNL7p8aidH3SrfugkXVfBdpz8j6Pj4tNf7TghrxbrMtzsfKe8U3/w2zu15y1sS3FB3kIwMMA7mfmYytZASIKAHwCwM0ADgK4j4j+jZl/tJL3MYy06UVG/g2AKoA3ALibiO4korGznHMuXAdgHzM/wcwNAJ8F8NoVvL5hDAS9yMgdAL4N4K+ZuUVEbwHwIQC/vUJ12ArggDo+COCFfqHYiMib3BBoWtJQS6xdtborxWbq4nlSzqoQCYHrVKul44SSigHcdW/aAjmmjGxlcl/xCImUuqp4wMl7YlKCzz45LWu95re4XjJpeVTEhWN4wTv/3CkXVES+z9fE8Trw9rRrtaVcUJOX5oe6WC300rNdDuCPATxGRP8H4OUAXk9ENxHRSoQmXs7+e0YwDouIbAw7vfRsHz89TiOiKQC7ALwCwFsA/AmAay+wDgcBbFfH2wA8fYHXNIyBo5fG9o2lRvYogB8DuALAZ5j5Pd1P65n7AFxGRD8H4BCAXwPw5hW6tmEMDGdtbMz8vCWL4RUAdgL4JoCvrVQFlsaB7wbwdXRM/59k5r29nh8W4h9BjwOaNbfczIJ4OZSy8eHHtUfJWCDTAtpjBACaLIpce7eXM673R4GlHs/KzTp5V40cjNI/2Cj7WT+z0R2zVa/eGqX7GdbOJy7MXfENL3LK6R1uGmqsPOd55zQX5V3o7cWDmlsuyWfslfOpU0/zbMzcBrB36d+Kw8xfBfDVflzbMAYF8yAxjIQYCA+Sc+Unjz4tCyo9GdlWqi3TFDnDDW8HFRV+7VTBW4Cq0DLSWUiacb0htJPyZChxN0bJXbRaIJFV0xm37pfmj0bpKyYk/dSOi5xy809J3Tds3YK0KZ5wPWay82K6b14kv4E29QMAh8s7IhdmXVk/KNJRcz51sp7NMBLCGpthJMRQykg028gsbRBYU5Y5AFBGQajtnEGejGwr6+TJebFMNkO3XF2tg8vEOCUDQFFFRC6STruOzRMZ1/NEs05ZOLcUZ6N0ecL19K2tcy2cp/E3wuin/HIsk9e7ofZ0NEBW+2i3Gt5mHGq9oTL0ojk2nP8tz4b1bIaRENbYDCMhhr6/bpVdaaKtkZxVk6Mxli/AnVytwNvvrbn8K2pxb9+pNh91jjd5ofI0lVDWtxWULJ0acWXkkQ0SamDuReKUfc/n01n3pffXBoCc2j+jVpPfhz1HZOg93lRWed/sCtZucLCezTASwhqbYSSENTbDSIihH7PVx92xmPZ1JWV2Rs411bNaWAqSb05jruCUaxVlXBEqU7Ufhq2l5hmaLOla6C5aPdYW7xI9XQAAJ1oyZptpql1dvBDm7TF5lso2+Qn7bfqPc77NVjwPkqqMezNqXBYW3WmPjPoN8vJasHiJG+Jv0PH3AI/DejbDSAhrbIaREEMvIwtzrjm5Pqnlocprut8V7UfMeSlHTS+CsfI8WVAeELWaKw8Xx+W41pbX+kzRjY20Pi+h98qBK7/m2yK/nlyQ2Co6jgcAQNVJKVYEz7ncKbbS68B6vUZQV+9Tq/ea9xssLv+tz59y5fUgrmfT6DgwRB+PLWc9m2EkhDU2w0iIVGUkEb0RwB8CeA6A65j5/l7Ou+zq7bjr/k7XfdONH3by2gWRXI3JeK+RsKCkjpIznn8xaFFZIKtSrj3qfqdOKMvnqaqsjztSdmWkjgpczLqRfxebIkVPVsU5ujJbdsqVjsi9x58UC5/e87rf+JZPDV+zO0o7DjPkevtoFZ2ryu/h75P+9QGUjudD2j3bIwBeD+A7KdfDMPpOqj0bMz8KAETxPZBhrBaGxhqpIyIXUY5kTPYFVznlwqzIyKwKgJWtug1aR4AKu7wFvTZLOzmHdfek8JSa1FYS9Zi3t9rxrJrY9b4x3FJCQ00G50+68ktPAAcNuV7Gfxdq7+1u9Grti5OOWS80w+jTYk2srVObZ2TdB86rrfX0e+613sNG3xsbEX0TwKZlsm5l5i/1eh1mvgPAHQAwTtNnREw2jEGn742NmV/W73sYxjCQtoHEMNYMaZv+fwXAXwJYD+ArRPQQM99ytvMuf8GluPv+5ccWL73lI1G6nZfxUmPcK6jMzgUV/6LlRbUrzIhirU1r7xS3XG5e/rC4UYXQm3XHbK2SUsDepy5XUY65ajxTOu6q5tGDMn2Qe0Y8Us7X9N+rV0ZcuT3PvdU51gt6Ax3xr+IUc+KO5ObjY7OsFtK2Rt4J4M4062AYSWEy0jASYmhM/3H45mj9QIU3yoYPo4dc15DKdjFJBw2RaZP7XCfYxfVSbvKnco3qRtccr/ZCxMghSfuy1Flj5zHyjEipbFXuVTzqxiA5HcYPAFqHZHetXs32K+3M29zo7g1ObRXyT02/6DQAJ+5I6V+/v6J1GkSsZzOMhLDGZhgJMfQy0kd7M4w/InLLlzpjajtrLdNansdHYUb0Yaal9mp7yttTW11j7jK51/iTrlNtY1xeefmwGy05mBGvXW1Z9B1x4yRhkhGRdSiA8ErXg2TkgNaLEt6hNuV+29c9LObJ1eJs3A3r2QwjIayxGUZCEPPwuRnu2rWL77//7EvftKzynXSpKrPaXBbpqC19ANC6eMOy126NuAq8cFC8gytXSkiD8XufcsqFmyRP1wFwpWNa8vB88DfWaEyIBTe7KFZVf8MMHfmY5mRyXocZGDaI6AFm3rVcnvVshpEQ1tgMIyGssRlGQqw607+m27jHD/t2Gu2RAQCI8dDwo+DWlPn7ns+/L0r7Trp6YeSFOgD3g/NZSFp/zXVOXn1CvGty8/I9z7Tj7QPDPE7rFevZDCMhrLEZRkKsahnZjV7XfmmPFG3i9l9c9m6ZiugmxXxZqRmEyL+9Skcn7sgJbwojL5Gd20VxvM4sDN8000piPZthJIQ1NsNIiLTDItwO4DXoBCn4KYDfYubZJO6tZdAZFkhFnJVMh18AgKy6hh/aTfO1vbfF5g0a3aIea/w9tfV245wRGZmruE7ZWsoPgoTuN2n3bHcDuIqZrwbwOID3p1wfw+gbqTY2Zv4GM5/+3N0LYFua9TGMfjJI1si3A/hcXKaOiLxjx47Yi8RJnzMmsZXjq87TDrGAaz3UsqfgXY9iZGk3KTbocsmvn34W7aDtOxi3C8tHmw5z7rc9p97hMMnr82UgIiIT0a0AWgA+HXcdHRF5165da9uGbAwlqUdEJqK3AXg1gJt4GNf7GEaPpG2N3APgdwHcwMzVs5U3jGEm7THbXwEoALh7aduoe5n5XSt1cW2Cb3gxSDJTssGg3nyvPeUtFv3eQ1FSj+18DxS9ODUYl3v5Y8Bu0wyDPoZzpkvUewlf90KnXEuN2Vj9D8s03bgtetHuWiDtiMjPTvP+hpEkac+zGcaaIW0ZecH06uWQabgSZnGDbJo4cqC17N8BoLT7GjlQoebClzzfKReoDS7md05F6dEfu/XIIt67ZNDR3jR6SiSoe2H9ZkVGtkqSrk/mnHJ6UxDzIDEMY8WwxmYYCTH0MtL3DNFWwsq12+XvRXdDi5Y6bheV9dDdfwNzl8rOGMUTIoMaE+7GGqMNsW7q/br9SMzZVWKB0x4fL37t7U5edYN8w7UHSfGEKzf1u/nW3tUpHTXWsxlGQlhjM4yEGHoZ6Tuw7v7VP43SeiMHf5+0lqg+BPX4PdP0VrStgshIfX6nnFgxtSPu7LOKTrmi2vZ30CMdd0PXvfHmX4wtFyoD5MIm97/b1N/fF6V1tLLVGmnLejbDSAhrbIaRENbYDCMhhn7MdtONH3aOF66SMdLCdlmxE3pPylnJC2pqzBa647ec8iMOldWevWFec0ymAkI1K0DeoqFWWb5vtXfsdvKGddySr7jzJfUJedn6vTdH3ZfWeMW1Ufrur/xOfyo3QFjPZhgJYY3NMBJiKGXk448djqITn9rp2vTnLxbd1l4vkXqzRTeMWiajyrXlm9NuuJ4hzboca0noy8i6rxdPn+PJUi03S4c9r5aYjRcHET1Nccu1H3LyTl08HqXb+ufxXlFzVN7FMMdq6RXr2QwjIayxGUZCpB2D5I8AvBZACOAogN9k5vi4AafJUBSFd+5SV4rxlsUovW39bJReV4wPcVJtiZtDtek6Ctda8opqDSkXZFynWk02kLzQ05uL6+QalQlXAhPELeWGV300Sn97AC11XfcrV69GR0duu840CHPybnTIhWGyxALnsKayz/U4G7cz89XMfA2ALwP4g5TrYxh9I+2IyHPqcARnDKENY/WQujWSiG4D8BsATgF4aZdyUUTkQmEykboZxkqSekRkZr4VwK1E9H4A7wbwweWuoyMilzds5xNXdgYAja3uRnw7txyN0leMPxOlN+QrsXWstmWcNuO58y+0xJu/1nZjaGjyGZlaKASSznid9YxaZLq/NOXkHc2oY46/1yDQzfQf1CWtx2yBtzC3ul7GbGNq2mMYVkNobx9nAfOP4s9JPSKy4p8AfAUxjc0whp1Ux2xEdJk6/GUAj6VVF8PoN5RmeH0i+gKAneiY/p8E8C5mPtTDeceWyifNRQCOp3DffrManyutZ7qYmdcvl5FqYxs2iOh+Zt6Vdj1WmtX4XIP4TGnPsxnGmsEam2EkhDW2c+OOtCvQJ1bjcw3cM9mYzTASwno2w0gIa2yGkRDW2M4RIrqdiB4jooeJ6E4imky7TucLEe0hoh8T0T4i+r2063OhENF2IvpPInqUiPYS0XvOflZy2JjtHCGilwP4D2ZuEdFHAICZfzflap0zRBQAeBzAzQAOArgPwJuYuYt332BDRJsBbGbmB4loDMADAF43KM9kPds5wszfYObTnsb3AtiWZn0ugOsA7GPmJ5i5AeCz6CzkHVqY+TAzP7iUrgB4FMDWdGslWGO7MN4O4K60K3GebAVwQB0fxAD9x7xQiOgSAM8D8P2UqxKR+nq2QeRsy4KWytwKoAXg00nWbQVZbjeRVTGmIKJRAF8A8F5vgXKqWGNbhrMtCyKitwF4NYCbeHgHvQcBbFfH2wCcPf7LgENEOXQa2qeZ+Ytp10djBpJzhIj2APgzADcw87G063O+EFEWHQPJTQAOoWMgeTMz7021YhcAERGATwE4yczvTbk6Z2CN7Rwhon0ACgBOLP3pXmZ+V4pVOm+I6JUA/gJAAOCTzHxb9zMGGyK6HsB/A/ghOsu2AOD3mfmr6dVKsMZmGAlh1kjDSAhrbIaRENbYDCMhrLEZRkJYYzOMhLDGZhgJYY3NMBLCGpsBIno9EX087XqsdqyxGUDHO/7BtCux2jFH5DUMEV0O4BMAXgTgBBFNMvNfpFur1Yu5a61RiKiAzlqvtwL4EoDd6OzBsoWZa2nWbbViMnLtcjOAH6CzrGaOmY8AqKHjlGz0AWtsa5dr0PGO/wUADxPRBgAVZl5ItVarGBuzrV3mAFyNzmrzhwF8CJ3xm9EnbMy2RiGiaQB3Avh5ADPoBPz5wBCvPB94rLGtcYjoBwBuZOYTZy1sXBA2ZlvDLFkkx6yhJYP1bIaRENazGUZCWGMzjISwxmYYCWGNzTASwhqbYSSENTbDSAhrbIaREP8P0pXIjbYfuYMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 216x216 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "if main:\n",
    "    fig, ax = plt.subplots(figsize=(3,3))\n",
    "    _ = plot_phi_psi(ax, dataset.trajectory, system)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interior-halifax",
   "metadata": {},
   "source": [
    "## Split Data and Randomly Permute Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "spiritual-execution",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train = len(dataset)//2\n",
    "n_test = len(dataset) - n_train\n",
    "permutation = np.random.permutation(n_train)\n",
    "\n",
    "all_data = coordinates.reshape(-1, dataset.dim)\n",
    "training_data = torch.tensor(all_data[permutation]).to(ctx)\n",
    "test_data = torch.tensor(all_data[permutation + n_train]).to(ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "associate-battle",
   "metadata": {},
   "source": [
    "## Define the Internal Coordinate Transform\n",
    "\n",
    "Rather than generating all-Cartesian coordinates, we use a mixed internal coordinate transform.\n",
    "The five central alanine atoms will serve as a Cartesian \"anchor\", from which all other atoms are placed with respect to internal coordinates (IC) defined through a z-matrix. We have deposited a valid `z_matrix` and the corresponding `rigid_block` in the `dataset.system` from `bgmol`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "impressed-boost",
   "metadata": {},
   "outputs": [],
   "source": [
    "import bgflow as bg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "comic-structure",
   "metadata": {},
   "outputs": [],
   "source": [
    "# throw away 6 degrees of freedom (rotation and translation)\n",
    "dim_cartesian = len(system.rigid_block) * 3 - 6\n",
    "dim_bonds = len(system.z_matrix)\n",
    "dim_angles = dim_bonds\n",
    "dim_torsions = dim_bonds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "marine-investor",
   "metadata": {},
   "outputs": [],
   "source": [
    "coordinate_transform = bg.MixedCoordinateTransformation(\n",
    "    data=training_data, \n",
    "    z_matrix=system.z_matrix,\n",
    "    fixed_atoms=system.rigid_block,\n",
    "    keepdims=dim_cartesian, \n",
    "    normalize_angles=True,\n",
    ").to(ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-magnitude",
   "metadata": {},
   "source": [
    "For demonstration, we transform the first 3 samples from the training data set into internal coordinates as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eligible-conviction",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 17]),\n",
       " torch.Size([3, 17]),\n",
       " torch.Size([3, 17]),\n",
       " torch.Size([3, 9]),\n",
       " torch.Size([3, 1]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bonds, angles, torsions, cartesian, dlogp = coordinate_transform.forward(training_data[:3])\n",
    "bonds.shape, angles.shape, torsions.shape, cartesian.shape, dlogp.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "incomplete-equipment",
   "metadata": {},
   "source": [
    "## Prior Distribution\n",
    "\n",
    "The next step is to define a prior distribution that we can easily sample from. The normalizing flow will be trained to transform such latent samples into molecular coordinates. Here, we just take a normal distribution, which is a rather naive choice for reasons that will be discussed in other notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "revised-democracy",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim_ics = dim_bonds + dim_angles + dim_torsions + dim_cartesian\n",
    "mean = torch.zeros(dim_ics).to(ctx) \n",
    "# passing the mean explicitly to create samples on the correct device\n",
    "prior = bg.NormalDistribution(dim_ics, mean=mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dental-conspiracy",
   "metadata": {},
   "source": [
    "## Normalizing Flow\n",
    "\n",
    "Next, we set up the normalizing flow by stacking together different neural networks. For now, we will do this in a rather naive way, not distinguishing between bonds, angles, and torsions. Therefore, we will first define a flow that splits the output from the prior into the different IC terms.\n",
    "\n",
    "### Split Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "prepared-contrast",
   "metadata": {},
   "outputs": [],
   "source": [
    "split_into_ics_flow = bg.SplitFlow(dim_bonds, dim_angles, dim_torsions, dim_cartesian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "polish-spell",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 66])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test\n",
    "_ics = split_into_ics_flow(prior.sample(3))[:-1]\n",
    "coordinate_transform.forward(*_ics, inverse=True)[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "maritime-drawing",
   "metadata": {},
   "source": [
    "### Coupling Layers\n",
    "\n",
    "Next, we will set up so-called RealNVP coupling layers, which split the input into two channels and then learn affine transformations of channel 1 conditioned on channel 2. Here we will do the split naively between the first and second half of the degrees of freedom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "detailed-confidentiality",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RealNVP(bg.SequentialFlow):\n",
    "    \n",
    "    def __init__(self, dim, hidden):\n",
    "        self.dim = dim\n",
    "        self.hidden = hidden\n",
    "        super().__init__(self._create_layers())\n",
    "    \n",
    "    def _create_layers(self):\n",
    "        dim_channel1 =  self.dim//2\n",
    "        dim_channel2 = self.dim - dim_channel1\n",
    "        split_into_2 = bg.SplitFlow(dim_channel1, dim_channel2)\n",
    "        \n",
    "        layers = [\n",
    "            # -- split\n",
    "            split_into_2,\n",
    "            # --transform\n",
    "            self._coupling_block(dim_channel1, dim_channel2),\n",
    "            bg.SwapFlow(),\n",
    "            self._coupling_block(dim_channel2, dim_channel1),\n",
    "            # -- merge\n",
    "            bg.InverseFlow(split_into_2)\n",
    "        ]\n",
    "        return layers\n",
    "        \n",
    "    def _dense_net(self, dim1, dim2):\n",
    "        return bg.DenseNet(\n",
    "            [dim1, *self.hidden, dim2],\n",
    "            activation=torch.nn.ReLU()\n",
    "        )\n",
    "    \n",
    "    def _coupling_block(self, dim1, dim2):\n",
    "        return bg.CouplingFlow(bg.AffineTransformer(\n",
    "            shift_transformation=self._dense_net(dim1, dim2),\n",
    "            scale_transformation=self._dense_net(dim1, dim2)\n",
    "        ))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "quality-anthropology",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 60])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "RealNVP(dim_ics, hidden=[128]).to(ctx).forward(prior.sample(3))[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "essential-tackle",
   "metadata": {},
   "source": [
    "### Boltzmann Generator\n",
    "\n",
    "Finally, we define the Boltzmann generator.\n",
    "It will sample molecular conformations by \n",
    "\n",
    "1. sampling in latent space from the normal prior distribution,\n",
    "2. transforming the samples into a more complication distribution through a number of RealNVP blocks (the parameters of these blocks will be subject to optimization),\n",
    "3. splitting the output of the network into blocks that define the internal coordinates, and\n",
    "4. transforming the internal coordinates into Cartesian coordinates through the inverse IC transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "capital-logan",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_realnvp_blocks = 5\n",
    "layers = []\n",
    "\n",
    "for i in range(n_realnvp_blocks):\n",
    "    layers.append(RealNVP(dim_ics, hidden=[128, 128, 128]))\n",
    "layers.append(split_into_ics_flow)\n",
    "layers.append(bg.InverseFlow(coordinate_transform))\n",
    "\n",
    "flow = bg.SequentialFlow(layers).to(ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "touched-cloud",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 66])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test\n",
    "flow.forward(prior.sample(3))[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "careful-location",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('#Parameters:', 817250)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print number of trainable parameters\n",
    "\"#Parameters:\", np.sum([np.prod(p.size()) for p in flow.parameters()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "documentary-demand",
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = bg.BoltzmannGenerator(\n",
    "    flow=flow,\n",
    "    prior=prior,\n",
    "    target=target_energy\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "regulated-chapel",
   "metadata": {},
   "source": [
    "## Train\n",
    "\n",
    "Boltzmann generators can be trained in two ways:\n",
    "1. by matching the density of samples from the training data via the negative log likelihood (NLL), and\n",
    "2. by matching the target density via the backward Kullback-Leibler loss (KLL).\n",
    "\n",
    "NLL-based training is faster, as it does not require the computation of molecular target energies. Therefore, we will first train the generator solely by density estimation.\n",
    "\n",
    "### NLL Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "private-speed",
   "metadata": {},
   "outputs": [],
   "source": [
    "nll_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-3)\n",
    "nll_trainer = bg.KLTrainer(\n",
    "    generator, \n",
    "    optim=nll_optimizer,\n",
    "    train_energy=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "grand-journal",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\tNLL: -8.4482\t\n",
      "1001\tNLL: -184.9325\t\n",
      "2001\tNLL: -201.8967\t\n",
      "3001\tNLL: -216.2626\t\n",
      "4001\tNLL: -223.9596\t\n",
      "5001\tNLL: -210.8960\t\n",
      "6001\tNLL: -227.0771\t\n",
      "7001\tNLL: -226.8064\t\n",
      "8001\tNLL: -231.2289\t\n",
      "9001\tNLL: -234.1555\t\n",
      "10001\tNLL: -234.2661\t\n",
      "11001\tNLL: -236.9761\t\n",
      "12001\tNLL: -229.9557\t\n",
      "13001\tNLL: -148.2766\t\n",
      "14001\tNLL: -228.4613\t\n",
      "15001\tNLL: -235.0358\t\n",
      "16001\tNLL: -221.6690\t\n",
      "17001\tNLL: -228.1348\t\n",
      "18001\tNLL: -229.6512\t\n",
      "19001\tNLL: -229.8495\t\n"
     ]
    }
   ],
   "source": [
    "if main:\n",
    "    nll_trainer.train(\n",
    "        n_iter=20000, \n",
    "        data=training_data,\n",
    "        batchsize=128,\n",
    "        n_print=1000, \n",
    "        w_energy=0.0\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cloudy-priority",
   "metadata": {},
   "source": [
    "To see what the generator has learned so far, let's first create a bunch of samples and compare their backbone angles with the molecular dynamics data. Let's also plot their energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "similar-pharmacology",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_energies(ax, samples, target_energy, test_data):\n",
    "    sample_energies = target_energy.energy(samples).cpu().detach().numpy()\n",
    "    md_energies = target_energy.energy(test_data[:len(samples)]).cpu().detach().numpy()\n",
    "    cut = max(np.percentile(sample_energies, 80), 20)\n",
    "    \n",
    "    ax.set_xlabel(\"Energy   [$k_B T$]\")\n",
    "    # y-axis on the right\n",
    "    ax2 = plt.twinx(ax)\n",
    "    ax.get_yaxis().set_visible(False)\n",
    "    \n",
    "    ax2.hist(sample_energies, range=(-50, cut), bins=40, density=False, label=\"BG\")\n",
    "    ax2.hist(md_energies, range=(-50, cut), bins=40, density=False, label=\"MD\")\n",
    "    ax2.set_ylabel(f\"Count   [#Samples / {len(samples)}]\")\n",
    "    ax2.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "annual-judge",
   "metadata": {},
   "outputs": [],
   "source": [
    "if main:\n",
    "    \n",
    "    n_samples = 10000\n",
    "    samples = generator.sample(n_samples)\n",
    "\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(6,3))\n",
    "    fig.tight_layout()\n",
    "\n",
    "    plot_phi_psi(axes[0], samples, system)\n",
    "    plot_energies(axes[1], samples, target_energy, test_data)\n",
    "\n",
    "    del samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baking-executive",
   "metadata": {},
   "source": [
    "### Mixed Training\n",
    "\n",
    "The next step is \"mixed\" training with a combination of NLL and KLL. To retain some of the progress made in the NLL phase, we decrease the learning rate and increase the batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "living-oriental",
   "metadata": {},
   "outputs": [],
   "source": [
    "mixed_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-4)\n",
    "mixed_trainer = bg.KLTrainer(\n",
    "    generator, \n",
    "    optim=mixed_optimizer,\n",
    "    train_energy=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "patient-spirituality",
   "metadata": {},
   "source": [
    "Mixed training will be considerably slower. \n",
    "To speed it up, you can change the settings for the OpenMM energy when creating the energy model. For example, consider not passing `n_workers=1`.\n",
    "\n",
    "To avoid large potential energy gradients from singularities, the components of the KL gradient are constrained to (-100, 100). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "apart-limitation",
   "metadata": {},
   "outputs": [],
   "source": [
    "if main:\n",
    "    mixed_trainer.train(\n",
    "        n_iter=2000, \n",
    "        data=training_data,\n",
    "        batchsize=1000,\n",
    "        n_print=100, \n",
    "        w_energy=0.1,\n",
    "        w_likelihood=0.9,\n",
    "        clip_forces=20.0\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "through-wiring",
   "metadata": {},
   "source": [
    "Plot the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "municipal-taste",
   "metadata": {},
   "outputs": [],
   "source": [
    "if main:\n",
    "    \n",
    "    n_samples = 10000\n",
    "    samples = generator.sample(n_samples)\n",
    "\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(6,3))\n",
    "    fig.tight_layout()\n",
    "\n",
    "    plot_phi_psi(axes[0], samples, system)\n",
    "    plot_energies(axes[1], samples, target_energy, test_data)\n",
    "\n",
    "    del samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adjustable-redhead",
   "metadata": {},
   "source": [
    "With nglview installed, we can also take a look at the samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "substantial-cathedral",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import nglview as nv\n",
    "    widget = nv.show_mdtraj(trajectory)\n",
    "except:\n",
    "    widget = None\n",
    "widget"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cardiac-triple",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "This tutorial has introduced the most basic concepts and implementations underlying Boltzmann generators and `bgflow`. That said, the trained networks did not do a particularly good job in reproducing the molecular Boltzmann distribution. Specifically, they only modeled the major modes of the $\\phi$ angle and still produced many samples with unreasonably large energies. Let's look at a few shortcomings of the present architecture:\n",
    "\n",
    "### 1) Unconstrained Internal Coordinates\n",
    "Bonds, angles, and torsions must not take arbitrary values in principle. Bond lengths need to be positive, angles live in $[0,\\pi],$ and torsions are periodic in $[-\\pi, \\pi].$ Neither those bounds nor the periodicity of torsions distributions have been taken into account by the present Boltzmann generator. The layers of the normalizing flow should be build in a way that preserves these constraints on the ICs.\n",
    "\n",
    "### 2)  Arbitrary Coupling\n",
    "The input for the coupling layers was split into two channels rather arbitrarily (first vs. second half). A partial remedy is to define the conditioning in a physically informed manner. Another solution is to augment the base space by momenta, which can be done with augmented normalizing flows (see for instance the notebook on temperature-steering flows).\n",
    "\n",
    "### 3) RealNVP Layers\n",
    "Affine coupling layers are well-known to perform poorly in separating modes. This explains that the metastable region around $\\phi \\approx \\pi/2$ was not captured by the generator. Other architectures such as augmented flows or neural spline flows do a better job for complicated, multimodal distributions.\n",
    "\n",
    "### 4) Training\n",
    "The generators were only trained for relatively few iterations and performance may improve with longer training and better choices of the learning rate and hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "broad-bidder",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}