[WIP] Linear Circular OT

README.md

@@ -389,3 +389,7 @@ Artificial Intelligence.
 [74] Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). [Gradient descent algorithms for Bures-Wasserstein barycenters](https://proceedings.mlr.press/v125/chewi20a.html). In Conference on Learning Theory (pp. 1276-1304). PMLR.
 
 [75] Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://papers.neurips.cc/paper_files/paper/2021/hash/b9acb4ae6121c941324b2b1d3fac5c30-Abstract.html). Advances in Neural Information Processing Systems, 34, 22132-22145.
+
+[76] 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.
+
+[77] 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

@@ -17,6 +17,8 @@
 - 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` (PR #736)
+- Added `ot.sliced.linear_sliced_wasserstein_sphere` (PR #736)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)

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/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.mean(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

@@ -26,6 +26,7 @@
     binary_search_circle,
     wasserstein_circle,
     semidiscrete_wasserstein2_unif_circle,
+    linear_circular_ot,
 )
 
 __all__ = [
@@ -42,6 +43,7 @@
     "binary_search_circle",
     "wasserstein_circle",
     "semidiscrete_wasserstein2_unif_circle",
+    "linear_circular_ot",
     "dmmot_monge_1dgrid_loss",
     "dmmot_monge_1dgrid_optimize",
     "check_number_threads",

ot/lp/solver_1d.py

@@ -791,10 +791,12 @@ def binary_search_circle(
                 -1, 1
             )
 
-            mask_end = mask * (nx.abs(dCptm - dCmtp) > 0.001)
-            tc[mask_end > 0] = (
-                (Ctp - Ctm + tm * dCptm - tp * dCmtp) / (dCptm - dCmtp)
-            )[mask_end > 0]
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", category=RuntimeWarning)
+                mask_end = mask * (nx.abs(dCptm - dCmtp) > 0.001)
+                tc[mask_end > 0] = (
+                    (Ctp - Ctm + tm * dCptm - tp * dCmtp) / (dCptm - dCmtp)
+                )[mask_end > 0]
             done[nx.prod(mask, axis=-1) > 0] = 1
         elif nx.any(1 - done):
             tm[((1 - mask) * (dCp < 0)) > 0] = tc[((1 - mask) * (dCp < 0)) > 0]
@@ -933,8 +935,8 @@ def wasserstein_circle(
     eps=1e-6,
     require_sort=True,
 ):
-    r"""Computes the Wasserstein distance on the circle using either [45] for p=1 or
-    the binary search algorithm proposed in [44] otherwise.
+    r"""Computes the Wasserstein distance on the circle using either :ref:`[45] <references-wasserstein-circle>` for p=1 or
+    the binary search algorithm proposed in :ref:`[44] <references-wasserstein-circle>` otherwise.
     Samples need to be in :math:`S^1\cong [0,1[`. If they are on :math:`\mathbb{R}`,
     takes the value modulo 1.
     If the values are on :math:`S^1\subset\mathbb{R}^2`, it requires to first find the coordinates
@@ -996,17 +998,19 @@ def wasserstein_circle(
     >>> wasserstein_circle(u.T, v.T)
     array([0.1])
 
+
+    .. _references-wasserstein-circle:
     References
     ----------
     .. [44] Hundrieser, Shayan, Marcel Klatt, and Axel Munk. "The statistics of circular optimal transport." Directional Statistics for Innovative Applications: A Bicentennial Tribute to Florence Nightingale. Singapore: Springer Nature Singapore, 2022. 57-82.
     .. [45] Delon, Julie, Julien Salomon, and Andrei Sobolevski. "Fast transport optimization for Monge costs on the circle." SIAM Journal on Applied Mathematics 70.7 (2010): 2239-2258.
     """
     assert p >= 1, "The OT loss is only valid for p>=1, {p} was given".format(p=p)
 
-    if p == 1:
-        return wasserstein1_circle(
-            u_values, v_values, u_weights, v_weights, require_sort
-        )
+    # if p == 1:
+    #     return wasserstein1_circle(
+    #         u_values, v_values, u_weights, v_weights, require_sort
+    #     )
 
     return binary_search_circle(
         u_values,
@@ -1042,7 +1046,7 @@ def semidiscrete_wasserstein2_unif_circle(u_values, u_weights=None):
     .. math::
         u = \frac{\pi + \mathrm{atan2}(-x_2,-x_1)}{2\pi},
 
-    using e.g. ot.utils.get_coordinate_circle(x)
+    using e.g. ot.utils.get_coordinate_circle(x).
 
     Parameters
     ----------
@@ -1095,3 +1099,150 @@ def semidiscrete_wasserstein2_unif_circle(u_values, u_weights=None):
     cpt2 = nx.sum(u_values * u_weights * ns, axis=0)
 
     return cpt1 - u_mean**2 + cpt2 + 1 / 12
+
+
+def linear_circular_embedding(x, u_values, u_weights=None, require_sort=True):
+    r"""Returns the embedding :math:`\hat{\mu}(x)` of Linear Circular OT with reference
+    :math:`\eta=\mathrm{Unif}(S^1)` evaluated in :math:`x`.
+
+    For any :math:`x\in [0,1[`, the embedding is given by (see :ref:`[76] <references-lcot>`)
+
+    .. math``
+        \hat{\mu}(x) = F_{\mu}^{-1}\big(x - \int z\mathrm{d}\mu(z) + \frac12) - x.
+
+    Parameters
+    ----------
+    x : ndary, shape (m,)
+        Points in [0,1[ where to evaluate the embedding
+    u_values : ndarray, shape (n, ...)
+        samples in the source domain (coordinates on [0,1[)
+    u_weights : ndarray, shape (n, ...), optional
+        samples weights in the source domain
+
+    Returns
+    -------
+    embedding: ndarray of shape (m, ...)
+        Embedding evaluated at :math:`x`
+
+    .. _references-lcot:
+    References
+    ----------
+    .. [76] Martin, R. D., Medri, I., Bai, Y., Liu, X., Yan, K., Rohde, G. K., & Kolouri, S. (2024). LCOT: Linear Circular Optimal Transport. International Conference on Learning Representations.
+    """
+    if u_weights is not None:
+        nx = get_backend(u_values, u_weights)
+    else:
+        nx = get_backend(u_values)
+
+    n = u_values.shape[0]
+    u_values = u_values % 1
+
+    if len(u_values.shape) == 1:
+        u_values = nx.reshape(u_values, (n, 1))
+
+    if u_weights is None:
+        u_weights = nx.full(u_values.shape, 1.0 / n, type_as=u_values)
+    elif u_weights.ndim != u_values.ndim:
+        u_weights = nx.repeat(u_weights[..., None], u_values.shape[-1], -1)
+
+    if require_sort:
+        u_sorter = nx.argsort(u_values, 0)
+        u_values = nx.take_along_axis(u_values, u_sorter, 0)
+        u_weights = nx.take_along_axis(u_weights, u_sorter, 0)
+
+    u_cdf = nx.cumsum(u_weights, 0)
+    u_cdf = nx.zero_pad(u_cdf, [(1, 0), (0, 0)])
+
+    q_s = (
+        x[:, None] - nx.sum(u_values * u_weights, axis=0)[None] + 0.5
+    )  # shape (m, ...)
+
+    u_quantiles = quantile_function(q_s % 1, u_cdf, u_values)
+
+    return (u_quantiles - x[:, None]) % 1
+
+
+def linear_circular_ot(u_values, v_values=None, u_weights=None, v_weights=None):
+    r"""Computes the Linear Circular Optimal Transport distance from :ref:`[76] <references-lcot>` using :math:`\eta=\mathrm{Unif}(S^1)`
+    as reference measure.
+    Samples need to be in :math:`S^1\cong [0,1[`. If they are on :math:`\mathbb{R}`,
+    takes the value modulo 1.
+    If the values are on :math:`S^1\subset\mathbb{R}^2`, it is required to first find the coordinates
+    using e.g. the atan2 function.
+
+    General loss returned:
+
+    .. math::
+        \mathrm{LCOT}_2^2(\mu, \nu) = \int_0^1 d_{S^1}\big(\hat{\mu}(t), \hat{\nu}(t)\big)^2\ \mathrm{d}t
+
+    where :math:`\hat{\mu}(x)=F_{\mu}^{-1}(x-\int z\mathrm{d}\mu(z)+\frac12) - x` for all :math:`x\in [0,1[`,
+    and :math:`d_{S^1}(x,y)=\min(|x-y|, 1-|x-y|)` for :math:`x,y\in [0,1[`.
+
+    Parameters
+    ----------
+    u_values : ndarray, shape (n, ...)
+        samples in the source domain (coordinates on [0,1[)
+    v_values : ndarray, shape (n, ...), optional
+        samples in the target domain (coordinates on [0,1[), if None, compute distance against uniform distribution
+    u_weights : ndarray, shape (n, ...), optional
+        samples weights in the source domain
+    v_weights : ndarray, shape (n, ...), optional
+        samples weights in the target domain
+
+    Returns
+    -------
+    loss: float
+        Cost associated to the linear optimal transportation
+
+    Examples
+    --------
+    >>> u = np.array([[0.2,0.5,0.8]])%1
+    >>> v = np.array([[0.4,0.5,0.7]])%1
+    >>> linear_circular_ot(u.T, v.T)
+    array([0.0127])
+
+
+    .. _references-lcot:
+    References
+    ----------
+    .. [76] Martin, R. D., Medri, I., Bai, Y., Liu, X., Yan, K., Rohde, G. K., & Kolouri, S. (2024). LCOT: Linear Circular Optimal Transport. International Conference on Learning Representations.
+    """
+    if u_weights is not None:
+        nx = get_backend(u_values, u_weights)
+    else:
+        nx = get_backend(u_values)
+
+    n = u_values.shape[0]
+    u_values = u_values % 1
+
+    if len(u_values.shape) == 1:
+        u_values = nx.reshape(u_values, (n, 1))
+
+    if u_weights is None:
+        u_weights = nx.full(u_values.shape, 1.0 / n, type_as=u_values)
+    elif u_weights.ndim != u_values.ndim:
+        u_weights = nx.repeat(u_weights[..., None], u_values.shape[-1], -1)
+
+    unif_s1 = nx.linspace(0, 1, 101, type_as=u_values)[:-1]
+
+    emb_u = linear_circular_embedding(unif_s1, u_values, u_weights)
+
+    if v_values is None:
+        dist_u = nx.minimum(nx.abs(emb_u), 1 - nx.abs(emb_u))
+        return nx.mean(dist_u**2, axis=0)
+    else:
+        m = v_values.shape[0]
+        if len(v_values.shape) == 1:
+            v_values = nx.reshape(v_values, (m, 1))
+
+        if u_values.shape[1] != v_values.shape[1]:
+            raise ValueError(
+                "u and v must have the same number of batchs {} and {} respectively given".format(
+                    u_values.shape[1], v_values.shape[1]
+                )
+            )
+
+    emb_v = linear_circular_embedding(unif_s1, v_values, v_weights)
+
+    dist_uv = nx.minimum(nx.abs(emb_u - emb_v), 1 - nx.abs(emb_u - emb_v))
+    return nx.mean(dist_uv**2, axis=0)