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
55
56
57
58
59
60
61
62
63
64
65
|
# -*- 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.random.randint(low=1, high=20) # nb samples
mu = np.random.normal(0., 1., (d,))
cov = np.random.uniform(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.cvx.free_support_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.random.uniform(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')
|