|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +""" |
| 3 | +====================================== |
| 4 | +Sinkhorn Divergence and Debiased OT solvers |
| 5 | +====================================== |
| 6 | +
|
| 7 | +This example shows how to use the debiased OT solvers in `ot.solve_sample` to |
| 8 | +compute Sinkhorn divergences and debiased Minibatch solutions. The debiased OT solvers |
| 9 | +can be used with balanced and unbalanced OT problems, and with different |
| 10 | +regularization types (entropic, L2, group lasso). |
| 11 | +""" |
| 12 | + |
| 13 | +# Author: Remi Flamary <remi.flamary@polytechnique.edu> |
| 14 | +# |
| 15 | +# License: MIT License |
| 16 | +# sphinx_gallery_thumbnail_number = 3 |
| 17 | + |
| 18 | +# %% |
| 19 | + |
| 20 | +import numpy as np |
| 21 | +import matplotlib.pylab as pl |
| 22 | +import ot |
| 23 | +import ot.plot |
| 24 | +from ot.datasets import make_1D_gauss as gauss |
| 25 | + |
| 26 | +############################################################################## |
| 27 | +# Generate data |
| 28 | +# ------------- |
| 29 | + |
| 30 | + |
| 31 | +# %% |
| 32 | +def sample_ball(n, radius=1.0, center=(0.0, 0.0)): |
| 33 | + np.random.seed(0) |
| 34 | + theta = 2 * np.pi * np.random.rand(n) |
| 35 | + r = radius * np.sqrt(np.random.rand(n)) |
| 36 | + |
| 37 | + x = r * np.cos(theta) + center[0] |
| 38 | + y = r * np.sin(theta) + center[1] |
| 39 | + |
| 40 | + return np.stack((x, y), axis=1) |
| 41 | + |
| 42 | + |
| 43 | +def sample_two_balls(n, radius=1.0, sep=1): |
| 44 | + assert n % 2 == 0, "n must be even" |
| 45 | + centers = ((-sep, -sep), (sep, sep)) |
| 46 | + n_half = n // 2 |
| 47 | + X1 = sample_ball(n_half, radius, centers[0]) |
| 48 | + X2 = sample_ball(n_half, radius, centers[1]) |
| 49 | + |
| 50 | + perm = np.random.permutation(n_half * 2) |
| 51 | + X = np.vstack((X1, X2)) |
| 52 | + X = X[perm] |
| 53 | + return X |
| 54 | + |
| 55 | + |
| 56 | +n = 50 |
| 57 | + |
| 58 | +x1 = sample_ball(n, radius=1.0, center=(0, 0)) |
| 59 | +x2 = sample_two_balls(n, radius=1.0, sep=1.5) |
| 60 | + |
| 61 | +pl.figure(1, figsize=(5, 5)) |
| 62 | +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", alpha=0.7) |
| 63 | +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", alpha=0.7) |
| 64 | +pl.legend() |
| 65 | +pl.title("Two distributions") |
| 66 | +ax = pl.axis() |
| 67 | + |
| 68 | + |
| 69 | +############################################################################## |
| 70 | +# Compute Sinkhorn divergence and visualize plans |
| 71 | +# ----------------------------------------------- |
| 72 | +# The Sinkhorn divergence is computed by setting the `debias` parameter to |
| 73 | +# `True` in the `ot.solve_sample` function. The resulting value is the Sinkhorn |
| 74 | +# divergence. The Sinkhorn divergences is computed as: |
| 75 | +# |
| 76 | +# .. math:: |
| 77 | +# S_\epsilon(\mu, \nu) = OT_\epsilon(\mu, \nu) - \frac{1}{2} OT_\epsilon(\mu, \mu) - \frac{1}{2} OT_\epsilon(\nu, \nu) |
| 78 | +# |
| 79 | +# The entropic OT plans for each of those terms can be accessed in the `log` |
| 80 | +# attribute of the result, and can be visualized using the |
| 81 | +# `ot.plot.plot2D_samples_mat` function. |
| 82 | + |
| 83 | +res = ot.solve_sample(x1, x2, reg=0.1, debias=True) |
| 84 | + |
| 85 | +print("Sinkhorn divergence: ", res.value) |
| 86 | + |
| 87 | +plan_11 = res.log["res_aa"].plan |
| 88 | +plan_12 = res.log["res_ab"].plan |
| 89 | +plan_22 = res.log["res_bb"].plan |
| 90 | + |
| 91 | +# |
| 92 | +pl.figure(2, figsize=(15, 5)) |
| 93 | + |
| 94 | +pl.subplot(1, 3, 1) |
| 95 | +ot.plot.plot2D_samples_mat(x1, x1, plan_11, thr=0.05) |
| 96 | +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) |
| 97 | +pl.axis(ax) |
| 98 | +pl.title("Plan between source and source") |
| 99 | +pl.subplot(1, 3, 2) |
| 100 | +ot.plot.plot2D_samples_mat(x1, x2, plan_12, thr=0.05) |
| 101 | +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) |
| 102 | +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", zorder=2) |
| 103 | +pl.axis(ax) |
| 104 | +pl.title("Plan between source and target") |
| 105 | +pl.subplot(1, 3, 3) |
| 106 | +ot.plot.plot2D_samples_mat(x2, x2, plan_22, thr=0.05) |
| 107 | +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", color="C1", zorder=2) |
| 108 | +pl.axis(ax) |
| 109 | +pl.title("Plan between target and target") |
| 110 | + |
| 111 | + |
| 112 | +############################################################################## |
| 113 | +# Debiased Minibatch OT |
| 114 | +# --------------------------------- |
| 115 | +# |
| 116 | +# Doing OT on minibatches leads to a similar bias than using entropic |
| 117 | +# regularization since the average OT plan is densified due to the stochasticity |
| 118 | +# of the minibatch sampling. On a given minibatch, the debiased loss can be |
| 119 | +# computed by setting the `debias` parameter to `'split'`that split the data |
| 120 | +# points in each distributions in two and computes the debias OT loss as: |
| 121 | +# |
| 122 | +# .. math:: |
| 123 | +# \tilde{OT}_m(\mu, \nu) = \frac{1}{2}(\hat{OT}_m(\mu_1, \nu_1) + \hat{OT}_m(\mu_2, \nu_2) - \hat{OT}_m(\nu_1, \nu_2) - \hat{OT}_m(\mu_1, \nu_2)) |
| 124 | +# |
| 125 | + |
| 126 | +# %% solve OT minibtach and visualize the plans |
| 127 | + |
| 128 | +res = ot.solve_sample(x1, x2, debias="split") |
| 129 | + |
| 130 | +print("Debiased minibatch OT loss: ", res.value) |
| 131 | + |
| 132 | +# recover the plans for each of the four terms in the debiased loss |
| 133 | +plan_11 = res.log["res_aa"].plan |
| 134 | +plan_12 = res.log["res_ab1"].plan |
| 135 | +plan_21 = res.log["res_ab2"].plan |
| 136 | +plan_22 = res.log["res_bb"].plan |
| 137 | +sel_a1 = res.log["sel_a1"] |
| 138 | +sel_a2 = res.log["sel_a2"] |
| 139 | +sel_b1 = res.log["sel_b1"] |
| 140 | +sel_b2 = res.log["sel_b2"] |
| 141 | + |
| 142 | +nb1 = plan_11.shape[0] |
| 143 | +nb2 = plan_22.shape[0] |
| 144 | + |
| 145 | +pl.figure(4, figsize=(15, 3)) |
| 146 | + |
| 147 | +pl.subplot(1, 4, 1) |
| 148 | +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="$\mu_1$", zorder=2) |
| 149 | +pl.scatter( |
| 150 | + x1[sel_a2, 0], x1[sel_a2, 1], label=r"$\mu_2$", zorder=2, color="C0", alpha=0.5 |
| 151 | +) |
| 152 | +pl.scatter(x2[sel_b1, 0], x2[sel_b1, 1], label=r"$\nu_1$", zorder=2, color="C1") |
| 153 | +pl.scatter( |
| 154 | + x2[sel_b2, 0], x2[sel_b2, 1], label=r"$\nu_2$", zorder=2, color="C1", alpha=0.5 |
| 155 | +) |
| 156 | +pl.title("Minibatch split") |
| 157 | +pl.axis(ax) |
| 158 | +pl.legend() |
| 159 | + |
| 160 | + |
| 161 | +pl.subplot(1, 4, 2) |
| 162 | +ot.plot.plot2D_samples_mat(x1[sel_a1], x1[sel_a2], plan_11, thr=0.05) |
| 163 | +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], zorder=2) |
| 164 | +pl.scatter( |
| 165 | + x1[sel_a2, 0], |
| 166 | + x1[sel_a2, 1], |
| 167 | + zorder=2, |
| 168 | + color="C0", |
| 169 | + alpha=0.5, |
| 170 | +) |
| 171 | +pl.axis(ax) |
| 172 | +pl.title("Plan between source and source") |
| 173 | +pl.subplot(1, 4, 3) |
| 174 | +ot.plot.plot2D_samples_mat(x1[sel_a1], x2[sel_b1], plan_12, thr=0.05) |
| 175 | +ot.plot.plot2D_samples_mat(x1[sel_a2], x2[sel_b2], plan_21, thr=0.05, alpha=0.5) |
| 176 | + |
| 177 | +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2) |
| 178 | +pl.scatter( |
| 179 | + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" |
| 180 | +) |
| 181 | +pl.scatter( |
| 182 | + x1[sel_a2, 0], |
| 183 | + x1[sel_a2, 1], |
| 184 | + label="Source distribution", |
| 185 | + zorder=2, |
| 186 | + color="C0", |
| 187 | + alpha=0.5, |
| 188 | +) |
| 189 | +pl.scatter( |
| 190 | + x2[sel_b2, 0], |
| 191 | + x2[sel_b2, 1], |
| 192 | + label="Target distribution", |
| 193 | + zorder=2, |
| 194 | + color="C1", |
| 195 | + alpha=0.5, |
| 196 | +) |
| 197 | +pl.axis(ax) |
| 198 | +pl.title("Plan between source and target") |
| 199 | + |
| 200 | +pl.subplot(1, 4, 4) |
| 201 | +ot.plot.plot2D_samples_mat(x2[sel_b1], x2[sel_b2], plan_22, thr=0.05) |
| 202 | +pl.scatter( |
| 203 | + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" |
| 204 | +) |
| 205 | +pl.scatter( |
| 206 | + x2[sel_b2, 0], |
| 207 | + x2[sel_b2, 1], |
| 208 | + label="Target distribution", |
| 209 | + zorder=2, |
| 210 | + color="C1", |
| 211 | + alpha=0.5, |
| 212 | +) |
| 213 | +pl.axis(ax) |
| 214 | +pl.title("Plan between target and target") |
| 215 | + |
| 216 | + |
| 217 | +############################################################################## |
| 218 | +# Comparison of the divergences |
| 219 | +# ------------------------------------------------- |
| 220 | + |
| 221 | +# %% move a distribution and compute Sinkhorn divergence and Sinkhorn distance |
| 222 | +reg = 0.1 |
| 223 | + |
| 224 | +sep_list = np.linspace(0, 1.0, 10) |
| 225 | +sink_list = [] |
| 226 | +sink_div_list = [] |
| 227 | +ot_mb_list = [] |
| 228 | +ot_mb_sink_list = [] |
| 229 | +for sep in sep_list: |
| 230 | + x2sep = sample_two_balls(n, radius=1.0, sep=sep) |
| 231 | + sink_list.append( |
| 232 | + ot.solve_sample( |
| 233 | + x1, |
| 234 | + x2sep, |
| 235 | + reg=reg, |
| 236 | + ).value |
| 237 | + ) |
| 238 | + sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value) |
| 239 | + ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value) |
| 240 | + |
| 241 | + ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value) |
| 242 | + |
| 243 | +pl.figure(3) |
| 244 | +pl.plot(sep_list, sink_list, label="Sinkhorn loss", color="C0") |
| 245 | +pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") |
| 246 | +pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2") |
| 247 | +pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3") |
| 248 | +pl.xlabel("Separation between distributions") |
| 249 | +pl.ylabel("Loss/Divergence") |
| 250 | +pl.title("Comparison of biased VS debiased OT losses") |
| 251 | +pl.grid() |
| 252 | +pl.legend() |
| 253 | + |
| 254 | +# %% |
0 commit comments