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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
|
# -*- coding: utf-8 -*-
"""
===========================================================
1D Wasserstein barycenter demo for Unbalanced distributions
===========================================================
This example illustrates the computation of regularized Wasserstein Barycenter
as proposed in [10] for Unbalanced inputs.
[10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
"""
# Author: Hicham Janati <hicham.janati@inria.fr>
#
# License: MIT License
# sphinx_gallery_thumbnail_number = 2
import numpy as np
import matplotlib.pylab as pl
import ot
# necessary for 3d plot even if not used
from mpl_toolkits.mplot3d import Axes3D # noqa
from matplotlib.collections import PolyCollection
##############################################################################
# Generate data
# -------------
# parameters
n = 100 # nb bins
# bin positions
x = np.arange(n, dtype=np.float64)
# Gaussian distributions
a1 = ot.datasets.make_1D_gauss(n, m=20, s=5) # m= mean, s= std
a2 = ot.datasets.make_1D_gauss(n, m=60, s=8)
# make unbalanced dists
a2 *= 3.0
# creating matrix A containing all distributions
A = np.vstack((a1, a2)).T
n_distributions = A.shape[1]
# loss matrix + normalization
M = ot.utils.dist0(n)
M /= M.max()
##############################################################################
# Plot data
# ---------
# plot the distributions
pl.figure(1, figsize=(6.4, 3))
for i in range(n_distributions):
pl.plot(x, A[:, i])
pl.title("Distributions")
pl.tight_layout()
##############################################################################
# Barycenter computation
# ----------------------
# non weighted barycenter computation
weight = 0.5 # 0<=weight<=1
weights = np.array([1 - weight, weight])
# l2bary
bary_l2 = A.dot(weights)
# wasserstein
reg = 1e-3
alpha = 1.0
bary_wass = ot.unbalanced.barycenter_unbalanced(A, M, reg, alpha, weights=weights)
pl.figure(2)
pl.clf()
pl.subplot(2, 1, 1)
for i in range(n_distributions):
pl.plot(x, A[:, i])
pl.title("Distributions")
pl.subplot(2, 1, 2)
pl.plot(x, bary_l2, "r", label="l2")
pl.plot(x, bary_wass, "g", label="Wasserstein")
pl.legend()
pl.title("Barycenters")
pl.tight_layout()
##############################################################################
# Barycentric interpolation
# -------------------------
# barycenter interpolation
n_weight = 11
weight_list = np.linspace(0, 1, n_weight)
B_l2 = np.zeros((n, n_weight))
B_wass = np.copy(B_l2)
for i in range(0, n_weight):
weight = weight_list[i]
weights = np.array([1 - weight, weight])
B_l2[:, i] = A.dot(weights)
B_wass[:, i] = ot.unbalanced.barycenter_unbalanced(
A, M, reg, alpha, weights=weights
)
# plot interpolation
pl.figure(3)
cmap = pl.get_cmap("viridis")
verts = []
zs = weight_list
for i, z in enumerate(zs):
ys = B_l2[:, i]
verts.append(list(zip(x, ys)))
ax = pl.gcf().add_subplot(projection="3d")
poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list])
poly.set_alpha(0.7)
ax.add_collection3d(poly, zs=zs, zdir="y")
ax.set_xlabel("x")
ax.set_xlim3d(0, n)
ax.set_ylabel(r"$\alpha$")
ax.set_ylim3d(0, 1)
ax.set_zlabel("")
ax.set_zlim3d(0, B_l2.max() * 1.01)
pl.title("Barycenter interpolation with l2")
pl.tight_layout()
pl.figure(4)
cmap = pl.get_cmap("viridis")
verts = []
zs = weight_list
for i, z in enumerate(zs):
ys = B_wass[:, i]
verts.append(list(zip(x, ys)))
ax = pl.gcf().add_subplot(projection="3d")
poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list])
poly.set_alpha(0.7)
ax.add_collection3d(poly, zs=zs, zdir="y")
ax.set_xlabel("x")
ax.set_xlim3d(0, n)
ax.set_ylabel(r"$\alpha$")
ax.set_ylim3d(0, 1)
ax.set_zlabel("")
ax.set_zlim3d(0, B_l2.max() * 1.01)
pl.title("Barycenter interpolation with Wasserstein")
pl.tight_layout()
pl.show()
|