{
  "nbformat_minor": 0, 
  "nbformat": 4, 
  "cells": [
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "%matplotlib inline"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "\n# 1D optimal transport\n\n\n@author: rflamary\n\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "import numpy as np\nimport matplotlib.pylab as pl\nimport ot\nfrom ot.datasets import get_1D_gauss as gauss\n\n\n#%% parameters\n\nn=100 # nb bins\nn_target=50 # nb target distributions\n\n\n# bin positions\nx=np.arange(n,dtype=np.float64)\n\nlst_m=np.linspace(20,90,n_target)\n\n# Gaussian distributions\na=gauss(n,m=20,s=5) # m= mean, s= std\n\nB=np.zeros((n,n_target))\n\nfor i,m in enumerate(lst_m):\n    B[:,i]=gauss(n,m=m,s=5)\n\n# loss matrix and normalization\nM=ot.dist(x.reshape((n,1)),x.reshape((n,1)),'euclidean')\nM/=M.max()\nM2=ot.dist(x.reshape((n,1)),x.reshape((n,1)),'sqeuclidean')\nM2/=M2.max()\n#%% plot the distributions\n\npl.figure(1)\npl.subplot(2,1,1)\npl.plot(x,a,'b',label='Source distribution')\npl.title('Source distribution')\npl.subplot(2,1,2)\npl.plot(x,B,label='Target distributions')\npl.title('Target distributions')\n\n#%% Compute and plot distributions and loss matrix\n\nd_emd=ot.emd2(a,B,M) # direct computation of EMD\nd_emd2=ot.emd2(a,B,M2)  # direct computation of EMD with loss M3\n\n\npl.figure(2)\npl.plot(d_emd,label='Euclidean EMD')\npl.plot(d_emd2,label='Squared Euclidean EMD')\npl.title('EMD distances')\npl.legend()\n\n#%%\nreg=1e-2\nd_sinkhorn=ot.sinkhorn(a,B,M,reg)\nd_sinkhorn2=ot.sinkhorn(a,B,M2,reg)\n\npl.figure(2)\npl.clf()\npl.plot(d_emd,label='Euclidean EMD')\npl.plot(d_emd2,label='Squared Euclidean EMD')\npl.plot(d_sinkhorn,'+',label='Euclidean Sinkhorn')\npl.plot(d_sinkhorn2,'+',label='Squared Euclidean Sinkhorn')\npl.title('EMD distances')\npl.legend()"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }
  ], 
  "metadata": {
    "kernelspec": {
      "display_name": "Python 2", 
      "name": "python2", 
      "language": "python"
    }, 
    "language_info": {
      "mimetype": "text/x-python", 
      "nbconvert_exporter": "python", 
      "name": "python", 
      "file_extension": ".py", 
      "version": "2.7.12", 
      "pygments_lexer": "ipython2", 
      "codemirror_mode": {
        "version": 2, 
        "name": "ipython"
      }
    }
  }
}