update tests

Author: rflamary

ot/gpu/bregman.py

@@ -90,14 +90,14 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
 
     """
 
-    a = cp.asarray(a, dtype=np.float64)
-    b = cp.asarray(b, dtype=np.float64)
-    M = cp.asarray(M, dtype=np.float64)
+    a = cp.asarray(a)
+    b = cp.asarray(b)
+    M = cp.asarray(M)
 
     if len(a) == 0:
-        a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0]
+        a = np.ones((M.shape[0],)) / M.shape[0]
     if len(b) == 0:
-        b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
+        b = np.ones((M.shape[1],)) / M.shape[1]
 
     # init data
     Nini = len(a)

test/test_gpu.py

@@ -17,63 +17,70 @@
 
 
 @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)