diff options
Diffstat (limited to 'ot/lp/__init__.py')
-rw-r--r-- | ot/lp/__init__.py | 10 |
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) |