free support barycenter

1 files changed, 66 insertions, 0 deletions

diff --git a/examples/plot_free_support_barycenter.py b/examples/plot_free_support_barycenter.py
new file mode 100644
index 0000000..a733745
--- /dev/null
+++ b/examples/plot_free_support_barycenter.py
@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+"""
+====================================================
+2D Wasserstein barycenters between empirical distributions
+====================================================
+
+Illustration of 2D Wasserstein barycenters between discributions that are weighted
+sum of diracs.
+
+"""
+
+# Author: Vivien Seguy <vivien.seguy@iip.ist.i.kyoto-u.ac.jp>
+#
+# License: MIT License
+
+import numpy as np
+import matplotlib.pylab as pl
+import ot.plot
+
+
+##############################################################################
+# Generate data
+# -------------
+#%% parameters and data generation
+N = 4
+d = 2
+measures_locations = []
+measures_weights = []
+
+for i in range(N):
+
+    n = np.rand.int(low=1, high=20)  # nb samples
+
+    mu = np.random.normal(0., 1., (d,))
+    cov = np.random.normal(0., 1., (d,d))
+
+    xs = ot.datasets.make_2D_samples_gauss(n, mu, cov)
+    b = np.random.uniform(0., 1., n)
+    b = b/np.sum(b)
+
+    measures_locations.append(xs)
+    measures_weights.append(b)
+
+k = 10
+X_init = np.random.normal(0., 1., (k,d))
+b_init = np.ones((k,)) / k
+
+
+##############################################################################
+# Compute free support barycenter
+# -------------
+X = ot.lp.barycenter(measures_locations, measures_weights, X_init, b_init)
+
+
+##############################################################################
+# Plot data
+# ---------
+
+#%% plot samples
+
+pl.figure(1)
+for (xs, b) in zip(measures_locations, measures_weights):
+    pl.scatter(xs[:, 0], xs[:, 1], s=b, c=np.tile(np.rand(0. ,255., size=(3,)), (1,b.size(0))) , label='Data measures')
+pl.scatter(xs[:, 0], xs[:, 1], s=b, c='black' , label='2-Wasserstein barycenter')
+pl.legend(loc=0)
+pl.title('Data measures and their barycenter')