1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
|
{
"nbformat_minor": 0,
"nbformat": 4,
"cells": [
{
"execution_count": null,
"cell_type": "code",
"source": [
"%matplotlib inline"
],
"outputs": [],
"metadata": {
"collapsed": false
}
},
{
"source": [
"\n# 1D Wasserstein barycenter demo\n\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 mpl_toolkits.mplot3d import Axes3D #necessary for 3d plot even if not used\nfrom matplotlib.collections import PolyCollection\n\n\n#%% parameters\n\nn=100 # nb bins\n\n# bin positions\nx=np.arange(n,dtype=np.float64)\n\n# Gaussian distributions\na1=ot.datasets.get_1D_gauss(n,m=20,s=5) # m= mean, s= std\na2=ot.datasets.get_1D_gauss(n,m=60,s=8)\n\n# creating matrix A containing all distributions\nA=np.vstack((a1,a2)).T\nnbd=A.shape[1]\n\n# loss matrix + normalization\nM=ot.utils.dist0(n)\nM/=M.max()\n\n#%% plot the distributions\n\npl.figure(1)\nfor i in range(nbd):\n pl.plot(x,A[:,i])\npl.title('Distributions')\n\n#%% barycenter computation\n\nalpha=0.2 # 0<=alpha<=1\nweights=np.array([1-alpha,alpha])\n\n# l2bary\nbary_l2=A.dot(weights)\n\n# wasserstein\nreg=1e-3\nbary_wass=ot.bregman.barycenter(A,M,reg,weights)\n\npl.figure(2)\npl.clf()\npl.subplot(2,1,1)\nfor i in range(nbd):\n pl.plot(x,A[:,i])\npl.title('Distributions')\n\npl.subplot(2,1,2)\npl.plot(x,bary_l2,'r',label='l2')\npl.plot(x,bary_wass,'g',label='Wasserstein')\npl.legend()\npl.title('Barycenters')\n\n\n#%% barycenter interpolation\n\nnbalpha=11\nalphalist=np.linspace(0,1,nbalpha)\n\n\nB_l2=np.zeros((n,nbalpha))\n\nB_wass=np.copy(B_l2)\n\nfor i in range(0,nbalpha):\n alpha=alphalist[i]\n weights=np.array([1-alpha,alpha])\n B_l2[:,i]=A.dot(weights)\n B_wass[:,i]=ot.bregman.barycenter(A,M,reg,weights)\n\n#%% plot interpolation\n\npl.figure(3,(10,5))\n\n#pl.subplot(1,2,1)\ncmap=pl.cm.get_cmap('viridis')\nverts = []\nzs = alphalist\nfor i,z in enumerate(zs):\n ys = B_l2[:,i]\n verts.append(list(zip(x, ys)))\n\nax = pl.gcf().gca(projection='3d')\n\npoly = PolyCollection(verts,facecolors=[cmap(a) for a in alphalist])\npoly.set_alpha(0.7)\nax.add_collection3d(poly, zs=zs, zdir='y')\n\nax.set_xlabel('x')\nax.set_xlim3d(0, n)\nax.set_ylabel('$\\\\alpha$')\nax.set_ylim3d(0,1)\nax.set_zlabel('')\nax.set_zlim3d(0, B_l2.max()*1.01)\npl.title('Barycenter interpolation with l2')\n\npl.show()\n\npl.figure(4,(10,5))\n\n#pl.subplot(1,2,1)\ncmap=pl.cm.get_cmap('viridis')\nverts = []\nzs = alphalist\nfor i,z in enumerate(zs):\n ys = B_wass[:,i]\n verts.append(list(zip(x, ys)))\n\nax = pl.gcf().gca(projection='3d')\n\npoly = PolyCollection(verts,facecolors=[cmap(a) for a in alphalist])\npoly.set_alpha(0.7)\nax.add_collection3d(poly, zs=zs, zdir='y')\n\nax.set_xlabel('x')\nax.set_xlim3d(0, n)\nax.set_ylabel('$\\\\alpha$')\nax.set_ylim3d(0,1)\nax.set_zlabel('')\nax.set_zlim3d(0, B_l2.max()*1.01)\npl.title('Barycenter interpolation with Wasserstein')\n\npl.show()"
],
"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"
}
}
}
}
|
