[MRG] Add backend-agnostic semi-discrete OT module and SGD-based solver in ot.semidiscrete

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'

README.md

@@ -468,3 +468,5 @@ Artificial Intelligence.
 [88] Bouveyron, C. & Corneli, M. (2026).  [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88.
 
 [89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622.
+
+[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [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.

RELEASES.md

@@ -26,6 +26,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s
 - Add `ot.utils.apply_scaler` helper that dispatches preprocessing to a scaler object,
   a callable, or a no-op (PR #808)
 - Add optional `scaler` parameter to `sliced_wasserstein_distance` and  `max_sliced_wasserstein_distance` (PR #808)
+- Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR #812)
 - Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723)
 - Add cost functions between linear operators following  
   [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920),  

examples/others/plot_semidiscrete.py

@@ -0,0 +1,235 @@
+# -*- 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 [90]_, 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.semidiscrete_c_transform`. The solver **maximises** this
+  objective.
+
+.. [90] 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 <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 :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.
+# 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,
+    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()
+
+
+##############################################################################
+# 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()
+
+
+##############################################################################
+# 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 = 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
+
+
+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 + semidiscrete_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}"
+)
+
+
+##############################################################################
+# 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()

ot/__init__.py

@@ -37,6 +37,7 @@
 from . import gaussian
 from . import lowrank
 from . import gmm
+from . import semidiscrete
 from . import sgot
 
 # OT functions
@@ -156,6 +157,7 @@
     "factored",
     "lowrank",
     "gmm",
+    "semidiscrete",
     "sgot",
     "binary_search_circle",
     "wasserstein_circle",