add test sinkhorn+log

1 files changed, 13 insertions, 1 deletions

diff --git a/ot/bregman.py b/ot/bregman.py
index 43340f7..013bc33 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -120,7 +120,8 @@ def sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000,
         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