update tests

1 files changed, 48 insertions, 41 deletions

diff --git a/test/test_gpu.py b/test/test_gpu.py
index 1e97c45..51a0cff 100644
--- a/test/test_gpu.py
+++ b/test/test_gpu.py
@@ -17,63 +17,70 @@ except ImportError:
 
 
 @pytest.mark.skipif(nogpu, reason="No GPU available")
-def test_gpu_sinkhorn():
+def test_gpu_dist():
 
     rng = np.random.RandomState(0)
 
-    def describe_res(r):
-        print("min:{:.3E}, max::{:.3E}, mean::{:.3E}, std::{:.3E}".format(
-            np.min(r), np.max(r), np.mean(r), np.std(r)))
-
     for n_samples in [50, 100, 500, 1000]:
         print(n_samples)
         a = rng.rand(n_samples // 4, 100)
         b = rng.rand(n_samples, 100)
-        time1 = time.time()
-        transport = ot.da.OTDA_sinkhorn()
-        transport.fit(a, b)
-        G1 = transport.G
-        time2 = time.time()
-        transport = ot.gpu.da.OTDA_sinkhorn()
-        transport.fit(a, b)
-        G2 = transport.G
-        time3 = time.time()
-        print("Normal sinkhorn, time: {:6.2f} sec ".format(time2 - time1))
-        describe_res(G1)
-        print("   GPU sinkhorn, time: {:6.2f} sec ".format(time3 - time2))
-        describe_res(G2)
-
-        np.testing.assert_allclose(G1, G2, rtol=1e-5, atol=1e-5)
+
+        M = ot.dist(a.copy(), b.copy())
+        M2 = ot.gpu.dist(a.copy(), b.copy())
+
+        np.testing.assert_allclose(M, M2, rtol=1e-10)
+
+        M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False)
 
 
 @pytest.mark.skipif(nogpu, reason="No GPU available")
-def test_gpu_sinkhorn_lpl1():
+def test_gpu_sinkhorn():
 
     rng = np.random.RandomState(0)
 
-    def describe_res(r):
-        print("min:{:.3E}, max:{:.3E}, mean:{:.3E}, std:{:.3E}"
-              .format(np.min(r), np.max(r), np.mean(r), np.std(r)))
+    for n_samples in [50, 100, 500, 1000]:
+        a = rng.rand(n_samples // 4, 100)
+        b = rng.rand(n_samples, 100)
+
+        wa = ot.unif(n_samples // 4)
+        wb = ot.unif(n_samples)
+
+        M = ot.dist(a.copy(), b.copy())
+        M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False)
+
+        reg = 1
+
+        G = ot.sinkhorn(wa, wb, M, reg)
+        G1 = ot.gpu.sinkhorn(wa, wb, M, reg)
+
+        np.testing.assert_allclose(G1, G, rtol=1e-10)
+
+        G2 = ot.gpu.sinkhorn(wa, wb, M2, reg, to_numpy=False)
+
+
+@pytest.mark.skipif(nogpu, reason="No GPU available")
+def test_gpu_sinkhorn_lpl1():
+
+    rng = np.random.RandomState(0)
 
     for n_samples in [50, 100, 500]:
         print(n_samples)
         a = rng.rand(n_samples // 4, 100)
         labels_a = np.random.randint(10, size=(n_samples // 4))
         b = rng.rand(n_samples, 100)
-        time1 = time.time()
-        transport = ot.da.OTDA_lpl1()
-        transport.fit(a, labels_a, b)
-        G1 = transport.G
-        time2 = time.time()
-        transport = ot.gpu.da.OTDA_lpl1()
-        transport.fit(a, labels_a, b)
-        G2 = transport.G
-        time3 = time.time()
-        print("Normal sinkhorn lpl1, time: {:6.2f} sec ".format(
-            time2 - time1))
-        describe_res(G1)
-        print("   GPU sinkhorn lpl1, time: {:6.2f} sec ".format(
-            time3 - time2))
-        describe_res(G2)
-
-        np.testing.assert_allclose(G1, G2, rtol=1e-3, atol=1e-3)
+
+        wa = ot.unif(n_samples // 4)
+        wb = ot.unif(n_samples)
+
+        M = ot.dist(a.copy(), b.copy())
+        M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False)
+
+        reg = 1
+
+        G = ot.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M, reg)
+        G1 = ot.gpu.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M, reg)
+
+        np.testing.assert_allclose(G1, G, rtol=1e-10)
+
+        G2 = ot.gpu.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M2, reg, to_numpy=False)