[MRG] add debiasing and all new solvers to ot.solve_sample

README.md

@@ -455,18 +455,20 @@ Artificial Intelligence.
 
 \[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] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920).
+\[83] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920).
 
-[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15.
+\[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15.
 
-[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385.
+\[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385.
 
-[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661.
+\[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661.
 
-[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations.
+\[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations.
 
-[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.
+\[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.
+\[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.
+\[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.
+
+\[91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & Courty, N. (2021). [Minibatch optimal transport distances; analysis and applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792.

RELEASES.md

@@ -38,7 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s
 - Update the geomloss wrapper to the new version and API (PR #826)
 - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823)
 - Reorganize all tests per backend (PR #828)
-- Update sgot cost function and example (PR #830)
+- Implement debiased OT solvers in `ot.solve_sample`.
 
 
 #### Closed issues
@@ -61,7 +61,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s
 - Fix documentation build on master with submodules (PR #818)
 - Fix failing test for unbalanced solver with generic regularization (PR #824)
 - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823)   
-
+- Update sgot cost function and example (PR #830)
 
 ## 0.9.6.post1
 

examples/plot_debias_sink_div.py

@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+"""
+======================================
+Sinkhorn Divergence and Debiased OT solvers
+======================================
+
+This example shows how to use the debiased OT solvers in `ot.solve_sample` to
+compute Sinkhorn divergences and debiased Minibatch solutions. The debiased OT solvers
+can be used with balanced and unbalanced OT problems, and with different
+regularization types (entropic, L2, group lasso).
+"""
+
+# Author: Remi Flamary <remi.flamary@polytechnique.edu>
+#
+# License: MIT License
+# sphinx_gallery_thumbnail_number = 3
+
+# %%
+
+import numpy as np
+import matplotlib.pylab as pl
+import ot
+import ot.plot
+from ot.datasets import make_1D_gauss as gauss
+
+##############################################################################
+# Generate data
+# -------------
+
+
+# %%
+def sample_ball(n, radius=1.0, center=(0.0, 0.0)):
+    np.random.seed(0)
+    theta = 2 * np.pi * np.random.rand(n)
+    r = radius * np.sqrt(np.random.rand(n))
+
+    x = r * np.cos(theta) + center[0]
+    y = r * np.sin(theta) + center[1]
+
+    return np.stack((x, y), axis=1)
+
+
+def sample_two_balls(n, radius=1.0, sep=1):
+    assert n % 2 == 0, "n must be even"
+    centers = ((-sep, -sep), (sep, sep))
+    n_half = n // 2
+    X1 = sample_ball(n_half, radius, centers[0])
+    X2 = sample_ball(n_half, radius, centers[1])
+
+    perm = np.random.permutation(n_half * 2)
+    X = np.vstack((X1, X2))
+    X = X[perm]
+    return X
+
+
+n = 50
+
+x1 = sample_ball(n, radius=1.0, center=(0, 0))
+x2 = sample_two_balls(n, radius=1.0, sep=1.5)
+
+pl.figure(1, figsize=(5, 5))
+pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", alpha=0.7)
+pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", alpha=0.7)
+pl.legend()
+pl.title("Two distributions")
+ax = pl.axis()
+
+
+##############################################################################
+# Compute Sinkhorn divergence and visualize plans
+# -----------------------------------------------
+# The Sinkhorn divergence is computed by setting the `debias` parameter to
+# `True` in the `ot.solve_sample` function. The resulting value is the Sinkhorn
+# divergence. The Sinkhorn divergences is computed as:
+#
+# .. math::
+#     S_\epsilon(\mu, \nu) = OT_\epsilon(\mu, \nu) - \frac{1}{2} OT_\epsilon(\mu, \mu) - \frac{1}{2} OT_\epsilon(\nu, \nu)
+#
+# The entropic OT plans for each of those terms can be accessed in the `log`
+# attribute of the result, and can be visualized using the
+# `ot.plot.plot2D_samples_mat` function.
+
+res = ot.solve_sample(x1, x2, reg=0.1, debias=True)
+
+print("Sinkhorn divergence: ", res.value)
+
+plan_11 = res.log["res_aa"].plan
+plan_12 = res.log["res_ab"].plan
+plan_22 = res.log["res_bb"].plan
+
+#
+pl.figure(2, figsize=(15, 5))
+
+pl.subplot(1, 3, 1)
+ot.plot.plot2D_samples_mat(x1, x1, plan_11, thr=0.05)
+pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2)
+pl.axis(ax)
+pl.title("Plan between source and source")
+pl.subplot(1, 3, 2)
+ot.plot.plot2D_samples_mat(x1, x2, plan_12, thr=0.05)
+pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2)
+pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", zorder=2)
+pl.axis(ax)
+pl.title("Plan between source and target")
+pl.subplot(1, 3, 3)
+ot.plot.plot2D_samples_mat(x2, x2, plan_22, thr=0.05)
+pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", color="C1", zorder=2)
+pl.axis(ax)
+pl.title("Plan between target and target")
+
+
+##############################################################################
+# Debiased Minibatch OT
+# ---------------------------------
+#
+# Doing OT on minibatches leads to a similar bias than using entropic
+# regularization since the average OT plan is densified due to the stochasticity
+# of the minibatch sampling.  On a given minibatch, the debiased loss can be
+# computed by setting the `debias` parameter to `'split'`that split the data
+# points in each distributions in two and computes the debias OT loss as:
+#
+# .. math::
+#   \tilde{OT}_m(\mu, \nu) = \frac{1}{2}(\hat{OT}_m(\mu_1, \nu_1) + \hat{OT}_m(\mu_2, \nu_2) - \hat{OT}_m(\nu_1, \nu_2) - \hat{OT}_m(\mu_1, \nu_2))
+#
+
+# %% solve OT minibtach and visualize the plans
+
+res = ot.solve_sample(x1, x2, debias="split")
+
+print("Debiased minibatch OT loss: ", res.value)
+
+# recover the plans for each of the four terms in the debiased loss
+plan_11 = res.log["res_aa"].plan
+plan_12 = res.log["res_ab1"].plan
+plan_21 = res.log["res_ab2"].plan
+plan_22 = res.log["res_bb"].plan
+sel_a1 = res.log["sel_a1"]
+sel_a2 = res.log["sel_a2"]
+sel_b1 = res.log["sel_b1"]
+sel_b2 = res.log["sel_b2"]
+
+nb1 = plan_11.shape[0]
+nb2 = plan_22.shape[0]
+
+pl.figure(4, figsize=(15, 3))
+
+pl.subplot(1, 4, 1)
+pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="$\mu_1$", zorder=2)
+pl.scatter(
+    x1[sel_a2, 0], x1[sel_a2, 1], label=r"$\mu_2$", zorder=2, color="C0", alpha=0.5
+)
+pl.scatter(x2[sel_b1, 0], x2[sel_b1, 1], label=r"$\nu_1$", zorder=2, color="C1")
+pl.scatter(
+    x2[sel_b2, 0], x2[sel_b2, 1], label=r"$\nu_2$", zorder=2, color="C1", alpha=0.5
+)
+pl.title("Minibatch split")
+pl.axis(ax)
+pl.legend()
+
+
+pl.subplot(1, 4, 2)
+ot.plot.plot2D_samples_mat(x1[sel_a1], x1[sel_a2], plan_11, thr=0.05)
+pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], zorder=2)
+pl.scatter(
+    x1[sel_a2, 0],
+    x1[sel_a2, 1],
+    zorder=2,
+    color="C0",
+    alpha=0.5,
+)
+pl.axis(ax)
+pl.title("Plan between source and source")
+pl.subplot(1, 4, 3)
+ot.plot.plot2D_samples_mat(x1[sel_a1], x2[sel_b1], plan_12, thr=0.05)
+ot.plot.plot2D_samples_mat(x1[sel_a2], x2[sel_b2], plan_21, thr=0.05, alpha=0.5)
+
+pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2)
+pl.scatter(
+    x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1"
+)
+pl.scatter(
+    x1[sel_a2, 0],
+    x1[sel_a2, 1],
+    label="Source distribution",
+    zorder=2,
+    color="C0",
+    alpha=0.5,
+)
+pl.scatter(
+    x2[sel_b2, 0],
+    x2[sel_b2, 1],
+    label="Target distribution",
+    zorder=2,
+    color="C1",
+    alpha=0.5,
+)
+pl.axis(ax)
+pl.title("Plan between source and target")
+
+pl.subplot(1, 4, 4)
+ot.plot.plot2D_samples_mat(x2[sel_b1], x2[sel_b2], plan_22, thr=0.05)
+pl.scatter(
+    x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1"
+)
+pl.scatter(
+    x2[sel_b2, 0],
+    x2[sel_b2, 1],
+    label="Target distribution",
+    zorder=2,
+    color="C1",
+    alpha=0.5,
+)
+pl.axis(ax)
+pl.title("Plan between target and target")
+
+
+##############################################################################
+# Comparison of the divergences
+# -------------------------------------------------
+
+# %% move a distribution and compute Sinkhorn divergence and Sinkhorn distance
+reg = 0.1
+
+sep_list = np.linspace(0, 1.0, 10)
+sink_list = []
+sink_div_list = []
+ot_mb_list = []
+ot_mb_sink_list = []
+for sep in sep_list:
+    x2sep = sample_two_balls(n, radius=1.0, sep=sep)
+    sink_list.append(
+        ot.solve_sample(
+            x1,
+            x2sep,
+            reg=reg,
+        ).value
+    )
+    sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value)
+    ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value)
+
+    ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value)
+
+pl.figure(3)
+pl.plot(sep_list, sink_list, label="Sinkhorn loss", color="C0")
+pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1")
+pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2")
+pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3")
+pl.xlabel("Separation between distributions")
+pl.ylabel("Loss/Divergence")
+pl.title("Comparison of biased VS debiased OT losses")
+pl.grid()
+pl.legend()
+
+# %%

examples/plot_quickstart_guide.py

@@ -467,8 +467,8 @@ def df(G):
 #      loss_fgw = ot.gromov.fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1)
 #      loss_efgw = ot.gromov.entropic_fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1, epsilon=reg)
 #
-# Large scale OT
-# --------------
+# Large scale OT and approximations
+# ---------------------------------
 #
 # We discuss here strategies to solve large scale OT problems using approximations
 # of the exact OT problem.
@@ -583,11 +583,67 @@ def df(G):
 print(f"Exact OT loss = {loss:1.3f}")
 print(f"Bures-Wasserstein distance = {bw_value:1.3f}")
 
+# %%
+# One can also use the HD gaussian assumption (low rank covariance + diagonal)
+# that has better properties in high dimension. The rank of the covariance
+# matrices can be controlled with the :code:`rank` parameter.
+
+hdbw_value = ot.solve_sample(x1, x2, a, b, method="gaussian_hd", rank=1).value
+
+print(f"Bures-Wasserstein distance = {bw_value:1.3f}")
+print(f"High Dimensional Bures-Wasserstein distance = {hdbw_value}")
+
 # %%
 # .. note::
 #    The Gaussian Wasserstein problem can be solved with the classic API using the
 #    :func:`ot.gaussian.empirical_bures_wasserstein_distance` function.
 #
+# Sliced Wasserstein
+# ~~~~~~~~~~~~~~~~~~
+#
+# The Sliced Wasserstein distance is a Wasserstein distance between
+# empirical distributions that is computed by projecting the samples on random
+# directions and averaging the Wasserstein distances between the projected
+# distributions. It can be used as an approximation of the Wasserstein distance
+# between empirical distributions.
+
+sw_value = ot.solve_sample(x1, x2, a, b, method="sliced", n_projections=10).value
+max_sw_value = ot.solve_sample(
+    x1, x2, a, b, method="max_sliced", n_projections=10
+).value
+
+print(f"Exact OT loss = {loss:1.3f}")
+print(f"Sliced Wasserstein distance = {sw_value:1.3f}")
+print(f"Max Sliced Wasserstein distance = {max_sw_value:1.3f}")
+
+# %%
+# Binary Space Partitioning (BSP) OT
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# One can also use the BSP OT approximation that is based on a recursive
+# partitioning of the space and computes the Wasserstein distance between the
+# empirical distributions by solving small OT problems between the samples in each
+# partition. The number of partitions can be controlled with the
+# `n_projections` parameter.
+
+# BSP can only find bijections so require same number of points
+x1_bsp = np.concatenate([x1, x1], axis=0)
+
+sol_bsp = ot.solve_sample(x1_bsp, x2, method="bsp", n_projections=10)
+bsp_value = sol_bsp.value
+bsp_sparse_plan = sol_bsp.sparse_plan
+
+# sphinx_gallery_start_ignore
+
+pl.figure(1, (3, 3))
+plot2D_samples_mat(x1_bsp, x2, bsp_sparse_plan)
+pl.plot(x1_bsp[:, 0], x1_bsp[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("BSP OT plan loss={:.3f}".format(bsp_value))
+pl.show()
+
+# sphinx_gallery_end_ignore
+# %%
 # Comparing all OT plans
 # ----------------------
 #