summaryrefslogtreecommitdiff
path: root/examples/plot_free_support_barycenter.py
blob: 274cf76d81dd4f33e514ea8150e359ca26c671e8 (plain)
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')