update sliced plans with sparse matrix for tf compatibility

Author: lchapel

ot/sliced.py

@@ -756,6 +756,12 @@ def sliced_plans(
     assert X.ndim == 2, f"X must be a 2d array, got {X.ndim}d array instead"
     assert Y.ndim == 2, f"Y must be a 2d array, got {Y.ndim}d array instead"
 
+    assert metric in ("minkowski", "euclidean", "cityblock", "sqeuclidean"), (
+        "Sliced plans work only with metrics "
+        + "from the following list: "
+        + "`['sqeuclidean', 'minkowski', 'cityblock', 'euclidean']`"
+    )
+
     assert (
         X.shape[1] == Y.shape[1]
     ), f"X ({X.shape}) and Y ({Y.shape}) must have the same number of columns"
@@ -776,7 +782,7 @@ def sliced_plans(
     do_draw_thetas = thetas is None
     if do_draw_thetas:  # create thetas (n_proj, d)
         assert n_proj is not None, "n_proj must be specified if thetas is None"
-        thetas = get_random_projections(d, n_proj, backend=nx).T
+        thetas = get_random_projections(d, n_proj, backend=nx, type_as=X).T
 
         if warm_theta is not None:
             thetas = nx.concatenate([thetas, warm_theta[:, None].T], axis=0)
@@ -787,8 +793,7 @@ def sliced_plans(
     X_theta = X @ thetas.T  # shape (n, n_proj)
     Y_theta = Y @ thetas.T  # shape (m, n_proj)
 
-    if is_perm:
-        # we compute maps (permutations)
+    if is_perm:  # we compute maps (permutations)
         # sigma[:, i_proj] is a permutation sorting X_theta[:, i_proj]
         sigma = nx.argsort(X_theta, axis=0)  # (n, n_proj)
         tau = nx.argsort(Y_theta, axis=0)  # (m, n_proj)
@@ -803,17 +808,12 @@ def sliced_plans(
                 )
                 for k in range(n_proj)
             ]
-        elif metric == "sqeuclidean":
+        else:  # metric = "sqeuclidean":
             costs = [
                 nx.sum((nx.sum((X[sigma[:, k]] - Y[tau[:, k]]) ** 2, axis=1)) / n)
                 for k in range(n_proj)
             ]
-        else:
-            raise ValueError(
-                "Sliced plans work only with metrics "
-                + "from the following list: "
-                + "`['sqeuclidean', 'minkowski', 'cityblock', 'euclidean']`"
-            )
+
         a = nx.ones(n) / n
         plan = [
             nx.coo_matrix(a, sigma[:, k], tau[:, k], shape=(n, m), type_as=a)
@@ -825,6 +825,35 @@ def sliced_plans(
             X_theta, Y_theta, a, b, p, require_sort=True, return_plan=True
         )
 
+        plan = plan.tocsr().tocoo()  # especially for tensorflow compatibility
+
+        if str(nx) == "jax":
+            plan = [nx.todense(plan[k]) for k in range(n_proj)]
+            if not dense:
+                warnings.warn(
+                    "JAX does not support sparse matrices, converting" "to dense"
+                )
+
+            costs = [
+                nx.sum(
+                    (
+                        (
+                            nx.sum(
+                                nx.abs(
+                                    X[np.nonzero(plan[k])[0]]
+                                    - Y[np.nonzero(plan[k])[1]]
+                                )
+                                ** p,
+                                axis=1,
+                            )
+                        )
+                        ** (1 / p)
+                    )
+                    * plan[np.nonzero(plan[k])]
+                )
+                for k in range(n_proj)
+            ]
+
         if metric in ("minkowski", "euclidean", "cityblock"):
             costs = [
                 nx.sum(
@@ -836,26 +865,17 @@ def sliced_plans(
                 )
                 for k in range(n_proj)
             ]
-        elif metric == "sqeuclidean":
+        else:  # metric = "sqeuclidean"
             costs = [
                 nx.sum(
                     (nx.sum((X[plan[k].row] - Y[plan[k].col]) ** 2, axis=1))
                     * plan[k].data
                 )
                 for k in range(n_proj)
             ]
-        else:
-            raise ValueError(
-                "Sliced plans work only with metrics "
-                + "from the following list: "
-                + "`['sqeuclidean', 'minkowski', 'cityblock', 'euclidean']`"
-            )
 
     if dense:
         plan = [nx.todense(plan[k]) for k in range(n_proj)]
-    elif str(nx) == "jax":
-        warnings.warn("JAX does not support sparse matrices, converting to dense")
-        plan = [nx.todense(plan[k]) for k in range(n_proj)]
 
     if log:
         log_dict = {"X_theta": X_theta, "Y_theta": Y_theta, "thetas": thetas}

test/test_sliced.py

@@ -782,6 +782,24 @@ def test_sliced_plans():
     # test with a warm theta
     ot.sliced.sliced_plans(x, y, n_proj=10, warm_theta=thetas[-1])
 
+    # test permutations
+    n = 5
+    m = 5
+    n_proj = 10
+    d = 2
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 2)
+    y = rng.randn(m, 2)
+
+    a = rng.uniform(0, 1, n)
+    a /= a.sum()
+    b = rng.uniform(0, 1, m)
+    b /= b.sum()
+
+    # test with the minkowski metric
+    ot.sliced.sliced_plans(x, y, metric="minkowski")
+
 
 def test_min_pivot_sliced():
     x = [1, 2]
@@ -924,7 +942,9 @@ def test_sliced_plans_backends(nx):
 
     x_b, y_b, a_b, b_b = nx.from_numpy(x, y, a, b)
 
-    thetas_b = ot.sliced.get_random_projections(d, n_proj, seed=0, backend=nx).T
+    thetas_b = ot.sliced.get_random_projections(
+        d, n_proj, seed=0, backend=nx, type_as=x
+    ).T
     thetas = nx.to_numpy(thetas_b)
 
     context = (