|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +""" |
| 3 | +================================= |
| 4 | +Debiased Sinkhorn barycenter demo |
| 5 | +================================= |
| 6 | +
|
| 7 | +This example illustrates the computation of the debiased Sinkhorn barycenter |
| 8 | +as proposed in [37]_. |
| 9 | +
|
| 10 | +
|
| 11 | +.. [37] Janati, H., Cuturi, M., Gramfort, A. Proceedings of the 37th |
| 12 | + International Conference on Machine Learning, PMLR 119:4692-4701, 2020 |
| 13 | +""" |
| 14 | + |
| 15 | +# Author: Hicham Janati <[email protected]> |
| 16 | +# |
| 17 | +# License: MIT License |
| 18 | +# sphinx_gallery_thumbnail_number = 3 |
| 19 | + |
| 20 | +import os |
| 21 | +from pathlib import Path |
| 22 | + |
| 23 | +import numpy as np |
| 24 | +import matplotlib.pyplot as plt |
| 25 | + |
| 26 | +import ot |
| 27 | +from ot.bregman import (barycenter, barycenter_debiased, |
| 28 | + convolutional_barycenter2d, |
| 29 | + convolutional_barycenter2d_debiased) |
| 30 | + |
| 31 | +############################################################################## |
| 32 | +# Debiased barycenter of 1D Gaussians |
| 33 | +# ------------------------------------ |
| 34 | + |
| 35 | +#%% parameters |
| 36 | + |
| 37 | +n = 100 # nb bins |
| 38 | + |
| 39 | +# bin positions |
| 40 | +x = np.arange(n, dtype=np.float64) |
| 41 | + |
| 42 | +# Gaussian distributions |
| 43 | +a1 = ot.datasets.make_1D_gauss(n, m=20, s=5) # m= mean, s= std |
| 44 | +a2 = ot.datasets.make_1D_gauss(n, m=60, s=8) |
| 45 | + |
| 46 | +# creating matrix A containing all distributions |
| 47 | +A = np.vstack((a1, a2)).T |
| 48 | +n_distributions = A.shape[1] |
| 49 | + |
| 50 | +# loss matrix + normalization |
| 51 | +M = ot.utils.dist0(n) |
| 52 | +M /= M.max() |
| 53 | + |
| 54 | +#%% barycenter computation |
| 55 | + |
| 56 | +alpha = 0.2 # 0<=alpha<=1 |
| 57 | +weights = np.array([1 - alpha, alpha]) |
| 58 | + |
| 59 | +epsilons = [5e-3, 1e-2, 5e-2] |
| 60 | + |
| 61 | + |
| 62 | +bars = [barycenter(A, M, reg, weights) for reg in epsilons] |
| 63 | +bars_debiased = [barycenter_debiased(A, M, reg, weights) for reg in epsilons] |
| 64 | +labels = ["Sinkhorn barycenter", "Debiased barycenter"] |
| 65 | +colors = ["indianred", "gold"] |
| 66 | + |
| 67 | +f, axes = plt.subplots(1, len(epsilons), tight_layout=True, sharey=True, |
| 68 | + figsize=(12, 4), num=1) |
| 69 | +for ax, eps, bar, bar_debiased in zip(axes, epsilons, bars, bars_debiased): |
| 70 | + ax.plot(A[:, 0], color="k", ls="--", label="Input data", alpha=0.3) |
| 71 | + ax.plot(A[:, 1], color="k", ls="--", alpha=0.3) |
| 72 | + for data, label, color in zip([bar, bar_debiased], labels, colors): |
| 73 | + ax.plot(data, color=color, label=label, lw=2) |
| 74 | + ax.set_title(r"$\varepsilon = %.3f$" % eps) |
| 75 | +plt.legend() |
| 76 | +plt.show() |
| 77 | + |
| 78 | + |
| 79 | +############################################################################## |
| 80 | +# Debiased barycenter of 2D images |
| 81 | +# --------------------------------- |
| 82 | +this_file = os.path.realpath('__file__') |
| 83 | +data_path = os.path.join(Path(this_file).parent.parent.parent, 'data') |
| 84 | +f1 = 1 - plt.imread(os.path.join(data_path, 'heart.png'))[:, :, 2] |
| 85 | +f2 = 1 - plt.imread(os.path.join(data_path, 'duck.png'))[:, :, 2] |
| 86 | + |
| 87 | +A = np.asarray([f1, f2]) + 1e-2 |
| 88 | +A /= A.sum(axis=(1, 2))[:, None, None] |
| 89 | + |
| 90 | +############################################################################## |
| 91 | +# Display the input images |
| 92 | + |
| 93 | +fig, axes = plt.subplots(1, 2, figsize=(7, 4), num=2) |
| 94 | +for ax, img in zip(axes, A): |
| 95 | + ax.imshow(img, cmap="Greys") |
| 96 | + ax.axis("off") |
| 97 | +fig.tight_layout() |
| 98 | +plt.show() |
| 99 | + |
| 100 | + |
| 101 | +############################################################################## |
| 102 | +# Barycenter computation and visualization |
| 103 | +# ---------------------------------------- |
| 104 | +# |
| 105 | + |
| 106 | +bars_sinkhorn, bars_debiased = [], [] |
| 107 | +epsilons = [5e-3, 7e-3, 1e-2] |
| 108 | +for eps in epsilons: |
| 109 | + bar = convolutional_barycenter2d(A, eps) |
| 110 | + bar_debiased, log = convolutional_barycenter2d_debiased(A, eps, log=True) |
| 111 | + bars_sinkhorn.append(bar) |
| 112 | + bars_debiased.append(bar_debiased) |
| 113 | + |
| 114 | +titles = ["Sinkhorn", "Debiased"] |
| 115 | +all_bars = [bars_sinkhorn, bars_debiased] |
| 116 | +fig, axes = plt.subplots(2, 3, figsize=(8, 6), num=3) |
| 117 | +for jj, (method, ax_row, bars) in enumerate(zip(titles, axes, all_bars)): |
| 118 | + for ii, (ax, img, eps) in enumerate(zip(ax_row, bars, epsilons)): |
| 119 | + ax.imshow(img, cmap="Greys") |
| 120 | + if jj == 0: |
| 121 | + ax.set_title(r"$\varepsilon = %.3f$" % eps, fontsize=13) |
| 122 | + ax.set_xticks([]) |
| 123 | + ax.set_yticks([]) |
| 124 | + ax.spines['top'].set_visible(False) |
| 125 | + ax.spines['right'].set_visible(False) |
| 126 | + ax.spines['bottom'].set_visible(False) |
| 127 | + ax.spines['left'].set_visible(False) |
| 128 | + if ii == 0: |
| 129 | + ax.set_ylabel(method, fontsize=15) |
| 130 | +fig.tight_layout() |
| 131 | +plt.show() |
0 commit comments