add test sinkhorn+log

Author: rflamary

ot/bregman.py

@@ -120,7 +120,8 @@ def sink():
         print('Warning : unknown method using classic Sinkhorn Knopp')
 
         def sink():
-            return sinkhorn_knopp(a, b, M, reg, **kwargs)
+            return sinkhorn_knopp(a, b, M, reg, numItermax=numItermax,
+                                  stopThr=stopThr, verbose=verbose, log=log, **kwargs)
 
     return sink()
 
@@ -499,6 +500,15 @@ def greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False, log=
 
     """
 
+    a = np.asarray(a, dtype=np.float64)
+    b = np.asarray(b, dtype=np.float64)
+    M = np.asarray(M, dtype=np.float64)
+
+    if len(a) == 0:
+        a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0]
+    if len(b) == 0:
+        b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
+
     n = a.shape[0]
     m = b.shape[0]
 
@@ -514,7 +524,9 @@ def greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False, log=
     viol = G.sum(1) - a
     viol_2 = G.sum(0) - b
     stopThr_val = 1
+
     if log:
+        log = dict()
         log['u'] = u
         log['v'] = v
 

test/test_bregman.py

@@ -81,6 +81,31 @@ def test_sinkhorn_variants():
     print(G0, G_green)
 
 
+def test_sinkhorn_variants_log():
+    # test sinkhorn
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 2)
+    u = ot.utils.unif(n)
+
+    M = ot.dist(x, x)
+
+    G0, log0 = ot.sinkhorn(u, u, M, 1, method='sinkhorn', stopThr=1e-10, log=True)
+    Gs, logs = ot.sinkhorn(u, u, M, 1, method='sinkhorn_stabilized', stopThr=1e-10, log=True)
+    Ges, loges = ot.sinkhorn(
+        u, u, M, 1, method='sinkhorn_epsilon_scaling', stopThr=1e-10, log=True)
+    Gerr, logerr = ot.sinkhorn(u, u, M, 1, method='do_not_exists', stopThr=1e-10, log=True)
+    G_green, loggreen = ot.sinkhorn(u, u, M, 1, method='greenkhorn', stopThr=1e-10, log=True)
+
+    # check values
+    np.testing.assert_allclose(G0, Gs, atol=1e-05)
+    np.testing.assert_allclose(G0, Ges, atol=1e-05)
+    np.testing.assert_allclose(G0, Gerr)
+    np.testing.assert_allclose(G0, G_green, atol=1e-5)
+    print(G0, G_green)
+
+
 def test_bary():
 
     n_bins = 100  # nb bins