update tests with backend

Author: lchapel

ot/sliced.py

@@ -695,7 +695,6 @@ def sliced_plans(
     n_proj=None,
     dense=False,
     log=False,
-    backend=None,
 ):
     r"""
     Computes all the permutations that sort the projections of two `(n, d)`
@@ -731,9 +730,6 @@ def sliced_plans(
         The number of projection directions. Required if thetas is None.
     log : bool, optional
         If True, returns additional logging information. Default is False.
-    backend : ot.backend, optional
-        Backend to use for computations. If None, the backend is inferred from
-        the input arrays. Default is None.
 
     Returns
     -------
@@ -746,7 +742,17 @@ def sliced_plans(
         Returned only if `log` is True.
     """
 
-    nx = get_backend(X, Y) if backend is None else backend
+    X, Y = list_to_array(X, Y)
+
+    if a is not None and b is not None and thetas is None:
+        nx = get_backend(X, Y, a, b)
+    elif a is not None and b is not None and thetas is not None:
+        nx = get_backend(X, Y, a, b, thetas)
+    elif a is None and b is None and thetas is not None:
+        nx = get_backend(X, Y, thetas)
+    else:
+        nx = get_backend(X, Y)
+
     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"
 
@@ -870,7 +876,6 @@ def min_pivot_sliced(
     dense=True,
     log=False,
     warm_theta=None,
-    backend=None,
 ):
     r"""
     Computes the cost and permutation associated to the min-Pivot Sliced
@@ -924,9 +929,6 @@ def min_pivot_sliced(
         If True, returns additional logging information. Default is False.
     warm_theta : array-like, shape (d,), optional
         A theta to add to the list of thetas. Default is None.
-    backend : ot.backend, optional
-        Backend to use for computations. If None, the backend is inferred from
-        the input arrays. Default is None.
 
     Returns
     -------
@@ -961,16 +963,24 @@ def min_pivot_sliced(
     2.125
     """
 
-    nx = get_backend(X, Y) if backend is None else backend
+    X, Y = list_to_array(X, Y)
+
+    if a is not None and b is not None and thetas is None:
+        nx = get_backend(X, Y, a, b)
+    elif a is not None and b is not None and thetas is not None:
+        nx = get_backend(X, Y, a, b, thetas)
+    elif a is None and b is None and thetas is not None:
+        nx = get_backend(X, Y, thetas)
+    else:
+        nx = get_backend(X, Y)
+
     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 (
         X.shape[1] == Y.shape[1]
     ), f"X ({X.shape}) and Y ({Y.shape}) must have the same number of columns"
 
