Semi-discrete OT: a toy 2D problem

This example shows the ot.semidiscrete solver on a small 2D problem: a uniform source on \([0, 1]^2\) and 15 random target atoms with uniform weights. With so few atoms the Laguerre cells can be drawn by brute force on a grid.

We call ot.semidiscrete.solve_semidiscrete() with its default arguments: the underlying algorithm is Projected Averaged SGD, and the default decreasing_reg=True adds the DRAG entropic-regularization schedule of [90], which improves convergence.

For the returned potential \(g\) we report:

the empirical Laguerre-cell masses (mean and max absolute deviation from \(1/15\));
the semi-dual objective \(\langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)]\) estimated by Monte Carlo, where the c-transform \(\varphi_g(x) = \min_j\big(c(x, y_j) - g_j\big)\) is computed by ot.semidiscrete.semidiscrete_c_transform(). The solver maximises this objective.

# Author: Ferdinand Genans <genans.ferdinand@gmail.com>
#
# License: MIT License

# sphinx_gallery_thumbnail_number = 1

import numpy as np
import matplotlib.pyplot as plt

from ot.semidiscrete import (
    solve_semidiscrete,
    semidiscrete_atom_weights,
    semidiscrete_c_transform,
    semidiscrete_ot_map,
)

Toy 2D problem

rng = np.random.default_rng(42)


def source_sampler(batch_size):
    return rng.random((batch_size, 2))


n_atoms = 15
target_positions = 0.1 + 0.8 * np.random.default_rng(0).random((n_atoms, 2))


def plot_laguerre_cells(target, g, ax, title, resolution=300, alpha=0.55):
    xs = np.linspace(0, 1, resolution)
    ys = np.linspace(0, 1, resolution)
    XX, YY = np.meshgrid(xs, ys)
    grid = np.stack([XX.ravel(), YY.ravel()], axis=1)
    labels = semidiscrete_atom_weights(target, grid, g, reg=0.0).argmax(axis=1)
    image = labels.reshape(resolution, resolution)
    cmap = plt.get_cmap("tab20", target.shape[0])
    ax.imshow(
        image,
        origin="lower",
        extent=(0, 1, 0, 1),
        cmap=cmap,
        alpha=alpha,
        vmin=-0.5,
        vmax=target.shape[0] - 0.5,
        interpolation="nearest",
    )
    # Target points share the colour of their Laguerre cell.
    ax.scatter(
        target[:, 0],
        target[:, 1],
        s=80,
        c=[cmap(i) for i in range(target.shape[0])],
        edgecolor="black",
        linewidths=1.2,
        zorder=3,
    )
    ax.set_title(title)
    ax.set_aspect("equal")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)

Solve and visualise

A single call to solve_semidiscrete() runs DRAG with the default arguments (decreasing_reg=True). We show the initial Voronoi cells (\(g = 0\)) next to the Laguerre cells at the optimum. With the squared-Euclidean cost (default metric='sqeuclidean'), the cost between a source point in \([0, 1]^2\) and an atom is \(\|x - y\|^2 \le 2\). We clip the potential to [-max_cost, max_cost] = [-2, 2], the localizing set where an optimal potential lies ([90], Lemma 1), which speeds up convergence.

g_drag = solve_semidiscrete(
    target_positions,
    source_sampler,
    max_iter=20_000,
    batch_size=32,
    max_cost=2.0,
)

fig, axes = plt.subplots(1, 2, figsize=(11, 5.5))
plot_laguerre_cells(target_positions, np.zeros(n_atoms), axes[0], "Voronoi (g = 0)")
plot_laguerre_cells(target_positions, g_drag, axes[1], "Approximated OT Laguerre cells")
plt.tight_layout()
plt.show()