diff options
author | Antoine Rolet <antoine.rolet@gmail.com> | 2017-07-13 15:30:39 +0900 |
---|---|---|
committer | Antoine Rolet <antoine.rolet@gmail.com> | 2017-07-13 15:30:39 +0900 |
commit | 55a38f8253e5831105d2c329f4d8ed77686d1330 (patch) | |
tree | 6e8890ca1e0817c9043ce3927495c1f8679d5d6c /ot/lp/__init__.py | |
parent | 0e86d1bdbc0dcf7ffdb943637f62df5de4612ad0 (diff) |
Added optional maximal number of iteration
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) |