{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.distributions as dist\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transforming probability distributions\n",
    "Normalizing flows uses transformation to obtain complex posteriors. \n",
    "This notebook explores transformations of distributions.\n",
    "\n",
    "We start with a basis Gaussian $\\mathcal{N}(\\mu=1, \\sigma=0.1)$. These parameters are chose so we can look at transformations in the positive domain and don't have trouble with square rooting negative numbers.\n",
    "\n",
    "Below we plot the **pdf** and integrate to validate $\\int P(x)dx=1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f90e9cb7208>]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18156 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(0.01, 2, num=100)\n",
    "\n",
    "base = stats.norm(1, 0.1)\n",
    "print(np.trapz(base.pdf(x), x))\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "plt.xlabel(r'$x$')\n",
    "plt.ylabel(r'$p(x)$')\n",
    "plt.plot(x, base.pdf(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformation\n",
    "\n",
    "As we apply $f(x) = y = x^2$ we see that the result of the transformation doesn't suffice the requirements of a probability distribution, as $\\int P(y)dy \\neq 1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.641895835477563\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f90e98c6588>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 25572 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f(x):\n",
    "    return x**2\n",
    "\n",
    "y = f(x)\n",
    "transformed = f(base.pdf(x))\n",
    "print(np.trapz(transformed, y))\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "plt.xlim(0, 2)\n",
    "plt.plot(y, transformed, label=r'$P(y)$')\n",
    "plt.plot(x, base.pdf(x), label='$P(x)$', ls='-.')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change of variables\n",
    "\n",
    "We want to transform $P(x)$ to $P(y)$ with $f: \\mathbb{R}^n \\mapsto \\mathbb{R}^n $ and $f(x) = y$\n",
    "The transformation blows up or shrinks the area/ volume (in N-D) of the distribution, so we need to determine a transformation so that:\n",
    "\n",
    "\n",
    "$$ \n",
    "\\begin{eqnarray}\n",
    "\\int P(x)dx &=&\\int P(y)dy = 1 \\\\\n",
    "P(x)dx &=&P(y)dy \\\\\n",
    "P(y) &=&P(x)\\frac{dx}{dy}\n",
    "\\end{eqnarray} $$\n",
    "\n",
    "So to keep an integral of 1 for $P(y)$ (the definition of a probability distribution), we need to multiply $P(x)$ with its derivative w.r.t. $y$.\n",
    "Therefore we need to express $x$ in $y$. This can be done if we only have invertible transformations.\n",
    "\n",
    "$$\\begin{eqnarray}\n",
    "f(x) &=& y \\\\\n",
    "f^{-1}(y) &=& x\n",
    "\\end{eqnarray} $$\n",
    "\n",
    "$$\\begin{eqnarray}\n",
    " P(y) &=& P(f^{-1}(y))\\frac{\\text{d}f^{-1}(y)}{\\text{d}y} \\\\\n",
    " &=& P(f^{-1}(y))f'^{-1}(y)\n",
    "\\end{eqnarray} $$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9987379589284238\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f90e9c442e8>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 27664 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f_i(y):\n",
    "    return y**0.5\n",
    "\n",
    "def f_i_prime(y):\n",
    "    return 0.5*y**-0.5\n",
    "\n",
    "assert np.allclose(f_i(f(x)),  x)\n",
    "\n",
    "y = f(x)\n",
    "px = base.pdf(x)\n",
    "transformed = px * f_i_prime(f_i(y))\n",
    "print(np.trapz(transformed, y))\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "plt.xlim(0, 2)\n",
    "plt.plot(y, transformed, label=r'$P(y)$')\n",
    "plt.plot(x, base.pdf(x), label='$P(x)$', ls='-.')\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9987894616589734"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12920 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Same principle, other transformation\n",
    "\n",
    "def f(x):\n",
    "    return x**0.5\n",
    "\n",
    "def f_i(y):\n",
    "    return y**2\n",
    "\n",
    "def f_i_prime(y):\n",
    "    return 2*y\n",
    "\n",
    "assert np.allclose(f_i(f(x)),  x)\n",
    "\n",
    "y = f(x)\n",
    "px = base.pdf(x)\n",
    "transformed = px * f_i_prime(f_i(y))\n",
    "\n",
    "plt.plot(x, px)\n",
    "plt.plot(y, transformed)\n",
    "plt.xlim(0, 2)\n",
    "\n",
    "np.trapz(transformed, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conditions\n",
    "The transformations we've used were strictly increasing. \n",
    "\n",
    "In the example we've used above $\\frac{dx}{dy}$ was positive, as the transformation was strictly increasing, in the strictly decreasing case, $\\frac{dx}{dy}$ is negative and $P(y)=-P(x)\\frac{dx}{dy}$. If we take the absolute value of the derivative we have one equation for both cases.\n",
    "\n",
    "$$ \n",
    "\\begin{eqnarray}\n",
    "P(y)&=&P(x) \\cdot \\left| \\frac{dx}{dy} \\right| \\\\\n",
    " &=& P(f^{-1}(y)) \\cdot \\left| f'^{-1}(y)\\right|\n",
    "\\end{eqnarray}\n",
    "$$\n",
    "\n",
    "The intuition of taking the modulus is that if the (local) rate of change of $x + dx = y + dy > 0$, then we can ensure ourselves that the rate of amount of probability is increasing, and the total amount of probability is preserved. This is ensured by conditioning $\\frac{dx}{dy} \\geq 0$.\n",
    "\n",
    "![](img/absolute.png)\n",
    "\n",
    "## Multiple dimensions\n",
    "In N-D $\\frac{dx}{dy}$ is expressed in the Jacobian matrix, i.e. the matrix with first order partial derivatives $\\frac{\\partial(x)}{\\partial{y}}=\\mathbf{J}$. This absolute value, is now a determinant, such that;\n",
    "\n",
    "$$ \n",
    "\\begin{eqnarray}\n",
    "P(y)&=&P(f^{-1}(y)) \\cdot \\left| \\det \\frac{\\partial{ f(y)^{-1} }}{\\partial{y}} \\right| \\\\\n",
    "&=&P(f^{-1}(y)) \\cdot \\left| \\det \\mathbf{J}_{f^{-1}(y)} \\right| \\\\\n",
    "\\end{eqnarray}\n",
    "$$\n",
    "\n",
    "The geometric meaning of the determinant is a scaling factor. $\\det A$ tells us how much an n-dimensional volume is scaled by applying the transformation of $A$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$Q(z') = Q(z) \\left| \\det \\frac{\\partial{f^{-1}}}{\\partial{z}} \\right| $$\n",
    "\n",
    "$$ f(x) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9905751293230018\n",
      "(50, 50, 2)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.9905751293230018"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 13820 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x1 = np.linspace(1, 5, num=50)\n",
    "x2 = np.linspace(1, 5, num=50)\n",
    "x1_s, x2_s = np.meshgrid(x1 ,x2)\n",
    "\n",
    "base = stats.multivariate_normal(mean=[3, 3], cov=0.5)\n",
    "assert base.pdf([3, 3]) > base.pdf([1, 1])\n",
    "\n",
    "# field of all possible events X\n",
    "x_field = np.concatenate([x1_s[..., None], x2_s[..., None]], axis=-1)\n",
    "\n",
    "print(np.trapz(np.trapz(base.pdf(x_field), x_field[:, 0, 1], axis=0), x_field[0, :, 0]))\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.contourf(x1_s, x2_s, base.pdf(x_field))\n",
    "np.trapz(np.trapz(base.pdf(x_field), x_field[:, 0, 1], axis=0), x_field[0, :, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jacobian matrix.\n",
    "\n",
    "Assume the function $f: \\mathbb{R}^n \\mapsto \\mathbb{R}^m$. The jacobian is a 2D matrix, that stores on the partial derivatives of all the outputs and all the inputs. The width of the jacobian is equal to the functions input $n$. The height of the jacobian is equal to the functions output $m$.\n",
    "\n",
    "$$ \\mathbf{J} = \\begin{bmatrix}\n",
    "\\frac{\\partial f_1}{\\partial x_1} & \\dots & \\frac{\\partial f_1}{\\partial x_n} \\\\[6pt]\n",
    "\\vdots & \\ddots & \\vdots \\\\[6pt]\n",
    "\\frac{\\partial f_m}{\\partial x_1} & \\dots & \\frac{\\partial f_m}{\\partial x_n} \\\\[6pt]\n",
    "\\end{bmatrix}$$\n",
    "\n",
    "We will map 2D events to another domain. So we will have two inputs $n$ and two outputs $m$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x1, x2):\n",
    "    return torch.exp(x1 / 3), x2**2\n",
    "\n",
    "def f_i(y1, y2):\n",
    "    return 3 * torch.log(y1), y2**0.5\n",
    "\n",
    "x_field = torch.tensor(x_field)\n",
    "# Transform x events to y events\n",
    "y_field = np.concatenate(f(x_field[..., 0, None], x_field[..., 1, None]), axis=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_det_jac(y_field):\n",
    "    # create for every y1, y2 combination the determinant of the jacobian f_i(y1, y2)\n",
    "    det_jac = np.zeros((y_field.shape[0], y_field.shape[1]))\n",
    "    for i in range(y_field.shape[0]):\n",
    "        for j in range(y_field.shape[1]):\n",
    "            y_field = torch.tensor(y_field)\n",
    "            y_field.requires_grad_(True);\n",
    "            fiy = torch.cat(f_i(y_field[..., 0, None], y_field[..., 1, None]), dim=-1)\n",
    "            fiy[i, j].sum().backward()\n",
    "            \n",
    "            # Ouputs of the partial derivatives are independent.  I.e. f1 is dependent of y1 and not y2\n",
    "            # and vice versa f2 is dependent of y2 and not y1\n",
    "            # therefore the multiplication w/ 0 \n",
    "            row1 = y_field.grad[i, j].data.numpy() * np.array([1., 0.])\n",
    "            row2 = y_field.grad[i, j].data.numpy() * np.array([0., 1.])\n",
    "            \n",
    "            det = np.linalg.det(np.array([row1, row2]))\n",
    "            det_jac[i, j] = det\n",
    "    return det_jac\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:6: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.contour.QuadContourSet at 0x7f90e4f80208>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11720 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "px = base.pdf(x_field)\n",
    "transformed = px * np.abs(create_det_jac(y_field))\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.contourf(y_field[..., 0], y_field[..., 1], transformed)\n",
    "\n",
    "print(np.trapz(np.trapz(transformed, y_field[:, 0, 1], axis=0), y_field[0, :, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 278,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9907110850291531"
      ]
     },
     "execution_count": 278,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.trapz(np.trapz(transformed, y_field[:, 0, 1], axis=0), y_field[0, :, 0])"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
