tests lcot and lssw

Author: clbonet

examples/plot_compute_wasserstein_circle.py

@@ -172,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

@@ -58,6 +58,7 @@
     max_sliced_wasserstein_distance,
     sliced_wasserstein_sphere,
     sliced_wasserstein_sphere_unif,
+    linear_sliced_wasserstein_sphere,
 )
 from .gromov import (
     gromov_wasserstein,
@@ -106,6 +107,7 @@
     "sinkhorn_unbalanced2",
     "sliced_wasserstein_distance",
     "sliced_wasserstein_sphere",
+    "linear_sliced_wasserstein_sphere",
     "gromov_wasserstein",
     "gromov_wasserstein2",
     "gromov_barycenters",

ot/lp/solver_1d.py

@@ -1228,7 +1228,19 @@ def linear_circular_ot(u_values, v_values=None, u_weights=None, v_weights=None):
     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)

ot/sliced.py

@@ -454,6 +454,8 @@ def sliced_wasserstein_sphere(
     ----------
     .. [46] Bonet, C., Berg, P., Courty, N., Septier, F., Drumetz, L., & Pham, M. T. (2023). Spherical sliced-wasserstein. International Conference on Learning Representations.
     """
+    d = X_s.shape[-1]
+
     if a is not None and b is not None:
         nx = get_backend(X_s, X_t, a, b)
     else:
@@ -470,11 +472,16 @@ def sliced_wasserstein_sphere(
     if nx.any(nx.abs(nx.sum(X_t**2, axis=-1) - 1) > 10 ** (-4)):
         raise ValueError("X_t is not on the sphere.")
 
-    Xps_coords, projections = projection_sphere_to_circle(
+    if projections is None:
+        projections = get_projections_sphere(
+            d, n_projections, seed=seed, backend=nx, type_as=X_s
+        )
+
+    Xps_coords, _ = projection_sphere_to_circle(
         X_s, n_projections=n_projections, projections=projections, seed=seed, backend=nx
     )
 
-    Xpt_coords, projections = projection_sphere_to_circle(
+    Xpt_coords, _ = projection_sphere_to_circle(
         X_t, n_projections=n_projections, projections=projections, seed=seed, backend=nx
     )
 
@@ -536,6 +543,8 @@ def sliced_wasserstein_sphere_unif(
     -----------
     .. [46] Bonet, C., Berg, P., Courty, N., Septier, F., Drumetz, L., & Pham, M. T. (2023). Spherical sliced-wasserstein. International Conference on Learning Representations.
     """
+    d = X_s.shape[-1]
+
     if a is not None:
         nx = get_backend(X_s, a)
     else:
@@ -544,7 +553,12 @@ def sliced_wasserstein_sphere_unif(
     if nx.any(nx.abs(nx.sum(X_s**2, axis=-1) - 1) > 10 ** (-4)):
         raise ValueError("X_s is not on the sphere.")
 
-    Xps_coords, projections = projection_sphere_to_circle(
+    if projections is None:
+        projections = get_projections_sphere(
+            d, n_projections, seed=seed, backend=nx, type_as=X_s
+        )
+
+    Xps_coords, _ = projection_sphere_to_circle(
         X_s, n_projections=n_projections, projections=projections, seed=seed, backend=nx
     )
 
@@ -616,6 +630,7 @@ def linear_sliced_wasserstein_sphere(
     ----------
     .. [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. International Conference on Learning Representations.
     """
+    d = X_s.shape[-1]
 
     if a is not None and b is not None:
         nx = get_backend(X_s, X_t, a, b)
@@ -633,12 +648,17 @@ def linear_sliced_wasserstein_sphere(
     if nx.any(nx.abs(nx.sum(X_t**2, axis=-1) - 1) > 10 ** (-4)):
         raise ValueError("X_t is not on the sphere.")
 
-    Xps_coords, projections = projection_sphere_to_circle(
+    if projections is None:
+        projections = get_projections_sphere(
+            d, n_projections, seed=seed, backend=nx, type_as=X_s
+        )
+
+    Xps_coords, _ = projection_sphere_to_circle(
         X_s, n_projections=n_projections, projections=projections, seed=seed, backend=nx
     )
 
     if X_t is not None:
-        Xpt_coords, projections = projection_sphere_to_circle(
+        Xpt_coords, _ = projection_sphere_to_circle(
             X_t,
             n_projections=n_projections,
             projections=projections,

test/test_1d_solver.py

@@ -355,3 +355,54 @@ def test_wasserstein_circle_bad_shape():
 
     with pytest.raises(ValueError):
         _ = ot.wasserstein_circle(u, v, p=1)
+
+
+def test_linear_circular_ot_devices(nx):
+    rng = np.random.RandomState(0)
+
+    n = 10
+    x = np.linspace(0, 1, n)
+    rho_u = np.abs(rng.randn(n))
+    rho_u /= rho_u.sum()
+    rho_v = np.abs(rng.randn(n))
+    rho_v /= rho_v.sum()
+
+    for tp in nx.__type_list__:
+        print(nx.dtype_device(tp))
+
+        xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v, type_as=tp)
+
+        lcot = ot.linear_circular_ot(xb, xb, rho_ub, rho_vb)
+
+        nx.assert_same_dtype_device(xb, lcot)
+
+
+def test_linear_circular_ot_bad_shape():
+    n = 20
+    m = 30
+    rng = np.random.RandomState(0)
+    u = rng.rand(n, 2)
+    v = rng.rand(m, 1)
+
+    with pytest.raises(ValueError):
+        _ = ot.linear_circular_ot(u, v)
+
+
+def test_linear_circular_ot_same_dist():
+    n = 20
+    rng = np.random.RandomState(0)
+    u = rng.rand(n)
+
+    lcot = ot.linear_circular_ot(u, u)
+    np.testing.assert_almost_equal(lcot, 0.0)
+
+
+def test_linear_circular_ot_different_dist():
+    n = 20
+    m = 30
+    rng = np.random.RandomState(0)
+    u = rng.rand(n)
+    v = rng.rand(m)
+
+    lcot = ot.linear_circular_ot(u, v)
+    assert lcot > 0.0

test/test_sliced.py

@@ -289,6 +289,30 @@ def test_projections_stiefel():
         np.matmul(P_T, P), np.array([np.eye(2) for k in range(n_projs)])
     )
 
+    rng = np.random.RandomState(0)
+
+    projections = ot.sliced.get_projections_sphere(3, n_projs, seed=rng)
+    projections_T = np.transpose(projections, [0, 2, 1])
+
+    np.testing.assert_almost_equal(
+        np.matmul(projections_T, projections),
+        np.array([np.eye(2) for k in range(n_projs)]),
+    )
+
+    # np.testing.assert_almost_equal(projections, P)
+
+
+def test_projections_sphere_to_circle():
+    rng = np.random.RandomState(0)
+
+    n_projs = 500
+    x = rng.randn(100, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    x_projs, _ = ot.sliced.projection_sphere_to_circle(x, n_projs)
+    assert x_projs.shape == (n_projs, 100)
+    assert np.all(x_projs >= 0) and np.all(x_projs < 1)
+
 
 def test_sliced_sphere_same_dist():
     n = 100
@@ -506,3 +530,153 @@ def test_sliced_sphere_unif_backend_type_devices(nx):
         valb = ot.sliced_wasserstein_sphere_unif(xb)
 
         nx.assert_same_dtype_device(xb, valb)
+
+
+def test_linear_sliced_sphere_same_dist():
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+    u = ot.utils.unif(n)
+
+    res = ot.linear_sliced_wasserstein_sphere(x, x, u, u, 10, seed=rng)
+    np.testing.assert_almost_equal(res, 0.0)
+
+
+def test_linear_sliced_sphere_same_proj():
+    n_projections = 10
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    y = rng.randn(n, 3)
+    y = y / np.sqrt(np.sum(y**2, -1, keepdims=True))
+
+    seed = 42
+
+    cost1, log1 = ot.linear_sliced_wasserstein_sphere(
+        x, y, seed=seed, n_projections=n_projections, log=True
+    )
+    cost2, log2 = ot.linear_sliced_wasserstein_sphere(
+        x, y, seed=seed, n_projections=n_projections, log=True
+    )
+
+    assert np.allclose(log1["projections"], log2["projections"])
+    assert np.isclose(cost1, cost2)
+
+
+def test_linear_sliced_sphere_bad_shapes():
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    y = rng.randn(n, 4)
+    y = y / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    u = ot.utils.unif(n)
+
+    with pytest.raises(ValueError):
+        _ = ot.linear_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng)
+
+
+def test_linear_sliced_sphere_values_on_the_sphere():
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    y = rng.randn(n, 4)
+
+    u = ot.utils.unif(n)
+
+    with pytest.raises(ValueError):
+        _ = ot.linear_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng)
+
+
+def test_linear_sliced_sphere_log():
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 4)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+    y = rng.randn(n, 4)
+    y = y / np.sqrt(np.sum(y**2, -1, keepdims=True))
+    u = ot.utils.unif(n)
+
+    res, log = ot.linear_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng, log=True)
+    assert len(log) == 2
+    projections = log["projections"]
+    projected_emds = log["projected_emds"]
+
+    assert projections.shape[0] == len(projected_emds) == 10
+    for emd in projected_emds:
+        assert emd > 0
+
+
+def test_linear_sliced_sphere_different_dists():
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    u = ot.utils.unif(n)
+    y = rng.randn(n, 3)
+    y = y / np.sqrt(np.sum(y**2, -1, keepdims=True))
+
+    res = ot.linear_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng)
+    assert res > 0.0
+
+
+def test_1d_linear_sliced_sphere_equals_emd():
+    n = 100
+    m = 120
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 2)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+    x_coords = (np.arctan2(-x[:, 1], -x[:, 0]) + np.pi) / (2 * np.pi)
+    a = rng.uniform(0, 1, n)
+    a /= a.sum()
+
+    y = rng.randn(m, 2)
+    y = y / np.sqrt(np.sum(y**2, -1, keepdims=True))
+    y_coords = (np.arctan2(-y[:, 1], -y[:, 0]) + np.pi) / (2 * np.pi)
+    u = ot.utils.unif(m)
+
+    res = ot.linear_sliced_wasserstein_sphere(x, y, a, u, 100, seed=42)
+    expected = ot.linear_circular_ot(x_coords.T, y_coords.T, a, u)
+
+    np.testing.assert_almost_equal(res**2, expected, decimal=5)
+
+
+def test_linear_sliced_sphere_backend_type_devices(nx):
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 3)
+    x = x / np.sqrt(np.sum(x**2, -1, keepdims=True))
+
+    y = rng.randn(2 * n, 3)
+    y = y / np.sqrt(np.sum(y**2, -1, keepdims=True))
+
+    sw_np, log = ot.linear_sliced_wasserstein_sphere(x, y, log=True)
+    P = log["projections"]
+
+    for tp in nx.__type_list__:
+        print(nx.dtype_device(tp))
+
+        xb, yb = nx.from_numpy(x, y, type_as=tp)
+
+        valb = ot.linear_sliced_wasserstein_sphere(
+            xb, yb, projections=nx.from_numpy(P, type_as=tp)
+        )
+
+        nx.assert_same_dtype_device(xb, valb)
+        np.testing.assert_almost_equal(sw_np, nx.to_numpy(valb))