discrete/Gaussians in R2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "baking-christian",
   "metadata": {},
   "source": [
    "# Optimal transport in $\\mathbb{R}^2$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "handed-reporter",
   "metadata": {},
   "source": [
    "Let $X \\sim N(0_{N}, \\Sigma_{X})$ and $Y \\sim N(0_{N}, \\Sigma_{Y})$ be two bivariate gaussian distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bibliographic-converter",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Random\n",
    "import LinearAlgebra\n",
    "import CovarianceEstimation\n",
    "using Distributions\n",
    "using Plots\n",
    "using StatsPlots\n",
    "seed = 1234\n",
    "gen = Random.Xoshiro(seed);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "metric-latvia",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Covariance matrix of X\n",
    "ΣX = [0.5 0.; 0. 0.5];\n",
    "# Covariance matrix of Y\n",
    "ΣY = [2. 0.; 0. 2.];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "supreme-buddy",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Distributions of X ~ N(0, I_n) and Y ~ N(0, sqrt(2) * I_n)\n",
    "distμ = MvNormal(ΣX);\n",
    "distν = MvNormal(ΣY);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "mysterious-browse",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100;\n",
    "x = rand(distμ, N);\n",
    "y = rand(distν, N);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "rental-solution",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pdf of X and Y\n",
    "f = Distributions.pdf(distμ, x);\n",
    "g = Distributions.pdf(distν, y);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "resident-reception",
   "metadata": {},
   "source": [
    "**Density function of X**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "successful-trading",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface(x[1,:], x[2, :], f, c=:PuBu, colorbar = :none)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aggregate-chair",
   "metadata": {},
   "source": [
    "**Density function of Y**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prospective-dividend",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface(y[1,:], y[2, :], g, c=:OrRd, colorbar=:none)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "korean-origin",
   "metadata": {},
   "source": [
    "**Density function of X ($f$) and density function of Y ($g$)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "color-biodiversity",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface(y[1,:], y[2, :], g, c=:OrRd, colorbar = :none)\n",
    "surface!(x[1,:], x[2, :], f, c=:PuBu, alpha = 0.4, colorbar = :none)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "infrared-nurse",
   "metadata": {},
   "source": [
    "**Loss function**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "formal-burning",
   "metadata": {},
   "source": [
    "$$c_{T}(x, y) = \\frac{1}{2} \\|x - y\\|^2$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "plastic-radical",
   "metadata": {},
   "source": [
    "**Optimal transportation map of X to Y**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sustained-qatar",
   "metadata": {},
   "source": [
    "If $X \\sim N(0, \\Sigma_{X})$ and $Y \\sim N(0, \\Sigma_{Y})$, and the cost function is $c_{T}(x,y) = \\frac{1}{2} \\|Ax-y\\|^2$ where $A$ is invertible, the optimal transportation map $T$ of $X$ to $Y$ has a closed form solution given by\n",
    "\n",
    "$$\n",
    "T : X -> Y \\\\\n",
    "T(x) = (A^{T})^{-1} \\Sigma_{X}^{-1/2}(\\Sigma_{X}^{1/2} A^{T} \\Sigma_{Y} A \\Sigma_{X}^{1/2})^{1/2} \\Sigma_{X}^{-1/2} x\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pacific-protein",
   "metadata": {},
   "source": [
    "If $A = I_{n}$ then\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "T(x) &= \\Sigma_{X}^{-1/2}(\\Sigma_{X}^{1/2} \\Sigma_{Y} \\Sigma_{X}^{1/2})^{1/2} \\Sigma_{X}^{-1/2} x \\\\\n",
    "    &= \\Sigma_{X}^{-1}(\\Sigma_{X} \\Sigma_{Y})^{1/2} x\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "spread-medline",
   "metadata": {},
   "outputs": [],
   "source": [
    "function optimal_fn(x, μ, ν)\n",
    "    ΣX = cov(μ)\n",
    "    ΣY = cov(ν)\n",
    "    v = inv(ΣX) * ((ΣX * ΣY) ^ (1/2))\n",
    "    n = size(x, 2)\n",
    "    tx = zeros(2, n)\n",
    "    for i in 1:n\n",
    "        tx[:, i] = v * x[:, i]\n",
    "    end\n",
    "    return tx\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unable-reaction",
   "metadata": {},
   "outputs": [],
   "source": [
    "tx = optimal_fn(x, distμ, distν);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deluxe-burner",
   "metadata": {},
   "source": [
    "**Visual checking**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efficient-fancy",
   "metadata": {},
   "source": [
    "To plot the distribution of the map of $X$, we assume that $T(X) \\sim N(0_{N}, \\Sigma_{\\gamma})$\n",
    "where $\\Sigma_{\\gamma}$ is the estimated covariance matrix of $T(X)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "silver-board",
   "metadata": {},
   "outputs": [],
   "source": [
    "Σγ = CovarianceEstimation.cov(CovarianceEstimation.SimpleCovariance(), transpose(tx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "friendly-slope",
   "metadata": {},
   "outputs": [],
   "source": [
    "distγ = MvNormal(Σγ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "black-newfoundland",
   "metadata": {},
   "outputs": [],
   "source": [
    "h = Distributions.pdf(distγ, tx);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "boolean-kelly",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface(y[1,:], y[2, :], g, c=:OrRd, colorbar = :none)\n",
    "surface!(tx[1,:], tx[2, :], h, c=:PuBu, alpha = 0.4, colorbar = :none)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sporting-frost",
   "metadata": {},
   "source": [
    "The distributions of $T(X)$ and $Y$ have similar shapes. This can be improved by sampling more points but it requires a lot more computing ressources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dedicated-freeze",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.8.2",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}