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

Author: Ferdinand-Genans

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,20 +61,20 @@ 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(
         image,
         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,15 +175,17 @@ 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
 
 
 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()