-    nx = get_backend(X, Y) if backend is None else backend
-
     log_dict = {}
     G, costs, log_dict_plans = sliced_plans(
         X,
@@ -983,7 +993,6 @@ def min_pivot_sliced(
         n_proj=n_proj,
         warm_theta=warm_theta,
         log=True,
-        backend=nx,
     )
     pos_min = nx.argmin(costs)
     cost = costs[pos_min]
@@ -1024,7 +1033,6 @@ def expected_sliced(
     n_proj=None,
     dense=True,
     log=False,
-    backend=None,
     beta=0.0,
 ):
     r"""
@@ -1072,9 +1080,6 @@ def expected_sliced(
         format.
     log : bool, optional
         If True, returns additional logging information. Default is False.
-    backend : ot.backend, optional
-        Backend to use for computations. If None, the backend is inferred from
-        the input arrays. Default is None.
     beta : float, optional
         Inverse-temperature parameter which weights each projection's
         contribution to the expected plan. Default is 0 (uniform weighting).
@@ -1110,7 +1115,16 @@ def expected_sliced(
     2.625
     """
 
-    nx = get_backend(X, Y) if backend is None else backend
+    X, Y = list_to_array(X, Y)
+
+    if a is not None and b is not None and thetas is None:
+        nx = get_backend(X, Y, a, b)
+    elif a is not None and b is not None and thetas is not None:
+        nx = get_backend(X, Y, a, b, thetas)
+    elif a is None and b is None and thetas is not None:
+        nx = get_backend(X, Y, thetas)
+    else:
+        nx = get_backend(X, Y)
 
     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"
@@ -1130,7 +1144,7 @@ def expected_sliced(
 
     log_dict = {}
     G, costs, log_dict_plans = sliced_plans(
-        X, Y, a, b, metric, p, thetas, n_proj=n_proj, log=True, backend=nx
+        X, Y, a, b, metric, p, thetas, n_proj=n_proj, log=True
     )
     if log:
         log_dict = {"thetas": log_dict_plans["thetas"], "costs": costs, "G": G}

test/test_sliced.py

@@ -729,7 +729,7 @@ def test_linear_sliced_sphere_backend_type_devices(nx):
         np.testing.assert_almost_equal(sw_np, nx.to_numpy(valb))
 
 
-def test_sliced_permutations(nx):
+def test_sliced_permutations():
     n = 4
     n_proj = 10
     d = 2
@@ -738,15 +738,7 @@ def test_sliced_permutations(nx):
     x = rng.randn(n, 2)
     y = rng.randn(n, 2)
 
-    x_b, y_b = nx.from_numpy(x, y)
     thetas = ot.sliced.get_random_projections(d, n_proj, seed=0).T
-    thetas_b = nx.from_numpy(thetas)
-
-    plan, _ = ot.sliced.sliced_plans(x, y, thetas=thetas, dense=True)
-    plan_b, _, _ = ot.sliced.sliced_plans(
-        x_b, y_b, thetas=thetas_b, log=True, dense=True, backend=nx
-    )
-    np.testing.assert_almost_equal(plan, nx.to_numpy(plan_b))
 
     # test without provided thetas
     _, _ = ot.sliced.sliced_plans(x, y, n_proj=n_proj)
@@ -756,7 +748,7 @@ def test_sliced_permutations(nx):
         ot.sliced.sliced_plans(x[:, 1:], y, thetas=thetas)
 
 
-def test_sliced_plans(nx):
+def test_sliced_plans():
     x = [1, 2]
     with pytest.raises(AssertionError):
         ot.sliced.min_pivot_sliced(x, x, n_proj=2)
@@ -775,39 +767,23 @@ def test_sliced_plans(nx):
     b = rng.uniform(0, 1, m)
     b /= b.sum()
 
-    x_b, y_b = nx.from_numpy(x, y)
-    print(x_b)
-    t_X = torch.tensor(x_b)
-    t_Y = torch.tensor(y_b)
     thetas = ot.sliced.get_random_projections(d, n_proj, seed=0).T
-    thetas_b = nx.from_numpy(thetas)
+
+    # test with a and b not uniform
+    ot.sliced.sliced_plans(x, y, a, b, thetas=thetas, dense=True)
 
     # test with the minkowski metric
-    ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="minkowski")
+    ot.sliced.sliced_plans(x, y, thetas=thetas, metric="minkowski")
 
     # test with an unsupported metric
     with pytest.raises(ValueError):
-        ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="mahalanobis")
+        ot.sliced.sliced_plans(x, y, thetas=thetas, metric="mahalanobis")
 
     # test with a warm theta
-    ot.sliced.min_pivot_sliced(x, y, n_proj=10, warm_theta=thetas[-1])
+    ot.sliced.sliced_plans(x, y, n_proj=10, warm_theta=thetas[-1])
 
-    # test with a and b uniform
-    plan, _ = ot.sliced.sliced_plans(x, y, thetas=thetas, dense=True)
-    plan_b, _, _ = ot.sliced.sliced_plans(
-        x_b, y_b, thetas=thetas_b, log=True, dense=True, backend=nx
-    )
-    np.testing.assert_almost_equal(plan, nx.to_numpy(plan_b))
 
