{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 221,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 221,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12464 characters]\n",
      "text/plain": [
       "<Figure size 432x288 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",
    "\n",
    "plt.plot(x, base.pdf(x))\n",
    "\n",
    "np.trapz(base.pdf(x), 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": 222,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.641895835477563"
      ]
     },
     "execution_count": 222,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10340 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f(x):\n",
    "    return x**2\n",
    "\n",
    "transformed = f(base.pdf(x))\n",
    "plt.plot(f(x), transformed)\n",
    "\n",
    "np.trapz(transformed, f(x))"
   ]
  },
  {
   "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(x)\\frac{\\text{d}f^{-1}(y)}{\\text{d}y} \\\\\n",
    " &=& P(x)f'^{-1}(y) \\\\\n",
    " &=& P(f^{-1}(y))f'^{-1}(y)\n",
    "\\end{eqnarray} $$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9987379589284238"
      ]
     },
     "execution_count": 223,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16888 characters]\n",
      "text/plain": [
       "<Figure size 432x288 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",
    "\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": "code",
   "execution_count": 224,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9987894616589735"
      ]
     },
     "execution_count": 224,
     "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*x\n",
    "\n",
    "assert np.allclose(f_i(f(x)),  x)\n",
    "\n",
    "y = f(x)\n",
    "px = base.pdf(x)\n",
    "\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$, than 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(y)&=&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": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43716 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "m = dist.Normal(torch.zeros(2), scale=torch.ones(2))\n",
    "s = m.sample((1000,))\n",
    "\n",
    "plt.scatter(s[:, 0], s[:, 1], alpha=0.2)\n",
    "plt.axis('scaled');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    torch.exp(x) + 2\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
