From 2b5a589e54f6aab7d7fbd548fe245095bf8192ed Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 17:06:36 +0200 Subject: [PATCH 1/8] Add backend-agnostic semi-discrete OT solver (DRAG) Introduces ot.semidiscrete: Projected Averaged SGD on the semi-dual, with an optional decreasing entropic-regularization schedule (DRAG) from Genans et al. 2025. Works with NumPy, PyTorch, JAX, CuPy and TensorFlow via ot.backend. - ot/semidiscrete.py: solve_semidiscrete, atom_weights, ot_map, c_transform. Closed-form gradient, no autograd graph through the loop; quadratic cost by default with custom-callable override. - ot/__init__.py: register the new submodule. - test/test_semidiscrete.py: convergence on three toy problems with known optimal potentials, helper-function contracts (row- stochasticity of atom_weights, identity for c_transform at g=0, shape and finiteness of ot_map), and solver options (warm-start, projection, log, polyak_average off, entropic regime, custom cost). All tests parametrized over the nx fixture (NumPy + PyTorch). - examples/others/plot_semidiscrete.py: gallery example on a small 2D toy problem with Laguerre cells, empirical cell masses and a Monte Carlo estimate of the semi-dual cost. - RELEASES.md: new-features entry under 0.9.7.dev0. --- RELEASES.md | 4 + examples/others/plot_semidiscrete.py | 154 +++++++++++ ot/__init__.py | 2 + ot/semidiscrete.py | 277 +++++++++++++++++++ test/test_semidiscrete.py | 397 +++++++++++++++++++++++++++ 5 files changed, 834 insertions(+) create mode 100644 examples/others/plot_semidiscrete.py create mode 100644 ot/semidiscrete.py create mode 100644 test/test_semidiscrete.py diff --git a/RELEASES.md b/RELEASES.md index 0f8918cac..f540e3d14 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -14,6 +14,10 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397) - Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765) - Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765) +- Add backend-agnostic semi-discrete OT solver in `ot.semidiscrete` + (Projected Averaged SGD on the semi-dual, with optional decreasing + entropic-regularization schedule "DRAG" from Genans et al. 2025) with tests + and a gallery example (PR TBD) #### Closed issues diff --git a/examples/others/plot_semidiscrete.py b/examples/others/plot_semidiscrete.py new file mode 100644 index 000000000..b3d38bda7 --- /dev/null +++ b/examples/others/plot_semidiscrete.py @@ -0,0 +1,154 @@ +# -*- coding: utf-8 -*- +r""" +================================== +Semi-discrete OT: a toy 2D problem +================================== + +This example shows the :mod:`ot.semidiscrete` solver on a small 2D problem: +a uniform source on :math:`[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 :func:`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 [62]_, which improves convergence. + +For the returned potential :math:`g` we report: + +- the empirical Laguerre-cell masses (mean and max absolute deviation from + :math:`1/15`); +- the semi-dual objective + :math:`\langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)]` estimated by + Monte Carlo, where the c-transform + :math:`\varphi_g(x) = \min_j\big(c(x, y_j) - g_j\big)` is computed by + :func:`ot.semidiscrete.c_transform`. The solver **maximises** this + objective. + +.. [62] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). *Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport.* NeurIPS 2025. +""" + +# Author: Ferdinand Genans +# +# License: MIT License + +# sphinx_gallery_thumbnail_number = 1 + +import numpy as np +import matplotlib.pyplot as plt + +from ot.semidiscrete import ( + solve_semidiscrete, + atom_weights, + c_transform, +) + +############################################################################## +# 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): + 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 = 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=0.55, + 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 :func:`solve_semidiscrete` runs DRAG with the default +# arguments (``decreasing_reg=True``). We show the initial Voronoi cells +# (:math:`g = 0`) next to the Laguerre cells at the optimum. +g_drag = solve_semidiscrete( + target_positions, + source_sampler, + n_iter=20_000, + batch_size=16, + proj_bound=1.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], "DRAG") +plt.tight_layout() +plt.show() + + +############################################################################## +# Cell masses and Monte Carlo cost +# -------------------------------- +# +# At the optimum each Laguerre cell should carry mass :math:`1/15`. We report +# the empirical mass error and the semi-dual objective +# +# .. math:: +# \mathcal{S}(g) = \langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)] +# +# estimated by Monte Carlo. The solver maximises :math:`\mathcal{S}`. + + +def cell_masses(target, g, sampler, n_samples=100_000): + labels = atom_weights(target, sampler(n_samples), g, reg=0.0).argmax(axis=1) + counts = np.bincount(labels, minlength=target.shape[0]) + return counts / n_samples + + +def mc_cost(target, g, sampler, n_samples=100_000): + b = np.full(target.shape[0], 1.0 / target.shape[0]) + samples = sampler(n_samples) + return float(g @ b + c_transform(target, samples, g, reg=0.0).mean()) + + +target_mass = 1.0 / n_atoms +m_drag = cell_masses(target_positions, g_drag, source_sampler) +cost_drag = mc_cost(target_positions, g_drag, source_sampler) + +print(f"Target mass per cell: {target_mass:.4f}") +print( + f"DRAG — mean abs. mass error: " + f"{np.mean(np.abs(m_drag - target_mass)):.4f}" + f" max: {np.max(np.abs(m_drag - target_mass)):.4f}" + f" semi-dual cost (MC): {cost_drag:.5f}" +) diff --git a/ot/__init__.py b/ot/__init__.py index 75f17fed6..dd4e068b1 100644 --- a/ot/__init__.py +++ b/ot/__init__.py @@ -36,6 +36,7 @@ from . import gaussian from . import lowrank from . import gmm +from . import semidiscrete # OT functions from .lp import ( @@ -145,6 +146,7 @@ "factored", "lowrank", "gmm", + "semidiscrete", "binary_search_circle", "wasserstein_circle", "semidiscrete_wasserstein2_unif_circle", diff --git a/ot/semidiscrete.py b/ot/semidiscrete.py new file mode 100644 index 000000000..26975639e --- /dev/null +++ b/ot/semidiscrete.py @@ -0,0 +1,277 @@ +# -*- coding: utf-8 -*- +""" +Semi-discrete optimal transport: continuous source, discrete target. + +Backend-agnostic semi-dual solver based on the Projected Averaged SGD +of [1]_, with an optional decreasing entropic regularization schedule +(DRAG, [2]_). Works with any backend supported by :mod:`ot.backend` +(NumPy, PyTorch, JAX, CuPy, TensorFlow). + +References +---------- +.. [1] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). + "Stochastic Optimization in Semi-Discrete Optimal Transport: + Convergence Analysis and Minimax Rate." NeurIPS 2025. +.. [2] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). + "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. +""" + +# Author: Ferdinand Genans +# +# License: MIT License + +import math + +import numpy as np + +from .backend import get_backend + + +def _quadratic_cost(x, y, nx): + r"""Default cost: :math:`\tfrac{1}{2} \|x - y\|^2`.""" + x_sq = nx.sum(x**2, axis=1)[:, None] + y_sq = nx.sum(y**2, axis=1)[None, :] + cross = nx.einsum("ij,kj->ik", x, y) + return 0.5 * (x_sq + y_sq - 2.0 * cross) + + +def _setup(target_positions, target_weights, cost): + """Resolve backend, default weights and default cost.""" + nx = get_backend(target_positions) + m = target_positions.shape[0] + if target_weights is None: + target_weights = nx.full((m,), 1.0 / m, type_as=target_positions) + if cost is None: + + def cost(x, y): + return _quadratic_cost(x, y, nx) + + return nx, m, target_weights, nx.log(target_weights), cost + + +def _atom_weights(score, reg, log_b, nx): + """Row-stochastic weights ``(batch, m)`` from ``score = g - C``. + + Softmax of ``score / reg + log_b`` when ``reg > 0``, one-hot of + ``argmax(score, axis=1)`` when ``reg == 0``. + """ + if reg > 0: + log_w = score / reg + log_b[None, :] + log_w = log_w - nx.logsumexp(log_w, axis=1)[:, None] + return nx.exp(log_w) + m = score.shape[1] + idx = nx.argmax(score, axis=1) + arange_m = nx.from_numpy(np.arange(m), type_as=score) + mask = idx[:, None] == arange_m[None, :] + one = nx.full((1,), 1.0, type_as=score) + zero = nx.full((1,), 0.0, type_as=score) + return nx.where(mask, one, zero) + + +def atom_weights( + target_positions, + source_samples, + semi_dual_potential, + target_weights=None, + cost=None, + reg=0.0, +): + r"""Row-stochastic atom-assignment weights induced by ``semi_dual_potential``. + + Returns an array ``w`` of shape ``(n_samples, n_atoms)`` such that + ``w[i, j]`` is the (entropic) probability that sample ``x_i`` is + transported to atom ``y_j``. + """ + nx, _, _, log_b, cost_fn = _setup(target_positions, target_weights, cost) + score = semi_dual_potential[None, :] - cost_fn(source_samples, target_positions) + return _atom_weights(score, reg, log_b, nx) + + +def ot_map( + target_positions, + source_samples, + semi_dual_potential, + target_weights=None, + cost=None, + reg=0.0, +): + r"""Transport map :math:`T(x) = \sum_j w_j(x)\, y_j` induced by the potential.""" + w = atom_weights( + target_positions, + source_samples, + semi_dual_potential, + target_weights=target_weights, + cost=cost, + reg=reg, + ) + return w @ target_positions + + +def c_transform( + target_positions, + source_samples, + semi_dual_potential, + target_weights=None, + cost=None, + reg=0.0, +): + r"""Pointwise (entropic) c-transform of ``semi_dual_potential``. + + - ``reg == 0``: :math:`\varphi_g(x) = \min_j\, c(x, y_j) - g_j`. + - ``reg > 0``: :math:`\varphi_g(x) = -\varepsilon \log \sum_j b_j + \exp\!\big((g_j - c(x, y_j))/\varepsilon\big)`. + """ + nx, _, _, log_b, cost_fn = _setup(target_positions, target_weights, cost) + score = semi_dual_potential[None, :] - cost_fn(source_samples, target_positions) + if reg == 0: + return -nx.max(score, axis=1) + return -reg * nx.logsumexp(score / reg + log_b[None, :], axis=1) + + +def solve_semidiscrete( + target_positions, + source_sampler, + target_weights=None, + cost=None, + reg=0.0, + n_iter=10_000, + batch_size=32, + lr0=None, + lr_exponent=2.0 / 3.0, + init_potential=None, + decreasing_reg=True, + decreasing_reg_initial_eps=0.1, + decreasing_reg_exponent=0.5, + proj_bound=None, + polyak_average=True, + log=False, +): + r"""Solve semi-discrete OT by Polyak-averaged SGD on the semi-dual. + + Maximizes the semi-dual :math:`g \mapsto \langle g, b \rangle + \mathbb{E}_X[\varphi_g(X)]` + by averaged stochastic gradient ascent with projection and decreasing + regularization, which corresponds to the DRAG algorithm [1]_. + Here :math:`\varphi_g` denotes the (entropic) c-transform of :math:`g`, + + .. math:: + \varphi_g(x) = \begin{cases} + \min_j \big(c(x, y_j) - g_j\big) & \text{if } \mathrm{reg} = 0, \\ + -\varepsilon \log \sum_j b_j \exp\!\big((g_j - c(x, y_j))/\varepsilon\big) + & \text{if } \mathrm{reg} = \varepsilon > 0, + \end{cases} + + cf. :func:`c_transform`. + + With ``decreasing_reg=True`` the regularization at iteration ``t`` is + :math:`\varepsilon_t = \max(\text{reg},\, \varepsilon_0 / t^\alpha)` — large + at first for smoothness, then annealed towards ``reg``. This is the + DRAG schedule of [1]_. + + Parameters + ---------- + target_positions : array-like, shape (n_atoms, d) + Positions of the target atoms. The backend of this array drives + all subsequent computations. + source_sampler : callable + ``source_sampler(batch_size)`` returns a ``(batch_size, d)`` array + of source samples, in the same backend as ``target_positions``. + target_weights : array-like, shape (n_atoms,), optional + Atom weights. Defaults to uniform. + cost : callable, optional + ``cost(x, y)`` returns the ``(n_samples, n_atoms)`` cost matrix. + Defaults to ``0.5 * ||x - y||^2``. + reg : float, default=0.0 + Entropic regularization (target value when ``decreasing_reg=True``). + n_iter : int, default=10000 + batch_size : int, default=32 + lr0 : float, optional + Initial learning rate. Defaults to ``sqrt(n_atoms * batch_size)``. + lr_exponent : float, default=2/3 + Step size decays as ``lr0 / t**lr_exponent``. + init_potential : array-like, shape (n_atoms,), optional + Starting iterate; defaults to zero. Not mutated. + decreasing_reg : bool, default=True + Enable the DRAG decreasing-regularization schedule. + decreasing_reg_initial_eps : float, default=0.1 + Initial regularization in the DRAG schedule. + decreasing_reg_exponent : float, default=0.5 + Decay exponent of the DRAG schedule. + proj_bound : float, optional + If given, clip each iterate to ``[-proj_bound, proj_bound]``. + polyak_average : bool, default=True + If True, return the uniform average of the iterates; else the last. + log : bool, default=False + If True, also return a small ``dict`` with the last iterate. + + Returns + ------- + semi_dual_potential : array, shape (n_atoms,) + info : dict, optional + Returned only when ``log=True``. + + References + ---------- + .. [1] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). + "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. + + Examples + -------- + >>> import numpy as np + >>> from ot.semidiscrete import solve_semidiscrete + >>> rng = np.random.default_rng(0) + >>> target = np.linspace(0.0, 1.0, 10).reshape(-1, 1) + >>> g = solve_semidiscrete( + ... target, lambda b: rng.random((b, 1)), + ... n_iter=500, batch_size=32, proj_bound=1.0, + ... ) + """ + nx, m, b, log_b, cost_fn = _setup(target_positions, target_weights, cost) + + if init_potential is None: + g = nx.zeros((m,), type_as=target_positions) + else: + g = init_potential + nx.zeros((m,), type_as=target_positions) + + if lr0 is None: + lr0 = math.sqrt(m * batch_size) + + g_avg = nx.zeros((m,), type_as=target_positions) if polyak_average else None + + for t in range(1, n_iter + 1): + if decreasing_reg: + reg_t = max(reg, decreasing_reg_initial_eps / (t**decreasing_reg_exponent)) + else: + reg_t = reg + + x = source_sampler(batch_size) + score = g[None, :] - cost_fn(x, target_positions) + w = _atom_weights(score, reg_t, log_b, nx) + grad = nx.mean(w, axis=0) - b + + lr_t = lr0 / (t**lr_exponent) + g = g - lr_t * grad + if proj_bound is not None: + g = nx.clip(g, -proj_bound, proj_bound) + if polyak_average: + g_avg = g_avg + (g - g_avg) / t + + result = g_avg if polyak_average else g + if log: + return result, { + "n_iter": n_iter, + "batch_size": batch_size, + "proj_bound": proj_bound, + "polyak_average": polyak_average, + "last_potential": g, + } + return result + + +__all__ = [ + "solve_semidiscrete", + "atom_weights", + "ot_map", + "c_transform", +] diff --git a/test/test_semidiscrete.py b/test/test_semidiscrete.py new file mode 100644 index 000000000..e0dedcd3b --- /dev/null +++ b/test/test_semidiscrete.py @@ -0,0 +1,397 @@ +# -*- coding: utf-8 -*- +"""Tests for ``ot/semidiscrete.py``. + +We rely on three small toy problems whose optimal semi-dual potential is +known in closed form, and check that the solver converges close to that +optimum. All tests run across every POT backend (via the ``nx`` fixture). +""" + +# License: MIT License + +import numpy as np +import pytest + +from ot.semidiscrete import ( + atom_weights, + c_transform, + ot_map, + solve_semidiscrete, +) + +N_ITER = 2_000 +BATCH_SIZE = 16 +TOLERANCE = 0.05 + + +# --------------------------------------------------------------------- +# Three toy problems with known optimal potentials. +# Each builder returns numpy arrays so we can lift them onto any backend. +# --------------------------------------------------------------------- + + +def regular_grid_problem(): + """Uniform source on ``[0, 1]^3``, 10 target atoms regularly placed on axis 0. + + By symmetry the optimal centred potential is zero on every atom. + """ + m, d = 10, 3 + target = np.zeros((m, d)) + target[:, 0] = (np.arange(m) + 0.5) / m + target[:, 1:] = 0.5 + weights = np.full(m, 1.0 / m) + optimal = np.zeros(m) + return target, weights, optimal, 1.0, d, "uniform_cube" + + +def nonuniform_weights_problem(): + """Uniform source on ``[0, 1]^3``, fixed support, nonuniform weights. + + The weights were computed by Monte Carlo so that the optimal centred + potential is exactly ``optimal_potential`` (used as a regression target). + """ + d = 3 + target = np.array( + [ + [0.54488318, 0.4236548, 0.64589411], + [0.77815675, 0.87001215, 0.97861834], + [0.11827443, 0.63992102, 0.14335329], + [0.94466892, 0.52184832, 0.41466194], + ] + ) + weights = np.array([0.3806643, 0.012264, 0.5503486, 0.0567231]) + optimal_potential = np.array([0.77423369, 0.61209572, 0.94374808, 0.6818203]) + return target, weights, optimal_potential, 1.0, d, "uniform_cube" + + +def shifted_1d_problem(): + r"""Uniform source on ``[delta, 1 + delta]``, atoms at ``k/m`` on the line. + + For this 1D problem with uniform target weights the optimal potential + has the closed form + + .. math:: + g^*_j = j \left( \frac{1}{2 m^2} - \frac{\delta}{m} \right), + \qquad j = 0, \dots, m - 1. + + See Appendix E.1 in "Stochastic Optimization in Semi-Discrete Optimal + Transport: Convergence Analysis and Minimax Rate", Genans et al. + (NeurIPS 2025). + """ + m, delta = 10, 0.5 + target = np.linspace(1 / m, 1.0, m).reshape(m, 1) + weights = np.full(m, 1.0 / m) + optimal = np.zeros(m) + for j in range(1, m): + optimal[j] = optimal[j - 1] + 1 / (2 * m * m) - delta / m + return target, weights, optimal, 1.0 + delta, 1, ("shifted_1d", delta) + + +ALL_PROBLEMS = [ + pytest.param(regular_grid_problem, id="regular_grid"), + pytest.param(nonuniform_weights_problem, id="nonuniform_weights"), + pytest.param(shifted_1d_problem, id="shifted_1d"), +] + + +# --------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------- + + +def make_sampler(kind, d, nx, target): + """Return a backend-aware sampler with a fixed numpy RNG inside.""" + rng = np.random.default_rng(0) + if kind == "uniform_cube": + + def sampler(b): + return nx.from_numpy(rng.random((b, d)), type_as=target) + + else: + # ("shifted_1d", delta) + delta = kind[1] + + def sampler(b): + return nx.from_numpy(delta + rng.random((b, d)), type_as=target) + + return sampler + + +def centered_l2_error(estimated, reference): + """L2 distance between the centred (mean-zero) potentials. + + The semi-dual is invariant to a global additive constant, so we factor that out. + """ + estimated = estimated - estimated.mean() + reference = reference - reference.mean() + return float(np.linalg.norm(estimated - reference)) + + +def lift(nx, target_np, weights_np): + """Move numpy ``target`` and ``weights`` arrays onto backend ``nx``.""" + target = nx.from_numpy(target_np) + weights = nx.from_numpy(weights_np, type_as=target) + return target, weights + + +# --------------------------------------------------------------------- +# Convergence on every backend +# --------------------------------------------------------------------- + + +@pytest.mark.parametrize("build_problem", ALL_PROBLEMS) +def test_solve_converges(nx, build_problem): + """Plain SGD (no decreasing reg) reaches the optimum on every toy problem.""" + target_np, weights_np, optimal, max_cost, d, kind = build_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + g = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=N_ITER, + batch_size=BATCH_SIZE, + decreasing_reg=False, + proj_bound=max_cost, + ) + err = centered_l2_error(nx.to_numpy(g), optimal) + assert err < TOLERANCE, f"err={err:.4f}" + + +@pytest.mark.parametrize("build_problem", ALL_PROBLEMS) +def test_drag_converges(nx, build_problem): + """DRAG (decreasing entropic reg) reaches the optimum on every toy problem.""" + target_np, weights_np, optimal, max_cost, d, kind = build_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + g = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=N_ITER, + batch_size=BATCH_SIZE, + decreasing_reg=True, + proj_bound=max_cost, + ) + err = centered_l2_error(nx.to_numpy(g), optimal) + assert err < TOLERANCE, f"err={err:.4f}" + + +# --------------------------------------------------------------------- +# Entropic regime +# --------------------------------------------------------------------- + + +def test_entropic_solver_runs(nx): + """In the entropic regime, the solver and ``c_transform`` produce finite values.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + g = solve_semidiscrete( + target, + sampler, + target_weights=weights, + reg=0.05, + n_iter=N_ITER, + batch_size=BATCH_SIZE, + proj_bound=max_cost, + ) + samples = sampler(64) + phi = c_transform(target, samples, g, target_weights=weights, reg=0.05) + assert np.isfinite(nx.to_numpy(g)).all() + assert np.isfinite(nx.to_numpy(phi)).all() + + +# --------------------------------------------------------------------- +# Custom cost +# --------------------------------------------------------------------- + + +def test_custom_quadratic_cost_matches_default(nx): + """A user-supplied quadratic cost reaches the same optimum as the default.""" + target_np, weights_np, optimal, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + def cost(x, y): + diff = x[:, None, :] - y[None, :, :] + return 0.5 * nx.sum(diff**2, axis=2) + + g = solve_semidiscrete( + target, + sampler, + target_weights=weights, + cost=cost, + n_iter=N_ITER, + batch_size=BATCH_SIZE, + proj_bound=max_cost, + ) + err = centered_l2_error(nx.to_numpy(g), optimal) + assert err < TOLERANCE, f"err={err:.4f}" + + +# --------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------- + + +@pytest.mark.parametrize("reg", [0.0, 0.1]) +def test_atom_weights_are_row_stochastic(nx, reg): + """``atom_weights`` returns nonnegative weights that sum to 1 per row.""" + target_np, weights_np, _, _, d, kind = nonuniform_weights_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + samples = sampler(32) + g = nx.zeros((target_np.shape[0],), type_as=target) + w = atom_weights(target, samples, g, target_weights=weights, reg=reg) + w_np = nx.to_numpy(w) + assert w_np.shape == (32, target_np.shape[0]) + assert (w_np >= 0).all() + np.testing.assert_allclose(w_np.sum(axis=1), 1.0, atol=1e-10) + + +def test_ot_map_shape_and_finiteness(nx): + """``ot_map`` returns finite values with the source-sample shape.""" + target_np, weights_np, _, _, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + samples = sampler(16) + g = nx.zeros((target_np.shape[0],), type_as=target) + transported = ot_map(target, samples, g, target_weights=weights) + transported_np = nx.to_numpy(transported) + samples_np = nx.to_numpy(samples) + assert transported_np.shape == samples_np.shape + assert np.isfinite(transported_np).all() + + +def test_c_transform_minimum_for_zero_potential(nx): + """At ``g = 0``, ``phi_g(x) = -max_j(-c(x, y_j)) = min_j c(x, y_j)``.""" + target_np, weights_np, _, _, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + samples = sampler(8) + g = nx.zeros((target_np.shape[0],), type_as=target) + phi = c_transform(target, samples, g, target_weights=weights) + samples_np = nx.to_numpy(samples) + target_np = nx.to_numpy(target) + diff = samples_np[:, None, :] - target_np[None, :, :] + expected = 0.5 * (diff**2).sum(axis=2).min(axis=1) + np.testing.assert_allclose(nx.to_numpy(phi), expected, atol=1e-10) + + +# --------------------------------------------------------------------- +# Solver options +# --------------------------------------------------------------------- + + +def test_warm_start_converges(nx): + """Splitting one run into two warm-started halves still converges.""" + target_np, weights_np, optimal, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + half = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=N_ITER // 2, + batch_size=BATCH_SIZE, + proj_bound=max_cost, + ) + g = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=N_ITER // 2, + batch_size=BATCH_SIZE, + init_potential=half, + proj_bound=max_cost, + ) + err = centered_l2_error(nx.to_numpy(g), optimal) + assert err < TOLERANCE, f"err={err:.4f}" + + +def test_init_potential_is_not_mutated(nx): + """The ``init_potential`` array passed by the caller is left intact.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + init_np = np.full(target_np.shape[0], 0.5) + init = nx.from_numpy(init_np, type_as=target) + snapshot = nx.to_numpy(init).copy() + + solve_semidiscrete( + target, + sampler, + target_weights=weights, + init_potential=init, + n_iter=10, + batch_size=4, + proj_bound=max_cost, + ) + np.testing.assert_array_equal(nx.to_numpy(init), snapshot) + + +def test_projection_clamps_last_iterate(nx): + """With ``proj_bound=b``, every coordinate of the last iterate lies in ``[-b, b]``.""" + target_np, weights_np, _, _, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + bound = 0.05 + _, info = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=300, + batch_size=4, + proj_bound=bound, + log=True, + ) + last = nx.to_numpy(info["last_potential"]) + assert np.abs(last).max() <= bound + 1e-10 + + +def test_polyak_average_off_returns_last_iterate(nx): + """With ``polyak_average=False`` the returned potential equals the last iterate.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + final, info = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=20, + batch_size=4, + polyak_average=False, + proj_bound=max_cost, + log=True, + ) + np.testing.assert_array_equal( + nx.to_numpy(final), nx.to_numpy(info["last_potential"]) + ) + + +def test_log_returns_metadata(nx): + """``log=True`` returns an info dict with the expected fields.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = lift(nx, target_np, weights_np) + sampler = make_sampler(kind, d, nx, target) + + g, info = solve_semidiscrete( + target, + sampler, + target_weights=weights, + n_iter=50, + batch_size=4, + proj_bound=max_cost, + log=True, + ) + assert nx.to_numpy(g).shape == (target_np.shape[0],) + assert info["n_iter"] == 50 + assert info["batch_size"] == 4 + assert info["proj_bound"] == max_cost From 835c959c51b72e9b123fc634db40c03b76c518b0 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 18:22:19 +0200 Subject: [PATCH 2/8] modified: README.md modified: RELEASES.md modified: examples/others/plot_semidiscrete.py Final small doc modifications. --- README.md | 2 ++ RELEASES.md | 5 +---- examples/others/plot_semidiscrete.py | 4 ++-- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 3e69af448..dfce350f9 100644 --- a/README.md +++ b/README.md @@ -451,3 +451,5 @@ Artificial Intelligence. [81] Xu, H., Luo, D., & Carin, L. (2019). [Scalable Gromov-Wasserstein learning for graph partitioning and matching](https://proceedings.neurips.cc/paper/2019/hash/6e62a992c676f611616097dbea8ea030-Abstract.html). Neural Information Processing Systems (NeurIPS). [82] Bonet, C., Nadjahi, K., Séjourné, T., Fatras, K., & Courty, N. (2024). [Slicing Unbalanced Optimal Transport](https://openreview.net/forum?id=AjJTg5M0r8). Transactions on Machine Learning Research. + +[83] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2026). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. diff --git a/RELEASES.md b/RELEASES.md index f540e3d14..b5e6402cf 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -14,10 +14,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397) - Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765) - Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765) -- Add backend-agnostic semi-discrete OT solver in `ot.semidiscrete` - (Projected Averaged SGD on the semi-dual, with optional decreasing - entropic-regularization schedule "DRAG" from Genans et al. 2025) with tests - and a gallery example (PR TBD) +- Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR TBD) #### Closed issues diff --git a/examples/others/plot_semidiscrete.py b/examples/others/plot_semidiscrete.py index b3d38bda7..cea748fc2 100644 --- a/examples/others/plot_semidiscrete.py +++ b/examples/others/plot_semidiscrete.py @@ -12,7 +12,7 @@ We call :func:`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 [62]_, which improves convergence. +schedule of [83]_, which improves convergence. For the returned potential :math:`g` we report: @@ -25,7 +25,7 @@ :func:`ot.semidiscrete.c_transform`. The solver **maximises** this objective. -.. [62] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. +.. [83] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. (2025). *Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport.* NeurIPS 2025. """ From f8549f082a06173b20f46b82c7791a2a298cc4c5 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 19:03:30 +0200 Subject: [PATCH 3/8] Modified "proj_bound" to "max_cost" in the solve_semidiscrete, and added a more detailed explanation of the effect of this argument on convergence in the example scipt. --- examples/others/plot_semidiscrete.py | 2 ++ ot/semidiscrete.py | 14 +++++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/others/plot_semidiscrete.py b/examples/others/plot_semidiscrete.py index cea748fc2..347becd8f 100644 --- a/examples/others/plot_semidiscrete.py +++ b/examples/others/plot_semidiscrete.py @@ -101,6 +101,8 @@ def plot_laguerre_cells(target, g, ax, title, resolution=300): # A single call to :func:`solve_semidiscrete` runs DRAG with the default # arguments (``decreasing_reg=True``). We show the initial Voronoi cells # (:math:`g = 0`) next to the Laguerre cells at the optimum. +# In this problem, the maximum cost between samples is 1.0, so we pass it in +# ``decreasing_reg=True``. Knowing this bound speeds up convergence. g_drag = solve_semidiscrete( target_positions, source_sampler, diff --git a/ot/semidiscrete.py b/ot/semidiscrete.py index 26975639e..64ed5d1d1 100644 --- a/ot/semidiscrete.py +++ b/ot/semidiscrete.py @@ -143,7 +143,7 @@ def solve_semidiscrete( decreasing_reg=True, decreasing_reg_initial_eps=0.1, decreasing_reg_exponent=0.5, - proj_bound=None, + max_cost=None, polyak_average=True, log=False, ): @@ -197,8 +197,8 @@ def solve_semidiscrete( Initial regularization in the DRAG schedule. decreasing_reg_exponent : float, default=0.5 Decay exponent of the DRAG schedule. - proj_bound : float, optional - If given, clip each iterate to ``[-proj_bound, proj_bound]``. + max_cost : float, optional + If given, clip each iterate to ``[-max_cost, max_cost]``. polyak_average : bool, default=True If True, return the uniform average of the iterates; else the last. log : bool, default=False @@ -224,7 +224,7 @@ def solve_semidiscrete( >>> target = np.linspace(0.0, 1.0, 10).reshape(-1, 1) >>> g = solve_semidiscrete( ... target, lambda b: rng.random((b, 1)), - ... n_iter=500, batch_size=32, proj_bound=1.0, + ... n_iter=500, batch_size=32, max_cost=1.0, ... ) """ nx, m, b, log_b, cost_fn = _setup(target_positions, target_weights, cost) @@ -252,8 +252,8 @@ def solve_semidiscrete( lr_t = lr0 / (t**lr_exponent) g = g - lr_t * grad - if proj_bound is not None: - g = nx.clip(g, -proj_bound, proj_bound) + if max_cost is not None: + g = nx.clip(g, -max_cost, max_cost) if polyak_average: g_avg = g_avg + (g - g_avg) / t @@ -262,7 +262,7 @@ def solve_semidiscrete( return result, { "n_iter": n_iter, "batch_size": batch_size, - "proj_bound": proj_bound, + "max_cost": max_cost, "polyak_average": polyak_average, "last_potential": g, } From b5383f5291284a4dcd8013a9f06679407574d249 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 19:04:58 +0200 Subject: [PATCH 4/8] Edited the semi-discrete example script "max_cost" explanation. --- examples/others/plot_semidiscrete.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/others/plot_semidiscrete.py b/examples/others/plot_semidiscrete.py index 347becd8f..69d475c31 100644 --- a/examples/others/plot_semidiscrete.py +++ b/examples/others/plot_semidiscrete.py @@ -101,14 +101,16 @@ def plot_laguerre_cells(target, g, ax, title, resolution=300): # A single call to :func:`solve_semidiscrete` runs DRAG with the default # arguments (``decreasing_reg=True``). We show the initial Voronoi cells # (:math:`g = 0`) next to the Laguerre cells at the optimum. -# In this problem, the maximum cost between samples is 1.0, so we pass it in -# ``decreasing_reg=True``. Knowing this bound speeds up convergence. +# In this problem, the maximum cost between samples is 1.0, so we pass it as +# ``max_cost=1.0``. Knowing this bound, the potential values are clipped to +# [-max_cost, max_cost], where it is known that an optimal potential lies ([83]_, Lemma 1), +# which speeds up convergence. g_drag = solve_semidiscrete( target_positions, source_sampler, n_iter=20_000, batch_size=16, - proj_bound=1.0, + max_cost=1.0, ) fig, axes = plt.subplots(1, 2, figsize=(11, 5.5)) From 3b8b330940fe602002e6e111be687b80620c2440 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 19:52:37 +0200 Subject: [PATCH 5/8] Changed "proj_bound" to "max_cost" in test_semidiscrete.py --- test/test_semidiscrete.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/test/test_semidiscrete.py b/test/test_semidiscrete.py index e0dedcd3b..3ff048a1b 100644 --- a/test/test_semidiscrete.py +++ b/test/test_semidiscrete.py @@ -152,7 +152,7 @@ def test_solve_converges(nx, build_problem): n_iter=N_ITER, batch_size=BATCH_SIZE, decreasing_reg=False, - proj_bound=max_cost, + max_cost=max_cost, ) err = centered_l2_error(nx.to_numpy(g), optimal) assert err < TOLERANCE, f"err={err:.4f}" @@ -172,7 +172,7 @@ def test_drag_converges(nx, build_problem): n_iter=N_ITER, batch_size=BATCH_SIZE, decreasing_reg=True, - proj_bound=max_cost, + max_cost=max_cost, ) err = centered_l2_error(nx.to_numpy(g), optimal) assert err < TOLERANCE, f"err={err:.4f}" @@ -196,7 +196,7 @@ def test_entropic_solver_runs(nx): reg=0.05, n_iter=N_ITER, batch_size=BATCH_SIZE, - proj_bound=max_cost, + max_cost=max_cost, ) samples = sampler(64) phi = c_transform(target, samples, g, target_weights=weights, reg=0.05) @@ -226,7 +226,7 @@ def cost(x, y): cost=cost, n_iter=N_ITER, batch_size=BATCH_SIZE, - proj_bound=max_cost, + max_cost=max_cost, ) err = centered_l2_error(nx.to_numpy(g), optimal) assert err < TOLERANCE, f"err={err:.4f}" @@ -298,7 +298,7 @@ def test_warm_start_converges(nx): target_weights=weights, n_iter=N_ITER // 2, batch_size=BATCH_SIZE, - proj_bound=max_cost, + max_cost=max_cost, ) g = solve_semidiscrete( target, @@ -307,7 +307,7 @@ def test_warm_start_converges(nx): n_iter=N_ITER // 2, batch_size=BATCH_SIZE, init_potential=half, - proj_bound=max_cost, + max_cost=max_cost, ) err = centered_l2_error(nx.to_numpy(g), optimal) assert err < TOLERANCE, f"err={err:.4f}" @@ -330,13 +330,13 @@ def test_init_potential_is_not_mutated(nx): init_potential=init, n_iter=10, batch_size=4, - proj_bound=max_cost, + max_cost=max_cost, ) np.testing.assert_array_equal(nx.to_numpy(init), snapshot) def test_projection_clamps_last_iterate(nx): - """With ``proj_bound=b``, every coordinate of the last iterate lies in ``[-b, b]``.""" + """With ``max_cost=b``, every coordinate of the last iterate lies in ``[-b, b]``.""" target_np, weights_np, _, _, d, kind = regular_grid_problem() target, weights = lift(nx, target_np, weights_np) sampler = make_sampler(kind, d, nx, target) @@ -348,7 +348,7 @@ def test_projection_clamps_last_iterate(nx): target_weights=weights, n_iter=300, batch_size=4, - proj_bound=bound, + max_cost=bound, log=True, ) last = nx.to_numpy(info["last_potential"]) @@ -368,7 +368,7 @@ def test_polyak_average_off_returns_last_iterate(nx): n_iter=20, batch_size=4, polyak_average=False, - proj_bound=max_cost, + max_cost=max_cost, log=True, ) np.testing.assert_array_equal( @@ -388,10 +388,10 @@ def test_log_returns_metadata(nx): target_weights=weights, n_iter=50, batch_size=4, - proj_bound=max_cost, + max_cost=max_cost, log=True, ) assert nx.to_numpy(g).shape == (target_np.shape[0],) assert info["n_iter"] == 50 assert info["batch_size"] == 4 - assert info["proj_bound"] == max_cost + assert info["max_cost"] == max_cost From bb15a066e22f5589f7aea1c5661efb4221b207f2 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Sun, 17 May 2026 20:08:14 +0200 Subject: [PATCH 6/8] Add PR's number . --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index b5e6402cf..be95d5b57 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -14,7 +14,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397) - Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765) - Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765) -- Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR TBD) +- Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR #812) #### Closed issues From 0da9fc7cb7cf7d115e2fbae3368a93cfe2dcce61 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Mon, 29 Jun 2026 15:28:50 +0200 Subject: [PATCH 7/8] Address review on semi-discrete OT module: - Rename public functions to semidiscrete_{atom_weights,ot_map,c_transform} and align params with ot.solve_sample (X_target, sampler_source, a_target, metric, max_iter, max_cost) - metric accepts ot.dist strings (default 'sqeuclidean') or a callable - sampler_source accepts built-in strings ('unif', 'ball', 'normal', ...), default 'unif' - Full docstrings with equations and references to both papers - Examplel OT-map visualization with arrow and cell-mass plot - Tests: use nx.from_numpy, cover the string metric/sampler paths --- examples/others/plot_semidiscrete.py | 109 ++++++-- ot/semidiscrete.py | 384 ++++++++++++++++++++------- test/test_semidiscrete.py | 175 ++++++++---- 3 files changed, 500 insertions(+), 168 deletions(-) diff --git a/examples/others/plot_semidiscrete.py b/examples/others/plot_semidiscrete.py index dd6d8643a..ea58665ab 100644 --- a/examples/others/plot_semidiscrete.py +++ b/examples/others/plot_semidiscrete.py @@ -22,7 +22,7 @@ :math:`\langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)]` estimated by Monte Carlo, where the c-transform :math:`\varphi_g(x) = \min_j\big(c(x, y_j) - g_j\big)` is computed by - :func:`ot.semidiscrete.c_transform`. The solver **maximises** this + :func:`ot.semidiscrete.semidiscrete_c_transform`. The solver **maximises** this objective. .. [90] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. @@ -41,8 +41,9 @@ from ot.semidiscrete import ( solve_semidiscrete, - atom_weights, - c_transform, + semidiscrete_atom_weights, + semidiscrete_c_transform, + semidiscrete_ot_map, ) ############################################################################## @@ -60,12 +61,12 @@ def source_sampler(batch_size): 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): +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 = atom_weights(target, grid, g, reg=0.0).argmax(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( @@ -73,7 +74,7 @@ def plot_laguerre_cells(target, g, ax, title, resolution=300): origin="lower", extent=(0, 1, 0, 1), cmap=cmap, - alpha=0.55, + alpha=alpha, vmin=-0.5, vmax=target.shape[0] - 0.5, interpolation="nearest", @@ -101,21 +102,61 @@ def plot_laguerre_cells(target, g, ax, title, resolution=300): # A single call to :func:`solve_semidiscrete` runs DRAG with the default # arguments (``decreasing_reg=True``). We show the initial Voronoi cells # (:math:`g = 0`) next to the Laguerre cells at the optimum. -# In this problem, the maximum cost between samples is 1.0, so we pass it as -# ``max_cost=1.0``. Knowing this bound, the potential values are clipped to -# [-max_cost, max_cost], where it is known that an optimal potential lies ([90]_, Lemma 1), -# which speeds up convergence. +# With the squared-Euclidean cost (default ``metric='sqeuclidean'``), the cost +# between a source point in :math:`[0, 1]^2` and an atom is +# :math:`\|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, - n_iter=20_000, - batch_size=16, - max_cost=1.0, + 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], "DRAG") +plot_laguerre_cells(target_positions, g_drag, axes[1], "Approximated OT Laguerre cells") +plt.tight_layout() +plt.show() + + +############################################################################## +# Transport map over the Laguerre cells +# ------------------------------------- +# +# :func:`semidiscrete_ot_map` with ``reg=0`` is the hard Monge map: every +# source point is sent to the atom of its Laguerre cell. Overlaying the map +# (arrows on a source grid) on the *faded* cells shows each cell's mass +# collapsing onto its atom -- a direct illustration of the mapping function. + +gx = np.linspace(0.04, 0.96, 14) +grid = np.stack([a.ravel() for a in np.meshgrid(gx, gx)], axis=1) +mapped = semidiscrete_ot_map(target_positions, grid, g_drag, reg=0.0) +labels = semidiscrete_atom_weights(target_positions, grid, g_drag, reg=0.0).argmax( + axis=1 +) + +cmap = plt.get_cmap("tab20", n_atoms) +fig, ax = plt.subplots(figsize=(6.5, 6.5)) +plot_laguerre_cells( + target_positions, g_drag, ax, "Approximated OT map over Laguerre cells", alpha=0.22 +) +ax.quiver( + grid[:, 0], + grid[:, 1], + mapped[:, 0] - grid[:, 0], + mapped[:, 1] - grid[:, 1], + angles="xy", + scale_units="xy", + scale=1, + width=0.005, + headwidth=4, + headlength=5, + color=[cmap(i) for i in labels], + zorder=2, +) plt.tight_layout() plt.show() @@ -134,7 +175,9 @@ def plot_laguerre_cells(target, g, ax, title, resolution=300): def cell_masses(target, g, sampler, n_samples=100_000): - labels = atom_weights(target, sampler(n_samples), g, reg=0.0).argmax(axis=1) + labels = semidiscrete_atom_weights(target, sampler(n_samples), g, reg=0.0).argmax( + axis=1 + ) counts = np.bincount(labels, minlength=target.shape[0]) return counts / n_samples @@ -142,7 +185,7 @@ def cell_masses(target, g, sampler, n_samples=100_000): def mc_cost(target, g, sampler, n_samples=100_000): b = np.full(target.shape[0], 1.0 / target.shape[0]) samples = sampler(n_samples) - return float(g @ b + c_transform(target, samples, g, reg=0.0).mean()) + return float(g @ b + semidiscrete_c_transform(target, samples, g, reg=0.0).mean()) target_mass = 1.0 / n_atoms @@ -156,3 +199,37 @@ def mc_cost(target, g, sampler, n_samples=100_000): f" max: {np.max(np.abs(m_drag - target_mass)):.4f}" f" semi-dual cost (MC): {cost_drag:.5f}" ) + + +############################################################################## +# Laguerre-cell masses +# -------------------- +# +# At the optimum every cell carries the same mass :math:`1/15`. The bar plot +# shows the empirical mass per cell against this ground truth (dashed line): +# every cell sits close to the theoretical value. + +cmap = plt.get_cmap("tab20", n_atoms) +fig, ax = plt.subplots(figsize=(7.5, 4)) +ax.bar( + np.arange(n_atoms), + m_drag, + color=[cmap(i) for i in range(n_atoms)], + edgecolor="black", + linewidth=0.6, +) +ax.axhline( + target_mass, + ls="--", + color="black", + lw=1.5, + label="theoretical mass per cell at the optimum", +) +ax.set_ylim(0, 1.6 * target_mass) +ax.set_xticks(np.arange(n_atoms)) +ax.set_xlabel("atom index") +ax.set_ylabel("Laguerre-cell mass") +ax.set_title("Approximated OT: Laguerre-cell masses") +ax.legend() +plt.tight_layout() +plt.show() diff --git a/ot/semidiscrete.py b/ot/semidiscrete.py index 64ed5d1d1..45f37676f 100644 --- a/ot/semidiscrete.py +++ b/ot/semidiscrete.py @@ -9,11 +9,11 @@ References ---------- -.. [1] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). - "Stochastic Optimization in Semi-Discrete Optimal Transport: +.. [1] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Stochastic Optimization in Semi-Discrete Optimal Transport: Convergence Analysis and Minimax Rate." NeurIPS 2025. -.. [2] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). - "Decreasing Entropic Regularization Averaged Gradient for +.. [2] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport." NeurIPS 2025. """ @@ -26,28 +26,60 @@ import numpy as np from .backend import get_backend +from .utils import dist -def _quadratic_cost(x, y, nx): - r"""Default cost: :math:`\tfrac{1}{2} \|x - y\|^2`.""" - x_sq = nx.sum(x**2, axis=1)[:, None] - y_sq = nx.sum(y**2, axis=1)[None, :] - cross = nx.einsum("ij,kj->ik", x, y) - return 0.5 * (x_sq + y_sq - 2.0 * cross) +def _resolve_metric(metric): + r"""Turn ``metric`` into a callable ``(x, y) -> (n_samples, n_atoms)`` matrix. + + ``None`` defaults to ``'sqeuclidean'``. A string is forwarded to + :func:`ot.dist`; a callable is returned unchanged. + """ + if metric is None: + metric = "sqeuclidean" + if callable(metric): + return metric + return lambda x, y: dist(x, y, metric=metric) -def _setup(target_positions, target_weights, cost): - """Resolve backend, default weights and default cost.""" - nx = get_backend(target_positions) - m = target_positions.shape[0] - if target_weights is None: - target_weights = nx.full((m,), 1.0 / m, type_as=target_positions) - if cost is None: +def _setup(X_target, a_target, metric): + """Resolve backend, default weights and metric callable.""" + nx = get_backend(X_target) + m = X_target.shape[0] + if a_target is None: + a_target = nx.full((m,), 1.0 / m, type_as=X_target) + return nx, m, a_target, nx.log(a_target), _resolve_metric(metric) - def cost(x, y): - return _quadratic_cost(x, y, nx) - return nx, m, target_weights, nx.log(target_weights), cost +def _resolve_sampler(sampler_source, X_target, nx): + r"""Turn ``sampler_source`` into a callable ``batch_size -> (batch_size, d)``. + + A callable is returned unchanged. A string selects a built-in sampler + drawing in the same backend as ``X_target``: + + - ``'unif'`` / ``'unif_cube'``: uniform on the unit cube :math:`[0, 1]^d`; + - ``'ball'`` / ``'unif_ball'``: uniform on the unit ball; + - ``'normal'``: standard Gaussian :math:`\mathcal{N}(0, I_d)`. + """ + if callable(sampler_source): + return sampler_source + d = X_target.shape[1] + if sampler_source in ("unif", "unif_cube"): + return lambda b: nx.rand(b, d, type_as=X_target) + if sampler_source in ("ball", "unif_ball"): + + def sampler(b): + z = nx.randn(b, d, type_as=X_target) + r = nx.rand(b, 1, type_as=X_target) ** (1.0 / d) + return r * z / nx.sqrt(nx.sum(z**2, axis=1))[:, None] + + return sampler + if sampler_source == "normal": + return lambda b: nx.randn(b, d, type_as=X_target) + raise ValueError( + f"Unknown sampler_source {sampler_source!r}. Expected a callable or one " + "of 'unif', 'unif_cube', 'ball', 'unif_ball', 'normal'." + ) def _atom_weights(score, reg, log_b, nx): @@ -69,73 +101,227 @@ def _atom_weights(score, reg, log_b, nx): return nx.where(mask, one, zero) -def atom_weights( - target_positions, - source_samples, +def semidiscrete_atom_weights( + X_target, + X_source, semi_dual_potential, - target_weights=None, - cost=None, + a_target=None, + metric=None, reg=0.0, ): - r"""Row-stochastic atom-assignment weights induced by ``semi_dual_potential``. + r"""(Entropic) assignment weights of source samples to target atoms. + + For target atoms :math:`(y_j)_{j=1}^M` (``X_target``) with weights + :math:`(w_j)_j` (``a_target``), a ground cost :math:`c` (``metric``) and a + semi-dual potential :math:`g` (``semi_dual_potential``), this returns, for + each source sample :math:`x` (a row of ``X_source``), the (entropic) + assignment weights :math:`\chi^\varepsilon(x, g)` ([2]_, Sec. 2.2): + + .. math:: + \chi^{\mathrm{reg}}_j(x, g) = + \begin{cases} + \displaystyle + \frac{w_j \exp\!\big((g_j - c(x, y_j))/\mathrm{reg}\big)} + {\sum_{k=1}^M w_k \exp\!\big((g_k - c(x, y_k))/\mathrm{reg}\big)}, + & \mathrm{reg} > 0, \\[1.2em] + \mathbf{1}\big[\, j = \arg\min_k\, c(x, y_k) - g_k \,\big] + & \mathrm{reg} = 0. + \end{cases} + + :math:`\mathbb{E}_{X}[\chi^\varepsilon(X, g)]` is the atom marginal that + the semi-dual gradient matches to :math:`w` ([1]_); cf. + :func:`solve_semidiscrete`. + + Parameters + ---------- + X_target : array-like, shape (n_atoms, d) + Target atom positions :math:`y_j`. + X_source : array-like, shape (n_samples, d) + Source samples :math:`x` to assign. + semi_dual_potential : array-like, shape (n_atoms,) + Semi-dual potential :math:`g`, e.g. from :func:`solve_semidiscrete`. + a_target : array-like, shape (n_atoms,), optional + Atom weights :math:`w_j`. Defaults to uniform. + metric : str or callable, optional + Ground cost. A string is passed to :func:`ot.dist` (e.g. + ``'sqeuclidean'``, ``'euclidean'``); a callable ``metric(x, y)`` + must return the ``(n_samples, n_atoms)`` cost matrix. Defaults to + ``'sqeuclidean'`` (:math:`\|x - y\|^2`). + reg : float, default=0.0 + Entropic regularization :math:`\varepsilon`. ``0`` gives hard, + one-hot assignments; ``> 0`` the softmax above. + + Returns + ------- + w : array, shape (n_samples, n_atoms) + Row-stochastic assignment weights: ``w[i, j]`` is the (entropic) + probability that sample :math:`x_i` is sent to atom :math:`y_j`. + Each row sums to 1. - Returns an array ``w`` of shape ``(n_samples, n_atoms)`` such that - ``w[i, j]`` is the (entropic) probability that sample ``x_i`` is - transported to atom ``y_j``. + References + ---------- + .. [1] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Stochastic Optimization in Semi-Discrete Optimal Transport: + Convergence Analysis and Minimax Rate." NeurIPS 2025. + .. [2] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. """ - nx, _, _, log_b, cost_fn = _setup(target_positions, target_weights, cost) - score = semi_dual_potential[None, :] - cost_fn(source_samples, target_positions) + nx, _, _, log_b, metric_fn = _setup(X_target, a_target, metric) + score = semi_dual_potential[None, :] - metric_fn(X_source, X_target) return _atom_weights(score, reg, log_b, nx) -def ot_map( - target_positions, - source_samples, +def semidiscrete_ot_map( + X_target, + X_source, semi_dual_potential, - target_weights=None, - cost=None, + a_target=None, + metric=None, reg=0.0, ): - r"""Transport map :math:`T(x) = \sum_j w_j(x)\, y_j` induced by the potential.""" - w = atom_weights( - target_positions, - source_samples, + r"""Semi-discrete OT map (barycentric projection) induced by a potential. + + For each source sample :math:`x` (a row of ``X_source``), the transported + position uses the (entropic) assignment weights + :math:`\chi^\varepsilon(x, g)` of :func:`semidiscrete_atom_weights`: + + .. math:: + T(x) = \begin{cases} + \displaystyle \sum_{j=1}^M \chi^\varepsilon_j(x, g)\, y_j + & \varepsilon = \mathrm{reg} > 0, \\[1em] + y_j \ \text{ for } x \in \mathbb{L}_j(g) + & \varepsilon = 0. + \end{cases} + + For :math:`\varepsilon > 0` this is the smoothed barycentric projection; + for :math:`\varepsilon = 0` the weights are one-hot and :math:`T` is the + Monge map of the (generalized) Brenier theorem ([1]_), sending :math:`x` + to the atom of its Laguerre cell + + .. math:: + \mathbb{L}_j(g) = \big\{\, x : g^{c}(x) = c(x, y_j) - g_j \,\big\}, + + cf. [2]_, Sec. 2.2. + + Parameters + ---------- + X_target : array-like, shape (n_atoms, d) + Target atom positions :math:`y_j`. + X_source : array-like, shape (n_samples, d) + Source samples :math:`x` to transport. + semi_dual_potential : array-like, shape (n_atoms,) + Semi-dual potential :math:`g`, e.g. from :func:`solve_semidiscrete`. + a_target : array-like, shape (n_atoms,), optional + Atom weights :math:`w_j`. Defaults to uniform. + metric : str or callable, optional + Ground cost. A string is passed to :func:`ot.dist` (e.g. + ``'sqeuclidean'``, ``'euclidean'``); a callable ``metric(x, y)`` + must return the ``(n_samples, n_atoms)`` cost matrix. Defaults to + ``'sqeuclidean'`` (:math:`\|x - y\|^2`). + reg : float, default=0.0 + Entropic regularization :math:`\varepsilon`. ``0`` gives the hard + Monge map; ``> 0`` the smoothed barycentric map. + + Returns + ------- + T : array, shape (n_samples, d) + Transported source positions :math:`T(x_i)`. + + References + ---------- + .. [1] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Stochastic Optimization in Semi-Discrete Optimal Transport: + Convergence Analysis and Minimax Rate." NeurIPS 2025. + .. [2] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. + """ + w = semidiscrete_atom_weights( + X_target, + X_source, semi_dual_potential, - target_weights=target_weights, - cost=cost, + a_target=a_target, + metric=metric, reg=reg, ) - return w @ target_positions + return w @ X_target -def c_transform( - target_positions, - source_samples, +def semidiscrete_c_transform( + X_target, + X_source, semi_dual_potential, - target_weights=None, - cost=None, + a_target=None, + metric=None, reg=0.0, ): - r"""Pointwise (entropic) c-transform of ``semi_dual_potential``. + r"""(Entropic) :math:`c`-transform of a semi-dual potential. - - ``reg == 0``: :math:`\varphi_g(x) = \min_j\, c(x, y_j) - g_j`. - - ``reg > 0``: :math:`\varphi_g(x) = -\varepsilon \log \sum_j b_j - \exp\!\big((g_j - c(x, y_j))/\varepsilon\big)`. + The vectorial :math:`(c, \varepsilon)`-transform + :math:`g^{c,\varepsilon}` of the potential :math:`g` + (``semi_dual_potential``), evaluated at the source samples ([1]_, Eq. (3); + entropic form in [2]_, Eq. (4)): + + .. math:: + g^{c,\varepsilon}(x) = \begin{cases} + \min_{j}\, \big(c(x, y_j) - g_j\big) & \varepsilon = 0, \\[4pt] + -\varepsilon \log \sum_{j=1}^M w_j + \exp\!\big((g_j - c(x, y_j))/\varepsilon\big) & \varepsilon > 0. + \end{cases} + + For :math:`\varepsilon = 0` this is the standard :math:`c`-transform + :math:`g^{c}(x) = \min_j (c(x, y_j) - g_j)` whose expectation gives the + concave semi-dual :math:`H(g) = \mathbb{E}_X[g^{c}(X)] + \langle g, w\rangle` + maximized by :func:`solve_semidiscrete` ([1]_, Eq. (2)). + + Parameters + ---------- + X_target : array-like, shape (n_atoms, d) + Target atom positions :math:`y_j`. + X_source : array-like, shape (n_samples, d) + Source samples :math:`x` at which to evaluate the transform. + semi_dual_potential : array-like, shape (n_atoms,) + Semi-dual potential :math:`g`, e.g. from :func:`solve_semidiscrete`. + a_target : array-like, shape (n_atoms,), optional + Atom weights :math:`w_j`. Defaults to uniform. + metric : str or callable, optional + Ground cost. A string is passed to :func:`ot.dist` (e.g. + ``'sqeuclidean'``, ``'euclidean'``); a callable ``metric(x, y)`` + must return the ``(n_samples, n_atoms)`` cost matrix. Defaults to + ``'sqeuclidean'`` (:math:`\|x - y\|^2`). + reg : float, default=0.0 + Entropic regularization :math:`\varepsilon`. ``0`` gives the hard + :math:`\min`; ``> 0`` the soft log-sum-exp. + + Returns + ------- + phi : array, shape (n_samples,) + The :math:`c`-transform :math:`g^{c,\varepsilon}(x_i)` at each sample. + + References + ---------- + .. [1] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Stochastic Optimization in Semi-Discrete Optimal Transport: + Convergence Analysis and Minimax Rate." NeurIPS 2025. + .. [2] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. """ - nx, _, _, log_b, cost_fn = _setup(target_positions, target_weights, cost) - score = semi_dual_potential[None, :] - cost_fn(source_samples, target_positions) + nx, _, _, log_b, metric_fn = _setup(X_target, a_target, metric) + score = semi_dual_potential[None, :] - metric_fn(X_source, X_target) if reg == 0: return -nx.max(score, axis=1) return -reg * nx.logsumexp(score / reg + log_b[None, :], axis=1) def solve_semidiscrete( - target_positions, - source_sampler, - target_weights=None, - cost=None, + X_target, + sampler_source="unif", + a_target=None, + metric=None, reg=0.0, - n_iter=10_000, + max_iter=10_000, batch_size=32, lr0=None, lr_exponent=2.0 / 3.0, @@ -147,43 +333,47 @@ def solve_semidiscrete( polyak_average=True, log=False, ): - r"""Solve semi-discrete OT by Polyak-averaged SGD on the semi-dual. + r"""Solve semi-discrete OT by (projected) averaged SGD on the semi-dual. - Maximizes the semi-dual :math:`g \mapsto \langle g, b \rangle + \mathbb{E}_X[\varphi_g(X)]` - by averaged stochastic gradient ascent with projection and decreasing - regularization, which corresponds to the DRAG algorithm [1]_. - Here :math:`\varphi_g` denotes the (entropic) c-transform of :math:`g`, + Maximizes the concave semi-dual objective ([1]_, Eq. (4), in its negative convex form) .. math:: - \varphi_g(x) = \begin{cases} - \min_j \big(c(x, y_j) - g_j\big) & \text{if } \mathrm{reg} = 0, \\ - -\varepsilon \log \sum_j b_j \exp\!\big((g_j - c(x, y_j))/\varepsilon\big) - & \text{if } \mathrm{reg} = \varepsilon > 0, - \end{cases} + H(g) = \mathbb{E}_{X}[g^{c}(X)] + \langle g, w\rangle , + + over the potential :math:`g \in \mathbb{R}^M`, where :math:`g^{c}` is the + (entropic) :math:`c`-transform of :math:`g` (see + :func:`semidiscrete_c_transform`). - cf. :func:`c_transform`. + The base solver is the projected averaged SGD of [1]_ (Algorithm 1); the + projection ``max_cost`` clips each iterate to the localizing set + :math:`\{|g_j| \le \texttt{max\_cost}\}` ([1]_, Sec. 3.1). When + ``decreasing_reg=True``, the entropic regularization is annealed along the + iterations following the DRAG schedule of [2]_ (Algorithm 1), which + accelerates convergence. With ``decreasing_reg=True`` the regularization at iteration ``t`` is :math:`\varepsilon_t = \max(\text{reg},\, \varepsilon_0 / t^\alpha)` — large at first for smoothness, then annealed towards ``reg``. This is the - DRAG schedule of [1]_. + DRAG schedule of [2]_. Parameters ---------- - target_positions : array-like, shape (n_atoms, d) + X_target : array-like, shape (n_atoms, d) Positions of the target atoms. The backend of this array drives all subsequent computations. - source_sampler : callable - ``source_sampler(batch_size)`` returns a ``(batch_size, d)`` array - of source samples, in the same backend as ``target_positions``. - target_weights : array-like, shape (n_atoms,), optional + sampler_source : str or callable, default='unif' + Source distribution to sample from: either a callable + ``sampler_source(batch_size)`` returning a ``(batch_size, d)`` batch in + the backend of ``X_target``, or a built-in name -- one of ``'unif'`` + (``'unif_cube'``), ``'ball'`` (``'unif_ball'``) or ``'normal'``. + a_target : array-like, shape (n_atoms,), optional Atom weights. Defaults to uniform. - cost : callable, optional - ``cost(x, y)`` returns the ``(n_samples, n_atoms)`` cost matrix. - Defaults to ``0.5 * ||x - y||^2``. + metric : str or callable, optional + Ground cost, see ``metric`` in :func:`semidiscrete_atom_weights`. + Defaults to ``'sqeuclidean'``. reg : float, default=0.0 Entropic regularization (target value when ``decreasing_reg=True``). - n_iter : int, default=10000 + max_iter : int, default=10000 batch_size : int, default=32 lr0 : float, optional Initial learning rate. Defaults to ``sqrt(n_atoms * batch_size)``. @@ -212,9 +402,12 @@ def solve_semidiscrete( References ---------- - .. [1] Genans, Godichon-Baggioni, Vialard, Wintenberger (2025). - "Decreasing Entropic Regularization Averaged Gradient for - Semi-Discrete Optimal Transport." NeurIPS 2025. + .. [1] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Stochastic Optimization in Semi-Discrete Optimal Transport: + Convergence Analysis and Minimax Rate." NeurIPS 2025. + .. [2] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O. + (2025). "Decreasing Entropic Regularization Averaged Gradient for + Semi-Discrete Optimal Transport." NeurIPS 2025. Examples -------- @@ -224,29 +417,30 @@ def solve_semidiscrete( >>> target = np.linspace(0.0, 1.0, 10).reshape(-1, 1) >>> g = solve_semidiscrete( ... target, lambda b: rng.random((b, 1)), - ... n_iter=500, batch_size=32, max_cost=1.0, + ... max_iter=500, batch_size=32, max_cost=1.0, ... ) """ - nx, m, b, log_b, cost_fn = _setup(target_positions, target_weights, cost) + nx, m, b, log_b, metric_fn = _setup(X_target, a_target, metric) + sampler_source = _resolve_sampler(sampler_source, X_target, nx) if init_potential is None: - g = nx.zeros((m,), type_as=target_positions) + g = nx.zeros((m,), type_as=X_target) else: - g = init_potential + nx.zeros((m,), type_as=target_positions) + g = init_potential + nx.zeros((m,), type_as=X_target) if lr0 is None: lr0 = math.sqrt(m * batch_size) - g_avg = nx.zeros((m,), type_as=target_positions) if polyak_average else None + g_avg = nx.zeros((m,), type_as=X_target) if polyak_average else None - for t in range(1, n_iter + 1): + for t in range(1, max_iter + 1): if decreasing_reg: reg_t = max(reg, decreasing_reg_initial_eps / (t**decreasing_reg_exponent)) else: reg_t = reg - x = source_sampler(batch_size) - score = g[None, :] - cost_fn(x, target_positions) + x = sampler_source(batch_size) + score = g[None, :] - metric_fn(x, X_target) w = _atom_weights(score, reg_t, log_b, nx) grad = nx.mean(w, axis=0) - b @@ -260,7 +454,7 @@ def solve_semidiscrete( result = g_avg if polyak_average else g if log: return result, { - "n_iter": n_iter, + "max_iter": max_iter, "batch_size": batch_size, "max_cost": max_cost, "polyak_average": polyak_average, @@ -271,7 +465,7 @@ def solve_semidiscrete( __all__ = [ "solve_semidiscrete", - "atom_weights", - "ot_map", - "c_transform", + "semidiscrete_atom_weights", + "semidiscrete_ot_map", + "semidiscrete_c_transform", ] diff --git a/test/test_semidiscrete.py b/test/test_semidiscrete.py index 3ff048a1b..feb806b0d 100644 --- a/test/test_semidiscrete.py +++ b/test/test_semidiscrete.py @@ -11,10 +11,11 @@ import numpy as np import pytest +from ot.backend import get_backend from ot.semidiscrete import ( - atom_weights, - c_transform, - ot_map, + semidiscrete_atom_weights, + semidiscrete_c_transform, + semidiscrete_ot_map, solve_semidiscrete, ) @@ -25,7 +26,7 @@ # --------------------------------------------------------------------- # Three toy problems with known optimal potentials. -# Each builder returns numpy arrays so we can lift them onto any backend. +# Each builder returns numpy arrays so we can move them onto any backend. # --------------------------------------------------------------------- @@ -126,11 +127,15 @@ def centered_l2_error(estimated, reference): return float(np.linalg.norm(estimated - reference)) -def lift(nx, target_np, weights_np): - """Move numpy ``target`` and ``weights`` arrays onto backend ``nx``.""" - target = nx.from_numpy(target_np) - weights = nx.from_numpy(weights_np, type_as=target) - return target, weights +def half_sqeuclidean(x, y): + r"""Backend-aware ``0.5 * ||x - y||^2`` cost, used to pin closed forms. + + The closed-form optima below are derived for this cost; passing it + explicitly keeps the tests valid independently of the default metric. + """ + nx = get_backend(x, y) + diff = x[:, None, :] - y[None, :, :] + return 0.5 * nx.sum(diff**2, axis=2) # --------------------------------------------------------------------- @@ -142,14 +147,15 @@ def lift(nx, target_np, weights_np): def test_solve_converges(nx, build_problem): """Plain SGD (no decreasing reg) reaches the optimum on every toy problem.""" target_np, weights_np, optimal, max_cost, d, kind = build_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) g = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=N_ITER, + a_target=weights, + metric=half_sqeuclidean, + max_iter=N_ITER, batch_size=BATCH_SIZE, decreasing_reg=False, max_cost=max_cost, @@ -162,14 +168,15 @@ def test_solve_converges(nx, build_problem): def test_drag_converges(nx, build_problem): """DRAG (decreasing entropic reg) reaches the optimum on every toy problem.""" target_np, weights_np, optimal, max_cost, d, kind = build_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) g = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=N_ITER, + a_target=weights, + metric=half_sqeuclidean, + max_iter=N_ITER, batch_size=BATCH_SIZE, decreasing_reg=True, max_cost=max_cost, @@ -184,22 +191,22 @@ def test_drag_converges(nx, build_problem): def test_entropic_solver_runs(nx): - """In the entropic regime, the solver and ``c_transform`` produce finite values.""" + """In the entropic regime, the solver and ``semidiscrete_c_transform`` produce finite values.""" target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) g = solve_semidiscrete( target, sampler, - target_weights=weights, + a_target=weights, reg=0.05, - n_iter=N_ITER, + max_iter=N_ITER, batch_size=BATCH_SIZE, max_cost=max_cost, ) samples = sampler(64) - phi = c_transform(target, samples, g, target_weights=weights, reg=0.05) + phi = semidiscrete_c_transform(target, samples, g, a_target=weights, reg=0.05) assert np.isfinite(nx.to_numpy(g)).all() assert np.isfinite(nx.to_numpy(phi)).all() @@ -209,22 +216,64 @@ def test_entropic_solver_runs(nx): # --------------------------------------------------------------------- -def test_custom_quadratic_cost_matches_default(nx): - """A user-supplied quadratic cost reaches the same optimum as the default.""" - target_np, weights_np, optimal, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) +def test_string_metric_runs(nx): + """A string metric is routed through ``ot.dist`` and yields finite output.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) - def cost(x, y): - diff = x[:, None, :] - y[None, :, :] - return 0.5 * nx.sum(diff**2, axis=2) - g = solve_semidiscrete( target, sampler, - target_weights=weights, - cost=cost, - n_iter=N_ITER, + a_target=weights, + metric="euclidean", + max_iter=N_ITER, + batch_size=BATCH_SIZE, + max_cost=max_cost, + ) + samples = sampler(32) + phi = semidiscrete_c_transform( + target, samples, g, a_target=weights, metric="euclidean" + ) + assert np.isfinite(nx.to_numpy(g)).all() + assert np.isfinite(nx.to_numpy(phi)).all() + + +# --------------------------------------------------------------------- +# Built-in samplers +# --------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "name", ["unif", "unif_cube", "ball", "unif_ball", "normal"] +) +def test_string_sampler_runs(nx, name): + """Each built-in string sampler yields a finite potential of the right shape.""" + target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() + target, weights = nx.from_numpy(target_np, weights_np) + + g = solve_semidiscrete( + target, + name, + a_target=weights, + max_iter=200, + batch_size=BATCH_SIZE, + max_cost=max_cost, + ) + assert nx.to_numpy(g).shape == (target_np.shape[0],) + assert np.isfinite(nx.to_numpy(g)).all() + + +def test_default_sampler_unif_converges(nx): + """The default ``sampler_source='unif'`` matches the uniform-cube problem.""" + target_np, weights_np, optimal, max_cost, d, kind = regular_grid_problem() + target, weights = nx.from_numpy(target_np, weights_np) + + g = solve_semidiscrete( + target, + a_target=weights, + metric=half_sqeuclidean, + max_iter=N_ITER, batch_size=BATCH_SIZE, max_cost=max_cost, ) @@ -232,6 +281,14 @@ def cost(x, y): assert err < TOLERANCE, f"err={err:.4f}" +def test_unknown_sampler_raises(nx): + """An unknown sampler string is rejected with ``ValueError``.""" + target_np, weights_np, _, _, _, _ = regular_grid_problem() + target, _ = nx.from_numpy(target_np, weights_np) + with pytest.raises(ValueError): + solve_semidiscrete(target, "nope", max_iter=10) + + # --------------------------------------------------------------------- # Helper functions # --------------------------------------------------------------------- @@ -239,13 +296,13 @@ def cost(x, y): @pytest.mark.parametrize("reg", [0.0, 0.1]) def test_atom_weights_are_row_stochastic(nx, reg): - """``atom_weights`` returns nonnegative weights that sum to 1 per row.""" + """``semidiscrete_atom_weights`` returns nonnegative weights that sum to 1 per row.""" target_np, weights_np, _, _, d, kind = nonuniform_weights_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) samples = sampler(32) g = nx.zeros((target_np.shape[0],), type_as=target) - w = atom_weights(target, samples, g, target_weights=weights, reg=reg) + w = semidiscrete_atom_weights(target, samples, g, a_target=weights, reg=reg) w_np = nx.to_numpy(w) assert w_np.shape == (32, target_np.shape[0]) assert (w_np >= 0).all() @@ -253,13 +310,13 @@ def test_atom_weights_are_row_stochastic(nx, reg): def test_ot_map_shape_and_finiteness(nx): - """``ot_map`` returns finite values with the source-sample shape.""" + """``semidiscrete_ot_map`` returns finite values with the source-sample shape.""" target_np, weights_np, _, _, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) samples = sampler(16) g = nx.zeros((target_np.shape[0],), type_as=target) - transported = ot_map(target, samples, g, target_weights=weights) + transported = semidiscrete_ot_map(target, samples, g, a_target=weights) transported_np = nx.to_numpy(transported) samples_np = nx.to_numpy(samples) assert transported_np.shape == samples_np.shape @@ -269,11 +326,13 @@ def test_ot_map_shape_and_finiteness(nx): def test_c_transform_minimum_for_zero_potential(nx): """At ``g = 0``, ``phi_g(x) = -max_j(-c(x, y_j)) = min_j c(x, y_j)``.""" target_np, weights_np, _, _, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) samples = sampler(8) g = nx.zeros((target_np.shape[0],), type_as=target) - phi = c_transform(target, samples, g, target_weights=weights) + phi = semidiscrete_c_transform( + target, samples, g, a_target=weights, metric=half_sqeuclidean + ) samples_np = nx.to_numpy(samples) target_np = nx.to_numpy(target) diff = samples_np[:, None, :] - target_np[None, :, :] @@ -289,22 +348,24 @@ def test_c_transform_minimum_for_zero_potential(nx): def test_warm_start_converges(nx): """Splitting one run into two warm-started halves still converges.""" target_np, weights_np, optimal, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) half = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=N_ITER // 2, + a_target=weights, + metric=half_sqeuclidean, + max_iter=N_ITER // 2, batch_size=BATCH_SIZE, max_cost=max_cost, ) g = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=N_ITER // 2, + a_target=weights, + metric=half_sqeuclidean, + max_iter=N_ITER // 2, batch_size=BATCH_SIZE, init_potential=half, max_cost=max_cost, @@ -316,7 +377,7 @@ def test_warm_start_converges(nx): def test_init_potential_is_not_mutated(nx): """The ``init_potential`` array passed by the caller is left intact.""" target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) init_np = np.full(target_np.shape[0], 0.5) @@ -326,9 +387,9 @@ def test_init_potential_is_not_mutated(nx): solve_semidiscrete( target, sampler, - target_weights=weights, + a_target=weights, init_potential=init, - n_iter=10, + max_iter=10, batch_size=4, max_cost=max_cost, ) @@ -338,15 +399,15 @@ def test_init_potential_is_not_mutated(nx): def test_projection_clamps_last_iterate(nx): """With ``max_cost=b``, every coordinate of the last iterate lies in ``[-b, b]``.""" target_np, weights_np, _, _, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) bound = 0.05 _, info = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=300, + a_target=weights, + max_iter=300, batch_size=4, max_cost=bound, log=True, @@ -358,14 +419,14 @@ def test_projection_clamps_last_iterate(nx): def test_polyak_average_off_returns_last_iterate(nx): """With ``polyak_average=False`` the returned potential equals the last iterate.""" target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) final, info = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=20, + a_target=weights, + max_iter=20, batch_size=4, polyak_average=False, max_cost=max_cost, @@ -379,19 +440,19 @@ def test_polyak_average_off_returns_last_iterate(nx): def test_log_returns_metadata(nx): """``log=True`` returns an info dict with the expected fields.""" target_np, weights_np, _, max_cost, d, kind = regular_grid_problem() - target, weights = lift(nx, target_np, weights_np) + target, weights = nx.from_numpy(target_np, weights_np) sampler = make_sampler(kind, d, nx, target) g, info = solve_semidiscrete( target, sampler, - target_weights=weights, - n_iter=50, + a_target=weights, + max_iter=50, batch_size=4, max_cost=max_cost, log=True, ) assert nx.to_numpy(g).shape == (target_np.shape[0],) - assert info["n_iter"] == 50 + assert info["max_iter"] == 50 assert info["batch_size"] == 4 assert info["max_cost"] == max_cost From 6bfe7b8faad9783f992babb9cdeb91e5dfc14848 Mon Sep 17 00:00:00 2001 From: Ferdinand-Genans Date: Mon, 29 Jun 2026 15:53:14 +0200 Subject: [PATCH 8/8] Added my name in citation.cff --- CITATION.cff | 2 ++ test/test_semidiscrete.py | 4 +--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/CITATION.cff b/CITATION.cff index 2c73bd3bd..5fe2a31ab 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -84,6 +84,8 @@ authors: - given-names: Marco family-names: Corneli affiliation: Université Côte d'Azur + - given-names: Ferdinand Genans + affiliation: Sorbonne Université, LPSM, CNRS identifiers: - type: url value: 'https://github.com/PythonOT/POT' diff --git a/test/test_semidiscrete.py b/test/test_semidiscrete.py index feb806b0d..ce89c5aa7 100644 --- a/test/test_semidiscrete.py +++ b/test/test_semidiscrete.py @@ -244,9 +244,7 @@ def test_string_metric_runs(nx): # --------------------------------------------------------------------- -@pytest.mark.parametrize( - "name", ["unif", "unif_cube", "ball", "unif_ball", "normal"] -) +@pytest.mark.parametrize("name", ["unif", "unif_cube", "ball", "unif_ball", "normal"]) def test_string_sampler_runs(nx, name): """Each built-in string sampler yields a finite potential of the right shape.""" target_np, weights_np, _, max_cost, d, kind = regular_grid_problem()