dr +gpu numpy assert

2 files changed, 19 insertions, 19 deletions

diff --git a/test/test_dr.py b/test/test_dr.py
index e3d1e6b..bdb920e 100644
--- a/test/test_dr.py
+++ b/test/test_dr.py
@@ -29,7 +29,7 @@ def test_fda():
 
     projfda(xs)
 
-    assert np.allclose(np.sum(Pfda**2, 0), np.ones(p))
+    np.testing.assert_allclose(np.sum(Pfda**2, 0), np.ones(p))
 
 
 @pytest.mark.skipif(nogo, reason="Missing modules (autograd or pymanopt)")
@@ -51,4 +51,4 @@ def test_wda():
 
     projwda(xs)
 
-    assert np.allclose(np.sum(Pwda**2, 0), np.ones(p))
+    np.testing.assert_allclose(np.sum(Pwda**2, 0), np.ones(p))
diff --git a/test/test_gpu.py b/test/test_gpu.py
index 5184a6c..7ae159b 100644
--- a/test/test_gpu.py
+++ b/test/test_gpu.py
@@ -16,14 +16,14 @@ def test_gpu_sinkhorn():
 
     np.random.seed(0)
 
-    def describeRes(r):
+    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 in [50, 100, 500, 1000]:
-        print(n)
-        a = np.random.rand(n // 4, 100)
-        b = np.random.rand(n, 100)
+    for n_samples in [50, 100, 500, 1000]:
+        print(n_samples)
+        a = np.random.rand(n_samples // 4, 100)
+        b = np.random.rand(n_samples, 100)
         time1 = time.time()
         transport = ot.da.OTDA_sinkhorn()
         transport.fit(a, b)
@@ -34,26 +34,26 @@ def test_gpu_sinkhorn():
         G2 = transport.G
         time3 = time.time()
         print("Normal sinkhorn, time: {:6.2f} sec ".format(time2 - time1))
-        describeRes(G1)
+        describe_res(G1)
         print("   GPU sinkhorn, time: {:6.2f} sec ".format(time3 - time2))
-        describeRes(G2)
+        describe_res(G2)
 
-        assert np.allclose(G1, G2, rtol=1e-5, atol=1e-5)
+        np.testing.assert_allclose(G1, G2, rtol=1e-5, atol=1e-5)
 
 
 @pytest.mark.skipif(nogpu, reason="No GPU available")
 def test_gpu_sinkhorn_lpl1():
     np.random.seed(0)
 
-    def describeRes(r):
+    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 in [50, 100, 500]:
-        print(n)
-        a = np.random.rand(n // 4, 100)
-        labels_a = np.random.randint(10, size=(n // 4))
-        b = np.random.rand(n, 100)
+    for n_samples in [50, 100, 500]:
+        print(n_samples)
+        a = np.random.rand(n_samples // 4, 100)
+        labels_a = np.random.randint(10, size=(n_samples // 4))
+        b = np.random.rand(n_samples, 100)
         time1 = time.time()
         transport = ot.da.OTDA_lpl1()
         transport.fit(a, labels_a, b)
@@ -65,9 +65,9 @@ def test_gpu_sinkhorn_lpl1():
         time3 = time.time()
         print("Normal sinkhorn lpl1, time: {:6.2f} sec ".format(
             time2 - time1))
-        describeRes(G1)
+        describe_res(G1)
         print("   GPU sinkhorn lpl1, time: {:6.2f} sec ".format(
             time3 - time2))
-        describeRes(G2)
+        describe_res(G2)
 
-        assert np.allclose(G1, G2, rtol=1e-5, atol=1e-5)
+        np.testing.assert_allclose(G1, G2, rtol=1e-5, atol=1e-5)