-    # test with a and b not uniform
-    plan, _ = ot.sliced.sliced_plans(x, y, a, b, thetas=thetas, dense=True)
-    plan_b, _, _ = ot.sliced.sliced_plans(
-        x_b, y_b, a, b, thetas=thetas_b, log=True, dense=True, backend=nx
-    )
-    np.testing.assert_almost_equal(plan, nx.to_numpy(plan_b))
-
-
-def test_min_pivot_sliced(nx):
+def test_min_pivot_sliced():
     x = [1, 2]
     with pytest.raises(AssertionError):
         ot.sliced.min_pivot_sliced(x, x, n_proj=2)
@@ -825,17 +801,13 @@ def test_min_pivot_sliced(nx):
     b = rng.uniform(0, 1, m)
     b /= b.sum()
 
-    x_b, y_b = nx.from_numpy(x, y)
     thetas = ot.sliced.get_random_projections(d, n_proj, seed=0).T
-    thetas_b = nx.from_numpy(thetas)
 
-    G, min_cost = ot.sliced.min_pivot_sliced(x, y, a, b, thetas=thetas, dense=True)
-    G_b, min_cost_b, _ = ot.sliced.min_pivot_sliced(
-        x_b, y_b, a, b, thetas=thetas_b, log=True, dense=True
-    )
+    # identity of the indiscernibles
+    _, min_cost = ot.min_pivot_sliced(x, x, a, a, n_proj=10)
+    np.testing.assert_almost_equal(min_cost, 0.0)
 
-    np.testing.assert_almost_equal(G, nx.to_numpy(G_b))
-    np.testing.assert_almost_equal(min_cost, nx.to_numpy(min_cost_b))
+    _, min_cost = ot.sliced.min_pivot_sliced(x, y, a, b, thetas=thetas, dense=True)
 
     # result should be an upper-bound of W2 and relatively close
     w2 = ot.emd2(a, b, ot.dist(x, y))
@@ -849,8 +821,30 @@ def test_min_pivot_sliced(nx):
     with pytest.raises(AssertionError):
         ot.sliced.min_pivot_sliced(x[:, 1:], y, thetas=thetas)
 
+    # test the logs
+    _, min_cost, log = ot.sliced.min_pivot_sliced(
+        x, y, a, b, thetas=thetas, dense=False, log=True
+    )
+    assert len(log) == 5
+    costs = log["costs"]
+    assert len(costs) == thetas.shape[0]
+    assert len(log["min_theta"]) == d
+    assert (log["thetas"] == thetas).all()
+    for c in costs:
+        assert c > 0
+
+    # test with the minkowski metric
+    ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="minkowski")
+
+    # test with an unsupported metric
+    with pytest.raises(ValueError):
+        ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="mahalanobis")
 
-def test_expected_sliced(nx):
+    # test with a warm theta
+    ot.sliced.min_pivot_sliced(x, y, n_proj=10, warm_theta=thetas[-1])
+
+
+def test_expected_sliced():
     x = [1, 2]
     with pytest.raises(AssertionError):
         ot.sliced.min_pivot_sliced(x, x, n_proj=2)
@@ -868,9 +862,67 @@ def test_expected_sliced(nx):
     b = rng.uniform(0, 1, m)
     b /= b.sum()
 
-    x_b, y_b = nx.from_numpy(x, y)
     thetas = ot.sliced.get_random_projections(d, n_proj, seed=0).T
