[MRG] Linear Circular OT (#736)

* Linear Circular OT

* update LCOT

* update LCOT

* Add doc

* typo example

* Add LSSOT

* tests lcot and lssw

* Tests LCOT

* skip test tensorflow

* catch warnings divide by 0

* Update examples/plot_compute_wasserstein_circle.py

* fix doc

* fix number citations

* fix comment warning

---------

Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>

Author: clbonet

README.md

@@ -440,3 +440,7 @@ Artificial Intelligence.
 [76] Chapel, L., Tavenard, R. (2025). [One for all and all for one: Efficient computation of partial Wasserstein distances on the line](https://iclr.cc/virtual/2025/poster/28547). In International Conference on Learning Representations.
 
 [77] Tanguy, Eloi and Delon, Julie and Gozlan, Nathaël (2024). [Computing Barycentres of Measures for Generic Transport Costs](https://arxiv.org/abs/2501.04016). arXiv preprint 2501.04016 (2024)
+
+[78] Martin, R. D., Medri, I., Bai, Y., Liu, X., Yan, K., Rohde, G. K., & Kolouri, S. (2024). [LCOT: Linear Circular Optimal Transport](https://openreview.net/forum?id=49z97Y9lMq). International Conference on Learning Representations.
+
+[79] Liu, X., Bai, Y., Martín, R. D., Shi, K., Shahbazi, A., Landman, B. A., Chang, C., & Kolouri, S. (2025). [Linear Spherical Sliced Optimal Transport: A Fast Metric for Comparing Spherical Data](https://openreview.net/forum?id=fgUFZAxywx). International Conference on Learning Representations.

RELEASES.md

@@ -20,6 +20,7 @@
 - Backend implementation of `ot.dist` for (PR #701)
 - Updated documentation Quickstart guide and User guide with new API (PR #726)
 - Fix jax version for auto-grad (PR #732)
+- Added `ot.solver_1d.linear_circular_ot` and `ot.sliced.linear_sliced_wasserstein_sphere` (PR #736)
 - Implement 1d solver for partial optimal transport (PR #741)
 - Fix reg_div function compatibility with numpy in `ot.unbalanced.lbfgsb_unbalanced` via new function `ot.utils.fun_to_numpy` (PR #731)
 - Added to each example in the examples gallery the information about the release version in which it was introduced (PR #743)
@@ -36,6 +37,8 @@
 - Clean documentation for `ot.gromov.gromov_wasserstein` (PR #737)
 - Debug wheels building (PR #739)
 - Fix doc for projection sparse simplex (PR #734, PR #746)
+- Changed the default behavior of `ot.lp.solver_1d.wasserstein_circle` (Issue #738)
+- Avoid raising unnecessary warnings in `ot.lp.solver_1d.binary_search_circle` (Issue #738)
 
 ## 0.9.5
 

examples/backends/plot_ssw_unif_torch.py

@@ -10,7 +10,7 @@
 .. math::
      \min_{x} SSW_2(\nu, \frac{1}{n}\sum_{i=1}^n \delta_{x_i})
 
-where :math:`\nu=\mathrm{Unif}(S^1)`.
+where :math:`\nu=\mathrm{Unif}(S^{d-1})`.
 
 """
 
@@ -46,15 +46,18 @@
 
 
 def plot_sphere(ax):
-    xlist = np.linspace(-1.0, 1.0, 50)
-    ylist = np.linspace(-1.0, 1.0, 50)
-    r = np.linspace(1.0, 1.0, 50)
-    X, Y = np.meshgrid(xlist, ylist)
+    # Create a sphere using spherical coordinates
+    phi = np.linspace(0, 2 * np.pi, 100)
+    theta = np.linspace(0, np.pi, 100)
+    phi, theta = np.meshgrid(phi, theta)
 
-    Z = np.sqrt(np.maximum(r**2 - X**2 - Y**2, 0))
+    # Compute the spherical coordinates
+    X = np.sin(theta) * np.cos(phi)
+    Y = np.sin(theta) * np.sin(phi)
+    Z = np.cos(theta)
 
+    # Plot the wireframe
     ax.plot_wireframe(X, Y, Z, color="gray", alpha=0.3)
-    ax.plot_wireframe(X, Y, -Z, color="gray", alpha=0.3)  # Now plot the bottom half
 
 
 # plot the distributions

examples/sliced-wasserstein/plot_compute_wasserstein_circle.py

@@ -102,17 +102,23 @@ def pdf_von_Mises(theta, mu, kappa):
 
 L_w2_circle = np.zeros((n_try, 200))
 L_w2 = np.zeros((n_try, 200))
+L_lcot = np.zeros((n_try, 200))
 
 for i in range(n_try):
     w2_circle = ot.wasserstein_circle(xs2.T, xts2[i].T, p=2)
     w2 = ot.wasserstein_1d(xs2.T, xts2[i].T, p=2)
+    w_lcot = ot.linear_circular_ot(xs2.T, xts2[i].T)
 
     L_w2_circle[i] = w2_circle
     L_w2[i] = w2
+    L_lcot[i] = w_lcot
 
 m_w2_circle = np.mean(L_w2_circle, axis=0)
 std_w2_circle = np.std(L_w2_circle, axis=0)
 
+m_w2_lcot = np.mean(L_lcot, axis=0)
+std_w2_lcot = np.std(L_lcot, axis=0)
+
 m_w2 = np.mean(L_w2, axis=0)
 std_w2 = np.std(L_w2, axis=0)
 
@@ -128,6 +134,13 @@ def pdf_von_Mises(theta, mu, kappa):
 pl.fill_between(
     mu_targets / (2 * np.pi), m_w2 - 2 * std_w2, m_w2 + 2 * std_w2, alpha=0.5
 )
+pl.plot(mu_targets / (2 * np.pi), m_w2_lcot, label="Linear COT")
+pl.fill_between(
+    mu_targets / (2 * np.pi),
+    m_w2_lcot - 2 * std_w2_lcot,
+    m_w2_lcot + 2 * std_w2_lcot,
+    alpha=0.5,
+)
 pl.vlines(
     x=[mu1 / (2 * np.pi)],
     ymin=0,
@@ -159,15 +172,23 @@ def pdf_von_Mises(theta, mu, kappa):
         xts[i, k] = xt / (2 * np.pi)
 
 L_w2 = np.zeros((n_try, 100))
+L_lcot = np.zeros((n_try, 100))
 for i in range(n_try):
     L_w2[i] = ot.semidiscrete_wasserstein2_unif_circle(xts[i].T)
+    L_lcot[i] = ot.linear_circular_ot(xts[i].T)
 
 m_w2 = np.mean(L_w2, axis=0)
 std_w2 = np.std(L_w2, axis=0)
 
+m_lcot = np.mean(L_lcot, axis=0)
+std_lcot = np.std(L_lcot, axis=0)
+
 pl.figure(1)
-pl.plot(kappas, m_w2)
+pl.plot(kappas, m_w2, label="Wasserstein")
 pl.fill_between(kappas, m_w2 - std_w2, m_w2 + std_w2, alpha=0.5)
+pl.plot(kappas, m_lcot, label="LCOT")
+pl.fill_between(kappas, m_lcot - std_lcot, m_lcot + std_lcot, alpha=0.5)
+pl.legend()
 pl.title(r"Evolution of $W_2^2(vM(0,\kappa), Unif(S^1))$")
 pl.xlabel(r"$\kappa$")
 pl.show()

ot/__init__.py

@@ -48,6 +48,7 @@
     binary_search_circle,
     wasserstein_circle,
     semidiscrete_wasserstein2_unif_circle,
+    linear_circular_ot,
 )
 from .bregman import sinkhorn, sinkhorn2, barycenter
 from .unbalanced import sinkhorn_unbalanced, barycenter_unbalanced, sinkhorn_unbalanced2
@@ -57,6 +58,7 @@
     max_sliced_wasserstein_distance,
     sliced_wasserstein_sphere,
     sliced_wasserstein_sphere_unif,
+    linear_sliced_wasserstein_sphere,
 )
 from .gromov import (
     gromov_wasserstein,
@@ -105,6 +107,7 @@
     "sinkhorn_unbalanced2",
     "sliced_wasserstein_distance",
     "sliced_wasserstein_sphere",
+    "linear_sliced_wasserstein_sphere",
     "gromov_wasserstein",
     "gromov_wasserstein2",
     "gromov_barycenters",
@@ -129,6 +132,7 @@
     "binary_search_circle",
     "wasserstein_circle",
     "semidiscrete_wasserstein2_unif_circle",
+    "linear_circular_ot",
     "sliced_wasserstein_sphere_unif",
     "lowrank_sinkhorn",
     "lowrank_gromov_wasserstein_samples",

ot/lp/__init__.py

@@ -29,6 +29,7 @@
     binary_search_circle,
     wasserstein_circle,
     semidiscrete_wasserstein2_unif_circle,
+    linear_circular_ot,
 )
 
 __all__ = [
@@ -45,6 +46,7 @@
     "binary_search_circle",
     "wasserstein_circle",
     "semidiscrete_wasserstein2_unif_circle",
+    "linear_circular_ot",
     "dmmot_monge_1dgrid_loss",
     "dmmot_monge_1dgrid_optimize",
     "check_number_threads",