Schrodinger_Tutorial_CNF.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "# An introduction to Continuous Normalizing Flows"
      ],
      "metadata": {
        "id": "KDG7DXhWZNQh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Through this tutorial I want to explain in an intuitive and clear way the very interesting idea that lies behind Continuous Normalizing Flows(CNF for short) alongside with the code."
      ],
      "metadata": {
        "id": "3Frfh1D1TZwt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First of all, **CNF** is a deep generative model (as **VAE** or **GANS**) with the goal to seek the probabability density function that best fits a set of samples. But how do you find a probability just from knowing the samples? What loss function will the model have? \n",
        "\n",
        "We are going to look for an answer to those questions.\n",
        "\n",
        "This idea is very close to a method in Mathematical Statistics, called Maximum Likelihood Estimation, for finding the parameters of a probability distribution by assuming this distribution is indeed a good approximation for the samples.\n",
        "As an example, imagine you have some points on the X-axis (Fig.1) and you might try to guess that they are sampled from a Gaussian Distribution, thus you will search for it's mean and variance (Fig.2). In this \n",
        "particular case the blue curve is the gaussian distribution that fits the best.\n",
        "\n"
      ],
      "metadata": {
        "id": "9N_Y7hFsafCi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Figures { vertical-output: true }\n",
        "%%html\n",
        "<style>\n",
        "figure {\n",
        "    display: inline-block;\n",
        "}\n",
        "\n",
        "figure figcaption1 {\n",
        "    display:inline-block;\n",
        "    position:relative;\n",
        "    top:20px;\n",
        "    right:650px;\n",
        "}\n",
        "\n",
        "figure figcaption2 {\n",
        "    display:inline-block;\n",
        "    position:relative;\n",
        "    top:20px;\n",
        "    right:400px;\n",
        "}\n",
        "\n",
        "</style>\n",
        "\n",
        "<figure>\n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1H5rQegk6JMcHrLPGdNojbLtVG7Nsaui2\" width = '400'/>\n",
        "<img src = 'https://drive.google.com/uc?export=view&id=1EjU6HoDLPOP3vYp29uYJZbKy86S7hdmY' width = '400'/>\n",
        "<figcaption1> Fig.1 - Samples </figcaption1>\n",
        "<figcaption2>  Fig.2 - Some normal distributions </figcaption2>\n",
        "</figure>"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 333
        },
        "cellView": "form",
        "id": "TO1A8hz0n9jo",
        "outputId": "1fc4485e-7cf8-418f-d47e-7d2f7a0dfcf8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<style>\n",
              "figure {\n",
              "    display: inline-block;\n",
              "}\n",
              "\n",
              "figure figcaption1 {\n",
              "    display:inline-block;\n",
              "    position:relative;\n",
              "    top:20px;\n",
              "    right:650px;\n",
              "}\n",
              "\n",
              "figure figcaption2 {\n",
              "    display:inline-block;\n",
              "    position:relative;\n",
              "    top:20px;\n",
              "    right:400px;\n",
              "}\n",
              "\n",
              "</style>\n",
              "\n",
              "<figure>\n",
              "<img src = \"https://drive.google.com/uc?export=view&id=1H5rQegk6JMcHrLPGdNojbLtVG7Nsaui2\" width = '400'/>\n",
              "<img src = 'https://drive.google.com/uc?export=view&id=1EjU6HoDLPOP3vYp29uYJZbKy86S7hdmY' width = '400'/>\n",
              "<figcaption1> Fig.1 - Samples </figcaption1>\n",
              "<figcaption2>  Fig.2 - Some normal distributions </figcaption2>\n",
              "</figure>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The structure of the flow\n",
        "\n",
        "Let's look at what happens to a PDF $p_{X} (x)$ if you change the random variable $X$ to $Z$ through an invertible function $f : X → Z$, $f(X) = Z$ **(1)**. This change of variable shouldn't transform the PDF to an arbitrary function, therefore the PDF must remain a PDF. Consequently, $|p_{X}(x)dx| = |p_{Z}(z)dz|$ **(2)**, in other words, the probability density in a very small region should remain unchanged, which implies if $p_{X}(x)$ is a true PDF such that $\\int_{X}p_{X}(x)dx=1$ then $\\int_{Z}p_{Z}(z)dz=1$. Subsequently, if we divide by $dz$ in (2) and use (1) to map $x$ to $z$ we end up with $p_{Z}(z) = p_{X}(f^{-1}(z))|\\frac{df^{-1}(z)}{dz}|$.\n",
        "\n",
        "Moreover, the equality becomes a sum when you apply log on both sides\n",
        "\\begin{align} \n",
        "\\log p_{Z}(z) = \\log p_{X}(f^{-1}(z)) + \\log |\\frac{df^{-1}(z)}{dz}|. \n",
        "\\end{align}\n",
        "\n",
        "The word flow comes from the fact that after composing K bijective function, the result is still a bijective function:\n",
        "\\begin{align}\n",
        "y=f_{k}(x_{k})∘f_{k - 1}(x_{k-1})∘...∘f_{2}(x_{2})∘f_{1}(x_{1})\n",
        "\\end{align}\n",
        "\n",
        "What if instead of K functions there is an infinite number of them?\n",
        "\n",
        "Now, the random variable $y$ will change smoothly through time according to a differential equation:\n",
        "\\begin{align}\n",
        "\\frac{dz}{dt} = f(z, t, θ) \\tag*{(3)}\n",
        "\\end{align} where θ are the parameters of a neural network. \n",
        "\n",
        "You can think of $z$ evolving as a fluid through space(Fig.1).\n",
        "\n",
        "<div>\n",
        "<center>\n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1mJli8yECu812dhz_Mta68BTTE-kF224H\">\n",
        "</div>\n",
        "<center>\n",
        "<figcaption> Fig.1 - Vector field\n",
        "</center>\n",
        "\n",
        "Also, the change in the log probability density will be represented by another differential equation:\n",
        "\\begin{align}\n",
        "\\frac{\\partial \\log p(z(t))}{\\partial t}=-Tr(\\frac{\\partial f}{\\partial z(t)}) \\tag* {(4)}\n",
        "\\end{align}\n",
        "\n",
        "$Tr$ is the trace of $\\frac{df}{dz(t)}$ (the Jacobian $J$). Because $z$ is a multidimensional vector, $f$ is a vector-valued function which means the derivative becomes a matrix $A_{ij}=\\frac{\\partial f_{i}}{\\partial z_{j}(t)}$. Hence, the trace is the sum of the elements on the main diagonal:\n",
        "\\begin{align}\n",
        "\\sum_{i}\\frac{\\partial f_{i}}{\\partial z_{i}(t)} \\tag*{(5)}\n",
        "\\end{align}\n",
        "\n",
        "By integrating equations $(3)$ and $(4)$ from $t_{0}$ to $t_{1}$, $t_{1} > t_{0}$:\n",
        "\\begin{align}\n",
        "&z(t_{1})=z(t_{0})+\\int_{t_{0}}^{t_{1}}f(z, t, θ)\\,dt  \\tag*{(6)} \\\\\n",
        "\\\\\n",
        "&\\log p(z(t_{1})) = \\log p(z(t_{0})) - \\int_{t_{0}}^{t_{1}}Tr(\\frac{\\partial f}{\\partial z(t)})\\, dt \\tag*{(7)}\n",
        "\\end{align}\n",
        "\n",
        "I will rearrange a little bit equation $(7)$\n",
        "\\begin{align}\n",
        "\\log p(z(t_{0})) = \\log p(z(t_{1})) + \\int_{t_{0}}^{t_{1}}Tr(\\frac{\\partial f}{\\partial z(t)})\\, dt \\tag*{(8)}\n",
        "\\end{align}\n",
        "\n",
        "Let's pause for a moment and try to understand the meaning of the last equality.\n",
        "\n",
        "- $p(z(t_{0}))$ is the unknown probability we try to approximate, we only know samples from it, but not the true distribution.\n",
        "\n",
        "- $p(z(t_{1}))$ is the base distribution (usually a very common one, for example, a Gaussian), from where we want to find a function $f(z, t, θ)$, which transforms the Gaussian into $p(z(t_{0}))$ (Fig.2).  \n",
        "- The trace is a correction factor that accounts for the changes in volume of the variable\n",
        "\n",
        "<div>\n",
        "</div> \n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1iVfH_W_gEapnJyk5HyG0Fgpg_V9PwnFt\">\n",
        "<center > <figcaption> Fig.2 - $p(z(t_{1}))$ to $p(z(t_{0}))$"
      ],
      "metadata": {
        "id": "-EK2qoAeudEQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jGZAVZI_TYVC"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import time\n",
        "from tqdm import tqdm\n",
        "from torchdiffeq import odeint_adjoint as odeint_adjoint\n",
        "from torchdiffeq import odeint\n",
        "from math import pi as pi\n",
        "import random"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this cell, I wrote manually a PDF from where I sample. Even though the true distribution is known beforehand, I only sample from it."
      ],
      "metadata": {
        "id": "tKzrY7jry3fd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "scale1 = 0.0005\n",
        "a = 4 * pi * scale1\n",
        "def psi(x, y):\n",
        "    r = np.sqrt(x**2 + y**2)\n",
        "    if x == 0:\n",
        "        theta = pi / 2\n",
        "    else:\n",
        "        theta = np.arctan(y / x)\n",
        "    return np.abs(1 / (64 * np.sqrt(5) * a * np.sqrt(a)) * (1 - r / (12 * a)) * (r / a)**2 * np.exp(-r / (4 * a)) * np.sqrt(15 / (8 * pi)) * np.sin(2 * theta))"
      ],
      "metadata": {
        "id": "03P3uGf5X4MS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## How to construct the CNF model?\n",
        "\n",
        "Here, the `NN` class is where I built the neural network, a very short one with only one hidden layer consisting of 64 vertices. The activation function that I used is tanh(Relu won't work because the function need to be Lipschitz continuous in order to guarantee the uniqueness of the differential equation). My output layer, instead of predicting a single velocity on the x or y direction, it predicts the average of 128 velocities in this case (for stability).\n",
        "\n",
        "The `CNF` class computes $\\frac{dz}{dt}$, where the input of the neural network is $z(t)$.\n",
        "\n",
        "The `JacTrace` function calculates the trace of $A_{ij}$. At first glance the code looks confusing and I will show you how it works.\n",
        "\n",
        "Every sample has it's own dynamics. Therefore, we can think of $f$ for some $z_{i}$ (a 2D vector) as a mapping from $\\mathbb{R^2}$ to $\\mathbb{R^2}$:\n",
        "\\begin{align}\n",
        "f(x,y)&=\\begin{bmatrix}\n",
        "f_{1}(x, y) \\\\\n",
        "f_{2}(x, y) \\\\\n",
        "\\end{bmatrix}\n",
        "\\end{align}\n",
        "\n",
        "The Jacobian will be\n",
        "\n",
        "\\begin{align}\n",
        "J &= \\begin{bmatrix}\n",
        "\\frac{\\partial f_{1}(x, y)}{\\partial x} \\, \\frac{\\partial f_{1}(x, y)}{\\partial y} \\\\\n",
        "\\frac{\\partial f_{2}(x, y)}{\\partial x} \\, \\frac{\\partial f_{2}(x, y)}{\\partial y} \n",
        "\\end{bmatrix} \n",
        "\\end{align}\n",
        "\n",
        "The trace of $J$ is\n",
        "\\begin{align}\n",
        "Tr(J) = \\frac{\\partial f_{1}(x, y)}{\\partial x} + \\frac{\\partial f_{2}(x, y)}{\\partial y}\n",
        "\\end{align}\n",
        "\n",
        "The `torch.autograd.grad` function is calculating $\\frac{\\partial f_{1}(x, y)}{\\partial x}$ and $\\frac{\\partial f_{2}(x, y)}{\\partial y}$ for every $z_{i}$ and stores them into a row vector\n",
        "\n",
        "\\begin{align}\n",
        "&\\begin{bmatrix}\n",
        "\\frac{\\partial f_{11}(x, y)}{\\partial x}, \\frac{\\partial f_{21}(x, y)}{\\partial x}, ... ,\\frac{\\partial f_{n1}(x, y)}{\\partial x}\n",
        "\\end{bmatrix} \\\\\n",
        "&\\begin{bmatrix}\n",
        "\\frac{\\partial f_{12}(x, y)}{\\partial x}, \\frac{\\partial f_{22}(x, y)}{\\partial x}, ..., \\frac{\\partial f_{n2}(x, y)}{\\partial x}\n",
        "\\end{bmatrix}\n",
        "\\end{align}\n",
        "\n",
        "where $f_{ij}$ it's the j'th ($j \\in {1, 2}$) output of the i'th coresponding dynamics of $z_{i}$\n",
        "\n",
        "The trace is just the sum of the two row vectors."
      ],
      "metadata": {
        "id": "3c4jodtlz6r1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def JacTrace(f, z):\n",
        "    sum_d = 0\n",
        "    for i in range(z.shape[1]):\n",
        "        sum_d += torch.autograd.grad(f[:, i].sum(), z, create_graph = True)[0][:, i]\n",
        "    return sum_d\n",
        "\n",
        "class CNF(nn.Module):\n",
        "    def __init__(self, in_dim):\n",
        "        super().__init__()\n",
        "        self.model = NN(in_dim)\n",
        "\n",
        "    def forward(self, t, state):\n",
        "        z = state[0]\n",
        "        log_p = state[1]\n",
        "\n",
        "        batch_size = z.shape[0]\n",
        "\n",
        "        with torch.set_grad_enabled(True):\n",
        "            z.requires_grad_(True)\n",
        "\n",
        "            d_z_dt = self.model(z)\n",
        "\n",
        "            d_log_p_z_dt = -JacTrace(d_z_dt, z).view(batch_size)\n",
        "\n",
        "        return (d_z_dt, d_log_p_z_dt)\n",
        "        \n",
        "class NN(nn.Module):\n",
        "    def __init__(self, in_dim):\n",
        "        super().__init__()\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Linear(in_dim, 64),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(64, 64),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(64, 128 * 2)\n",
        "        )\n",
        "    \n",
        "    def forward(self, z):\n",
        "        z = self.layers(z)\n",
        "        z = z.reshape(z.shape[0], 128, 2)\n",
        "\n",
        "        return z.mean(dim = 1)"
      ],
      "metadata": {
        "id": "600CgJRoX4cY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## How do I sample from a PDF?\n",
        "\n",
        "The technique is called \"Rejection sampling\". In general, it can take samples from any PDF. If we evaluate a probability distribution $p_{X}(x)$ at some point $x_{0}$, the output it's the \"likelihood\" $p_{X}(X = x_{0})$. If the likelihood it's large we would expect to have many samples in the neighborhood of $x_{0}$ and if it's small we would expect only a few points near $x_{0}$. That's exactly how the method works. It takes an $x$ from a closed interval according to a uniform probability distribution (it's the same for every $x$) and a $y$ from the interval [0, the maximum value of the PDF] (Fig.1).\n",
        "\n",
        "The uniform distribution on an interval $[a, b]$, b > a, is $p = \\frac{1}{b - a}$\n",
        "\n",
        "<center>\n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1bQEoVUmOXWyBRa2fGMnWaP4LB1k0evN4\" width = \"400\"/>\n",
        "<figcaption> Fig.1 - Rejection sampling </figcaption>\n",
        "</center>\n",
        "\n",
        "In the figure, the blue dots are the rejected samples and the green dots the accepted ones.\n",
        "If the point $x_{0}$ is picked by the uniform distribution, then the chance of rejection is equal to the length of the blue line which lies above the red curve and the chance of acceptance is equal to the length of the green line (This analogy is based on the fact that the probability $p = \\frac{1}{max(PDF)}$, i.e a constant, hence, the fraction $ A = \\frac {length \\, blue \\, line}{length \\, green \\, line}$ tells us whether there is a more chance to be accepted ($A < 1$) or not ($A > 1$). For the point $x_{0}$ is more likely to be rejected comparted to $x_{1}$. This is how the algorithm will pick much more points that are closer to $x_{1}$ rather than closer to $x_{0}$.   "
      ],
      "metadata": {
        "id": "lVIPWpj3DCd4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "t0 = 0\n",
        "t1 = 10\n",
        "Gaussian_distribution = torch.distributions.MultivariateNormal(\n",
        "    loc = torch.tensor([0., 0.]).to(device), covariance_matrix = torch.tensor([[0.1, 0.], [0., 0.1]]).to(device))\n",
        "\n",
        "def sample(f, xmin = -0.1, xmax = 0.1, ymin = -0.1, ymax = 0.1):\n",
        "    while True:\n",
        "        x = np.random.uniform(low = xmin, high = xmax)\n",
        "        y = np.random.uniform(low = ymin, high = ymax)\n",
        "        z = np.random.uniform(low = 0, high = 250)\n",
        "        if z < f(x, y):\n",
        "            return x * 10, y * 10\n",
        "        \n",
        "def batch(batch_size):\n",
        "    samps = [sample(psi) for i in range(batch_size)]\n",
        "    # a = Gaussian_distribution.sample([batch_size]).type(torch.float32).to(device)\n",
        "    a = torch.tensor(samps).type(torch.float32).to(device)\n",
        "    b = torch.zeros(batch_size).type(torch.float32).to(device)\n",
        "    return a, b"
      ],
      "metadata": {
        "id": "iLPVTKoWX8Yy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## How a change in the samples will induce a change in the random variable?\n",
        "\n",
        "If you recall the definition of a random variable, i.e, a function $f : S → X$, where $S$ is the sample space and $X$ is a measurable space. This function it's useful because in Mathematics it would be much more convenient to use numbers rather than trying to write $P(\"heads\")$ instead of $P(1)$ if $f(\"heads\") = 1$. \n",
        "\n",
        "The random variables $X$ and $Y$ (Fig.1) transform their coresponding sampling spaces $S$ and $S'$, respectively, into $x$ and $y$. If there is a bijective function that maps directly from $S$ to $S'$, you will be able to follow the arrows to move from x to y and vice-versa producing a transformation between the two random variables. That's an astonishing idea.\n",
        "\n",
        "<center>\n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1JMan4gvO9b6ywqCa6wMZqPtK8J1St2bF\" width = \"400\"/>\n",
        "<figcaption> Fig.1 - Change in the random variable </figcaption>\n",
        "</center>\n"
      ],
      "metadata": {
        "id": "CXwjbjv6gcih"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## What is the loss function?\n",
        "\n",
        "CNF is an example of an unsupervised learning model, so we don't have labels for the outputs and it's different from choosing some loss function to minimize the difference between the prediction and the labels. We can find a loss function, but we need to think what are we trying to minimize. We have the distribution that our model predicted and we need a function to quantify or to measure how far away are the prediction and the target distribution. In Statistics there is such a function, called **KL-divergence**, that can tell us how different two PDFs are. I will denote the model prediction as $p_{\\theta}(z(0))$ (I put $\\theta$ as a subscript because the probability depends on the parameters of the neural network) and the target distribution as $q(z(0))$.\n",
        "\n",
        "\\begin{align}\n",
        "KL(q(z(0))\\ || \\ p_{\\theta}(z(0))) = \\int q(z(0))\\cdot \\log \\frac{q(z(0))}{p_{\\theta}(z(0))} \\ dz(0)\n",
        "\\end{align}\n",
        "\n",
        "However, we are working with a finite number of samples, so the integrals will turn into sums.\n",
        "\n",
        "\\begin{align}\n",
        "KL(q(z(0))\\ || \\ p_{\\theta}(z(0))) &= \\sum_{i}^N q(z_{i}(0))\\cdot \\log \\frac{q(z_{i}(0))}{p_{\\theta}(z_{i}(0))} \\\\\n",
        "&= \\sum_{i}^N q(z_{i}(0))\\cdot \\log q(z_{i}(0)) - \\sum_{i}^N q(z_{i}(0))\\cdot \\log p_{\\theta}(z_{i}(0)) \\\\\n",
        "&= const - \\sum_{i}^N q(z_{i}(0))\\cdot \\log p_{\\theta}(z_{i}(0))\n",
        "\\end{align}\n",
        "\n",
        "We take the derivative of the expression w.r.t the parameters $\\theta$\n",
        "\\begin{align}\n",
        "\\frac{\\partial KL(q(z(0))\\ || \\ p_{\\theta}(z(0)))}{\\partial \\theta} = - \\nabla_{\\theta}\\sum_{i}^N q(z_{i}(0))\\cdot \\log p_{\\theta}(z_{i}(0))\n",
        "\\end{align}\n",
        "\n",
        "There is still something in our equation that can't be calculated because we don't know it, i.e $q(z_{i}(0))$. Fortunately, you can estimate it by replacing with a mean\n",
        "\n",
        "\\begin{align}\n",
        "\\frac{\\partial KL(q(z(0))\\ || \\ p_{\\theta}(z(0)))}{\\partial \\theta} &= - \\frac{1}{N}\\sum_{i}^N \\nabla_{\\theta} \\log p_{\\theta}(z_{i}(0)) \\\\\n",
        "&= \\frac{1}{N}\\sum_{i}^N \\nabla_{\\theta} (\\log p(z_{i}(t_{1})) + \\int_{t_{0}}^{t_{1}} Tr(J) \\, dt)\n",
        "\\end{align}\n",
        "\n",
        "One last step, we replace $\\log p_{\\theta}(z_{i}(0))$ with equation $(8)$\n",
        "\n",
        "\\begin{align}\n",
        "\\frac{\\partial KL(q(z(0))\\ || \\ p_{\\theta}(z(0)))}{\\partial \\theta} = -\\frac{1}{N}\\sum_{i}^N \\nabla_{\\theta} (\\log p(z_{i}(t_{1})) + \\int_{t_{0}}^{t_{1}} Tr(J) \\, dt) \\tag*{(9)}\n",
        "\\end{align}\n",
        "\n",
        "Finally, we can do gradient descent by minimizing the KL-divergence which is the same as maximizing the log likelihood, because the $\\log p_{\\theta}(z_{i}(0))$ are the log likelihoods."
      ],
      "metadata": {
        "id": "BC2Q8pDfw-Cy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The training part\n",
        "\n",
        "Firstly, I sample from the PDF and initialize `log_p_z_t0` with 0 because, after the integration, I want to have only -$\\int_{t_{0}}^{t_{1}} Tr(J) \\ dt$. Here, `log_p_z_t0` it's not the same thing as the log probability of the target distribution. The size of your batch or sample points depends if they are enough to represent the key aspects of the PDF or not. After the ODE I calculate the `log_likelihood` by distributing the minus sign to the right hand side of the equation $(9)$. The base distribution that I use is a Gaussian with mean = 0 and variance given by the `covariance_matrix`."
      ],
      "metadata": {
        "id": "tJhfJBul_KL-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "f = CNF(in_dim = 2).to(device)\n",
        "\n",
        "optimizer = optim.Adam(f.parameters(), lr = 1e-3)\n",
        "\n",
        "def training(loops):\n",
        "    for i in tqdm(range(1, loops)):\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        z_t0, log_p_z_t0 = batch(2000)\n",
        "        \n",
        "        z_t1, log_p_z_t1 = odeint (\n",
        "            f,\n",
        "            (z_t0, log_p_z_t0),\n",
        "            torch.tensor([t0, t1]).type(torch.float32).to(device),\n",
        "            rtol = 1e-5,\n",
        "            atol = 1e-5,\n",
        "            method = 'dopri5'\n",
        "        )\n",
        "\n",
        "        z_t11, log_p_z_t11 = z_t1[-1], log_p_z_t1[-1]\n",
        "\n",
        "        log_likelihood = log_p_z_t11 - Gaussian_distribution.log_prob(z_t11)\n",
        "        loss = log_likelihood.mean(0)\n",
        "        \n",
        "        if i % 100 == 0:\n",
        "            print(loss)\n",
        "        #plt.hist2d(z_t11.detach().cpu().numpy()[:, 0], z_t11.detach().cpu().numpy()[:, 1], bins = 300)\n",
        "        #plt.show()\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "training(10001)"
      ],
      "metadata": {
        "id": "ij2Mm-K5X_Wm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "nr_samples = 30000\n",
        "time = 100\n",
        "true_distribution, _, = batch(nr_samples)"
      ],
      "metadata": {
        "id": "34_b-lfAYC02"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The visualisation part\n",
        "\n",
        "Through the training session, the neural network (the function $f$) learned where the samples should be maped at each timestep such that at $t_{1}$ they end up looking like samples from the Gaussian. Also, the random variable which initially was the $X-Y$ axis has become ,maybe, a very complex coordinate system (Fig.1). \n",
        "\n",
        "<center>\n",
        "<img src = \"https://drive.google.com/uc?export=view&id=1yYq3qbMNmMzOCdVNMSA_Ml6VhNj94S2m\">\n",
        "<figcaption> Fig.1 - Space transformation\n",
        "</center>\n",
        "\n",
        "A big advantage of CNF is the reversability of the transformation. If you run the ODE backward in time and sample from the base distribution you will get samples from the target distribution. In `z_t_sample` are the positions of the samples at time $t$ starting with the Gaussian. For computing the probability density of the model at time $t$, a second ODE solver is used for the density space (or the random variable space) and then stores the values in `z_t_density`. The density at $t_{0}$ is a grid (the $X-Y$ axis)."
      ],
      "metadata": {
        "id": "2x55O8tpC0HU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with torch.no_grad():\n",
        "    \n",
        "    z_t0 = Gaussian_distribution.sample([nr_samples]).to(device)\n",
        "    log_p_t1 = torch.zeros(nr_samples).type(torch.float32).to(device)\n",
        "    # v_t0 = torch.randn(nr_samples, 2).type(torch.float32).to(device)\n",
        "    \n",
        "    #aug_t0 = torch.cat((v_t0, z_t0), dim = 1)\n",
        "\n",
        "    z_t_sample, rr = odeint (\n",
        "        f,\n",
        "        (z_t0, log_p_t1),\n",
        "        torch.tensor(np.linspace(t1, t0, time)).type(torch.float32).to(device),\n",
        "        rtol = 1e-5,\n",
        "        atol = 1e-5,\n",
        "        method = 'dopri5'\n",
        "    )\n",
        "    \n",
        "    c = 1.5\n",
        "    \n",
        "    x = np.linspace(-c, c, 100)\n",
        "    y = np.linspace(-c, c, 100)\n",
        "    x, y = np.meshgrid(x, y)\n",
        "    x, y = x.reshape(-1), y.reshape(-1)\n",
        "    points = np.vstack((x, y)).T\n",
        "    grid = torch.tensor(points).type(torch.float32).to(device)\n",
        "    log_p_t0 = torch.zeros(grid.shape[0]).type(torch.float32).to(device)\n",
        "    #v_t1 = torch.randn(grid.shape[0], 2).type(torch.float32).to(device)\n",
        "    #aug_t0 = torch.cat((v_t1, grid), dim = 1)\n",
        "    z_t_density, log_p_t1 = odeint (\n",
        "        f,\n",
        "        (grid, log_p_t0),\n",
        "        torch.tensor(np.linspace(t0, t1, time)).type(torch.float32).to(device),\n",
        "        rtol = 1e-5,\n",
        "        atol = 1e-5,\n",
        "        method = 'dopri5'\n",
        "    )\n",
        "        \n",
        "    for (t, z_sample, z_density, log_p) in zip(\n",
        "            np.linspace(t0, t1, time),\n",
        "            z_t_sample, z_t_density, log_p_t\n",
        "    ):\n",
        "        fig = plt.figure(figsize = (12, 4), dpi = 200)\n",
        "        plt.axis('off')\n",
        "        plt.margins(0, 0)\n",
        "        img1 = fig.add_subplot(1, 3, 1)\n",
        "        img1.set_title('True distribution')\n",
        "        img1.get_xaxis().set_ticks([])\n",
        "        img1.get_yaxis().set_ticks([])\n",
        "        \n",
        "        img2 = fig.add_subplot(1, 3, 2)\n",
        "        \n",
        "        img2.set_title('Model sample')\n",
        "        img2.get_xaxis().set_ticks([])\n",
        "        img2.get_yaxis().set_ticks([])\n",
        "        \n",
        "        img3 = fig.add_subplot(1, 3, 3)\n",
        "        \n",
        "        img3.set_title('Model probability distribution')\n",
        "        img3.get_xaxis().set_ticks([])\n",
        "        img3.get_yaxis().set_ticks([])\n",
        "        \n",
        "        img1.hist2d(true_distribution.detach().cpu().numpy()[:, 0], true_distribution.detach().cpu().numpy()[:, 1], density = True, bins = (300, 300), range = [(-c, c), (-c, c)])\n",
        "        \n",
        "        img2.hist2d(z_sample.detach().cpu().numpy()[:, 0], z_sample.detach().cpu().numpy()[:, 1], density = True, bins = (300, 300), range = [(-c, c), (-c, c)])\n",
        "        logp = Gaussian_distribution.log_prob(z_density) - log_p.view(-1)\n",
        "        img3.tricontourf(grid.detach().cpu().numpy()[:, 0], grid.detach().cpu().numpy()[:, 1], np.exp(logp.detach().cpu().numpy()), levels = 200)\n",
        "        plt.show()"
      ],
      "metadata": {
        "id": "kJ5V1cGeYEcg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## References\n",
        "\n",
        "- https://arxiv.org/abs/1806.07366\n",
        "- https://deepgenerativemodels.github.io/notes/flow/\n",
        "- https://paperswithcode.com/method/normalizing-flows\n",
        "- https://jmlr.org/papers/volume22/19-1028/19-1028.pdf\n",
        "- https://arxiv.org/abs/1908.09257\n",
        "- https://pyro.ai/examples/normalizing_flows_i.html\n",
        "- https://janosh.dev/blog/normalizing-flows\n",
        "- http://proceedings.mlr.press/v37/rezende15.pdf\n",
        "- https://www.slideshare.net/GrigorisChrysos/an-introduction-on-normalizing-flows\n",
        "- https://siboehm.com/articles/19/normalizing-flow-network\n",
        "- https://towardsdatascience.com/probability-concepts-explained-maximum-likelihood-estimation-c7b4342fdbb1\n",
        "- https://arxiv.org/abs/2012.04228\n",
        "\n"
      ],
      "metadata": {
        "id": "dPf34hm3KySx"
      }
    }
  ]
}