-    thetas_b = nx.from_numpy(thetas)
+
+    _, expected_cost = ot.sliced.expected_sliced(x, y, a, b, dense=True, thetas=thetas)
+    # result should be a coarse upper-bound of W2
+    w2 = ot.emd2(a, b, ot.dist(x, y))
+    assert expected_cost >= w2
+    assert expected_cost <= 3 * w2
+
+    # test without provided thetas
+    ot.sliced.expected_sliced(x, y, n_proj=n_proj, log=True)
+
+    # test with invalid shapes
+    with pytest.raises(AssertionError):
+        ot.sliced.min_pivot_sliced(x[:, 1:], y, thetas=thetas)
+
+    # with a small temperature (i.e. large beta), the cost should be close
+    # to min_pivot
+    _, expected_cost = ot.sliced.expected_sliced(
+        x, y, a, b, thetas=thetas, dense=True, beta=100.0
+    )
+    _, min_cost = ot.sliced.min_pivot_sliced(x, y, a, b, thetas=thetas, dense=True)
+    np.testing.assert_almost_equal(expected_cost, min_cost, decimal=3)
+
+    # test the logs
+    _, min_cost, log = ot.sliced.expected_sliced(
+        x, y, a, b, thetas=thetas, dense=False, log=True
+    )
+    assert len(log) == 4
+    costs = log["costs"]
+    assert len(costs) == thetas.shape[0]
+    assert len(log["weights"]) == thetas.shape[0]
+    assert (log["thetas"] == thetas).all()
+    for c in costs:
+        assert c > 0
+
+    # test with the minkowski metric
+    ot.sliced.expected_sliced(x, y, thetas=thetas, metric="minkowski")
+
+    # test with an unsupported metric
+    with pytest.raises(ValueError):
+        ot.sliced.expected_sliced(x, y, thetas=thetas, metric="mahalanobis")
+
+
+def test_sliced_plans_backends(nx):
+    n = 10
+    m = 24
+    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()
+
+    x_b, y_b, a_b, b_b = nx.from_numpy(x, y, a, b)
+
+    thetas = ot.sliced.get_random_projections(d, n_proj, seed=0, backend=nx).T
+    thetas_t = nx.to_numpy(thetas)
 
     context = (
         nullcontext()
@@ -879,32 +931,25 @@ def test_expected_sliced(nx):
     )
 
     with context:
-        expected_plan, expected_cost = ot.sliced.expected_sliced(
-            x, y, a, b, dense=True, thetas=thetas
+        _, expected_cost_b = ot.sliced.expected_sliced(
+            x_b, y_b, a_b, b_b, dense=True, thetas=thetas_t
         )
-        expected_plan_b, expected_cost_b, _ = ot.sliced.expected_sliced(
-            x_b, y_b, a, b, thetas=thetas_b, dense=True, log=True
+        # result should be the same than numpy version
+        _, expected_cost = ot.sliced.expected_sliced(
+            x, y, a, b, dense=True, thetas=thetas
         )
+        np.testing.assert_almost_equal(expected_cost_b, expected_cost)
 
-        np.testing.assert_almost_equal(expected_plan, nx.to_numpy(expected_plan_b))
-        np.testing.assert_almost_equal(expected_cost, nx.to_numpy(expected_cost_b))
-
-        # result should be a coarse upper-bound of W2
-        w2 = ot.emd2(a, b, ot.dist(x, y))
-        assert expected_cost >= w2
-        assert expected_cost <= 3 * w2
-
-        # test without provided thetas
-        ot.sliced.expected_sliced(x, y, n_proj=n_proj, log=True)
+    # for min_pivot
+    _, min_cost_b = ot.sliced.min_pivot_sliced(
+        x_b, y_b, a_b, b_b, dense=True, thetas=thetas_t
+    )
+    # result should be the same than numpy version
+    _, min_cost = ot.sliced.min_pivot_sliced(x, y, a, b, dense=True, thetas=thetas)
+    np.testing.assert_almost_equal(min_cost_b, min_cost)
 
-        # test with invalid shapes
-        with pytest.raises(AssertionError):
-            ot.sliced.min_pivot_sliced(x[:, 1:], y, thetas=thetas)
+    # for sliced_plans
+    thetas = ot.sliced.get_random_projections(d, n_proj, seed=0, backend=nx).T
 
-        # with a small temperature (i.e. large beta), the cost should be close
-        # to min_pivot
-        _, expected_cost = ot.sliced.expected_sliced(
-            x, y, a, b, thetas=thetas, dense=True, beta=100.0
-        )
-        _, min_cost = ot.sliced.min_pivot_sliced(x, y, a, b, thetas=thetas, dense=True)
-        np.testing.assert_almost_equal(expected_cost, min_cost, decimal=3)
+    # test with the minkowski metric
+    ot.sliced.min_pivot_sliced(x, y, thetas=thetas, metric="minkowski")