|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +r""" |
| 3 | +================================== |
| 4 | +Semi-discrete OT: a toy 2D problem |
| 5 | +================================== |
| 6 | +
|
| 7 | +This example shows the :mod:`ot.semidiscrete` solver on a small 2D problem: |
| 8 | +a uniform source on :math:`[0, 1]^2` and 15 random target atoms with uniform |
| 9 | +weights. With so few atoms the Laguerre cells can be drawn by brute force on |
| 10 | +a grid. |
| 11 | +
|
| 12 | +We call :func:`ot.semidiscrete.solve_semidiscrete` with its default |
| 13 | +arguments: the underlying algorithm is **Projected Averaged SGD**, and the |
| 14 | +default ``decreasing_reg=True`` adds the **DRAG** entropic-regularization |
| 15 | +schedule of [90]_, which improves convergence. |
| 16 | +
|
| 17 | +For the returned potential :math:`g` we report: |
| 18 | +
|
| 19 | +- the empirical Laguerre-cell masses (mean and max absolute deviation from |
| 20 | + :math:`1/15`); |
| 21 | +- the semi-dual objective |
| 22 | + :math:`\langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)]` estimated by |
| 23 | + Monte Carlo, where the c-transform |
| 24 | + :math:`\varphi_g(x) = \min_j\big(c(x, y_j) - g_j\big)` is computed by |
| 25 | + :func:`ot.semidiscrete.semidiscrete_c_transform`. The solver **maximises** this |
| 26 | + objective. |
| 27 | +
|
| 28 | +.. [90] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. |
| 29 | + (2025). *Decreasing Entropic Regularization Averaged Gradient for |
| 30 | + Semi-Discrete Optimal Transport.* NeurIPS 2025. |
| 31 | +""" |
| 32 | + |
| 33 | +# Author: Ferdinand Genans <genans.ferdinand@gmail.com> |
| 34 | +# |
| 35 | +# License: MIT License |
| 36 | + |
| 37 | +# sphinx_gallery_thumbnail_number = 1 |
| 38 | + |
| 39 | +import numpy as np |
| 40 | +import matplotlib.pyplot as plt |
| 41 | + |
| 42 | +from ot.semidiscrete import ( |
| 43 | + solve_semidiscrete, |
| 44 | + semidiscrete_atom_weights, |
| 45 | + semidiscrete_c_transform, |
| 46 | + semidiscrete_ot_map, |
| 47 | +) |
| 48 | + |
| 49 | +############################################################################## |
| 50 | +# Toy 2D problem |
| 51 | +# -------------- |
| 52 | + |
| 53 | +rng = np.random.default_rng(42) |
| 54 | + |
| 55 | + |
| 56 | +def source_sampler(batch_size): |
| 57 | + return rng.random((batch_size, 2)) |
| 58 | + |
| 59 | + |
| 60 | +n_atoms = 15 |
| 61 | +target_positions = 0.1 + 0.8 * np.random.default_rng(0).random((n_atoms, 2)) |
| 62 | + |
| 63 | + |
| 64 | +def plot_laguerre_cells(target, g, ax, title, resolution=300, alpha=0.55): |
| 65 | + xs = np.linspace(0, 1, resolution) |
| 66 | + ys = np.linspace(0, 1, resolution) |
| 67 | + XX, YY = np.meshgrid(xs, ys) |
| 68 | + grid = np.stack([XX.ravel(), YY.ravel()], axis=1) |
| 69 | + labels = semidiscrete_atom_weights(target, grid, g, reg=0.0).argmax(axis=1) |
| 70 | + image = labels.reshape(resolution, resolution) |
| 71 | + cmap = plt.get_cmap("tab20", target.shape[0]) |
| 72 | + ax.imshow( |
| 73 | + image, |
| 74 | + origin="lower", |
| 75 | + extent=(0, 1, 0, 1), |
| 76 | + cmap=cmap, |
| 77 | + alpha=alpha, |
| 78 | + vmin=-0.5, |
| 79 | + vmax=target.shape[0] - 0.5, |
| 80 | + interpolation="nearest", |
| 81 | + ) |
| 82 | + # Target points share the colour of their Laguerre cell. |
| 83 | + ax.scatter( |
| 84 | + target[:, 0], |
| 85 | + target[:, 1], |
| 86 | + s=80, |
| 87 | + c=[cmap(i) for i in range(target.shape[0])], |
| 88 | + edgecolor="black", |
| 89 | + linewidths=1.2, |
| 90 | + zorder=3, |
| 91 | + ) |
| 92 | + ax.set_title(title) |
| 93 | + ax.set_aspect("equal") |
| 94 | + ax.set_xlim(0, 1) |
| 95 | + ax.set_ylim(0, 1) |
| 96 | + |
| 97 | + |
| 98 | +############################################################################## |
| 99 | +# Solve and visualise |
| 100 | +# ------------------- |
| 101 | +# |
| 102 | +# A single call to :func:`solve_semidiscrete` runs DRAG with the default |
| 103 | +# arguments (``decreasing_reg=True``). We show the initial Voronoi cells |
| 104 | +# (:math:`g = 0`) next to the Laguerre cells at the optimum. |
| 105 | +# With the squared-Euclidean cost (default ``metric='sqeuclidean'``), the cost |
| 106 | +# between a source point in :math:`[0, 1]^2` and an atom is |
| 107 | +# :math:`\|x - y\|^2 \le 2`. We clip the potential to |
| 108 | +# ``[-max_cost, max_cost] = [-2, 2]``, the localizing set where an optimal |
| 109 | +# potential lies ([90]_, Lemma 1), which speeds up convergence. |
| 110 | +g_drag = solve_semidiscrete( |
| 111 | + target_positions, |
| 112 | + source_sampler, |
| 113 | + max_iter=20_000, |
| 114 | + batch_size=32, |
| 115 | + max_cost=2.0, |
| 116 | +) |
| 117 | + |
| 118 | +fig, axes = plt.subplots(1, 2, figsize=(11, 5.5)) |
| 119 | +plot_laguerre_cells(target_positions, np.zeros(n_atoms), axes[0], "Voronoi (g = 0)") |
| 120 | +plot_laguerre_cells(target_positions, g_drag, axes[1], "Approximated OT Laguerre cells") |
| 121 | +plt.tight_layout() |
| 122 | +plt.show() |
| 123 | + |
| 124 | + |
| 125 | +############################################################################## |
| 126 | +# Transport map over the Laguerre cells |
| 127 | +# ------------------------------------- |
| 128 | +# |
| 129 | +# :func:`semidiscrete_ot_map` with ``reg=0`` is the hard Monge map: every |
| 130 | +# source point is sent to the atom of its Laguerre cell. Overlaying the map |
| 131 | +# (arrows on a source grid) on the *faded* cells shows each cell's mass |
| 132 | +# collapsing onto its atom -- a direct illustration of the mapping function. |
| 133 | + |
| 134 | +gx = np.linspace(0.04, 0.96, 14) |
| 135 | +grid = np.stack([a.ravel() for a in np.meshgrid(gx, gx)], axis=1) |
| 136 | +mapped = semidiscrete_ot_map(target_positions, grid, g_drag, reg=0.0) |
| 137 | +labels = semidiscrete_atom_weights(target_positions, grid, g_drag, reg=0.0).argmax( |
| 138 | + axis=1 |
| 139 | +) |
| 140 | + |
| 141 | +cmap = plt.get_cmap("tab20", n_atoms) |
| 142 | +fig, ax = plt.subplots(figsize=(6.5, 6.5)) |
| 143 | +plot_laguerre_cells( |
| 144 | + target_positions, g_drag, ax, "Approximated OT map over Laguerre cells", alpha=0.22 |
| 145 | +) |
| 146 | +ax.quiver( |
| 147 | + grid[:, 0], |
| 148 | + grid[:, 1], |
| 149 | + mapped[:, 0] - grid[:, 0], |
| 150 | + mapped[:, 1] - grid[:, 1], |
| 151 | + angles="xy", |
| 152 | + scale_units="xy", |
| 153 | + scale=1, |
| 154 | + width=0.005, |
| 155 | + headwidth=4, |
| 156 | + headlength=5, |
| 157 | + color=[cmap(i) for i in labels], |
| 158 | + zorder=2, |
| 159 | +) |
| 160 | +plt.tight_layout() |
| 161 | +plt.show() |
| 162 | + |
| 163 | + |
| 164 | +############################################################################## |
| 165 | +# Cell masses and Monte Carlo cost |
| 166 | +# -------------------------------- |
| 167 | +# |
| 168 | +# At the optimum each Laguerre cell should carry mass :math:`1/15`. We report |
| 169 | +# the empirical mass error and the semi-dual objective |
| 170 | +# |
| 171 | +# .. math:: |
| 172 | +# \mathcal{S}(g) = \langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)] |
| 173 | +# |
| 174 | +# estimated by Monte Carlo. The solver maximises :math:`\mathcal{S}`. |
| 175 | + |
| 176 | + |
| 177 | +def cell_masses(target, g, sampler, n_samples=100_000): |
| 178 | + labels = semidiscrete_atom_weights(target, sampler(n_samples), g, reg=0.0).argmax( |
| 179 | + axis=1 |
| 180 | + ) |
| 181 | + counts = np.bincount(labels, minlength=target.shape[0]) |
| 182 | + return counts / n_samples |
| 183 | + |
| 184 | + |
| 185 | +def mc_cost(target, g, sampler, n_samples=100_000): |
| 186 | + b = np.full(target.shape[0], 1.0 / target.shape[0]) |
| 187 | + samples = sampler(n_samples) |
| 188 | + return float(g @ b + semidiscrete_c_transform(target, samples, g, reg=0.0).mean()) |
| 189 | + |
| 190 | + |
| 191 | +target_mass = 1.0 / n_atoms |
| 192 | +m_drag = cell_masses(target_positions, g_drag, source_sampler) |
| 193 | +cost_drag = mc_cost(target_positions, g_drag, source_sampler) |
| 194 | + |
| 195 | +print(f"Target mass per cell: {target_mass:.4f}") |
| 196 | +print( |
| 197 | + f"DRAG — mean abs. mass error: " |
| 198 | + f"{np.mean(np.abs(m_drag - target_mass)):.4f}" |
| 199 | + f" max: {np.max(np.abs(m_drag - target_mass)):.4f}" |
| 200 | + f" semi-dual cost (MC): {cost_drag:.5f}" |
| 201 | +) |
| 202 | + |
| 203 | + |
| 204 | +############################################################################## |
| 205 | +# Laguerre-cell masses |
| 206 | +# -------------------- |
| 207 | +# |
| 208 | +# At the optimum every cell carries the same mass :math:`1/15`. The bar plot |
| 209 | +# shows the empirical mass per cell against this ground truth (dashed line): |
| 210 | +# every cell sits close to the theoretical value. |
| 211 | + |
| 212 | +cmap = plt.get_cmap("tab20", n_atoms) |
| 213 | +fig, ax = plt.subplots(figsize=(7.5, 4)) |
| 214 | +ax.bar( |
| 215 | + np.arange(n_atoms), |
| 216 | + m_drag, |
| 217 | + color=[cmap(i) for i in range(n_atoms)], |
| 218 | + edgecolor="black", |
| 219 | + linewidth=0.6, |
| 220 | +) |
| 221 | +ax.axhline( |
| 222 | + target_mass, |
| 223 | + ls="--", |
| 224 | + color="black", |
| 225 | + lw=1.5, |
| 226 | + label="theoretical mass per cell at the optimum", |
| 227 | +) |
| 228 | +ax.set_ylim(0, 1.6 * target_mass) |
| 229 | +ax.set_xticks(np.arange(n_atoms)) |
| 230 | +ax.set_xlabel("atom index") |
| 231 | +ax.set_ylabel("Laguerre-cell mass") |
| 232 | +ax.set_title("Approximated OT: Laguerre-cell masses") |
| 233 | +ax.legend() |
| 234 | +plt.tight_layout() |
| 235 | +plt.show() |
0 commit comments