Added optional maximal number of iteration

1 files changed, 5 insertions, 5 deletions

diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py
index db3da78..673242d 100644
--- a/ot/lp/__init__.py
+++ b/ot/lp/__init__.py
@@ -11,7 +11,7 @@ import multiprocessing
 
 
 
-def emd(a, b, M):
+def emd(a, b, M, max_iter=-1):
     """Solves the Earth Movers distance problem and returns the OT matrix
 
 
@@ -80,9 +80,9 @@ def emd(a, b, M):
     if len(b) == 0:
         b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1]
 
-    return emd_c(a, b, M)
+    return emd_c(a, b, M, max_iter)
 
-def emd2(a, b, M,processes=multiprocessing.cpu_count()):
+def emd2(a, b, M,processes=multiprocessing.cpu_count(), max_iter=-1):
     """Solves the Earth Movers distance problem and returns the loss 
 
     .. math::
@@ -151,12 +151,12 @@ def emd2(a, b, M,processes=multiprocessing.cpu_count()):
         b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1]
         
     if len(b.shape)==1:
-        return emd2_c(a, b, M)
+        return emd2_c(a, b, M, max_iter)
     else:
         nb=b.shape[1]
         #res=[emd2_c(a,b[:,i].copy(),M) for i in range(nb)]
         def f(b):
-            return emd2_c(a,b,M)
+            return emd2_c(a,b,M, max_iter)
         res= parmap(f, [b[:,i] for i in range(nb)],processes)
         return np.array(res)