diff options
-rw-r--r-- | ot/lp/__init__.py | 12 | ||||
-rw-r--r-- | test/test_ot.py | 6 |
2 files changed, 12 insertions, 6 deletions
diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index f2eaa2b..d0f682b 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -107,7 +107,7 @@ def emd(a, b, M, num_iter_max=100000, log=False): return G -def emd2(a, b, M, processes=multiprocessing.cpu_count(), num_iter_max=100000, log=False): +def emd2(a, b, M, processes=multiprocessing.cpu_count(), num_iter_max=100000, log=False, return_matrix=False): """Solves the Earth Movers distance problem and returns the loss .. math:: @@ -134,6 +134,11 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), num_iter_max=100000, lo num_iter_max : int, optional (default=100000) The maximum number of iterations before stopping the optimization algorithm if it has not converged. + log: boolean, optional (default=False) + If True, returns a dictionary containing the cost and dual + variables. Otherwise returns only the optimal transportation cost. + return_matrix: boolean, optional (default=False) + If True, returns the optimal transportation matrix in the log. Returns ------- @@ -181,12 +186,13 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), num_iter_max=100000, lo if len(b) == 0: b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1] - if log: + if log or return_matrix: def f(b): G, cost, u, v, resultCode = emd_c(a, b, M, num_iter_max) result_code_string = check_result(resultCode) log = {} - log['G'] = G + if return_matrix: + log['G'] = G log['u'] = u log['v'] = v log['warning'] = result_code_string diff --git a/test/test_ot.py b/test/test_ot.py index e05e8aa..ea6d9dc 100644 --- a/test/test_ot.py +++ b/test/test_ot.py @@ -103,7 +103,7 @@ def test_emd2_multi(): # emd loss multipro proc with log ot.tic() - emdn = ot.emd2(a, b, M, log=True) + emdn = ot.emd2(a, b, M, log=True, return_matrix=True) ot.toc('multi proc : {} s') for i in range(len(emdn)): @@ -140,8 +140,8 @@ def test_warnings(): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") print('Computing {} EMD '.format(1)) - ot.emd(a, b, M, num_iter_max=1) - assert "num_iter_max" in str(w[-1].message) + ot.emd(a, b, M, numItermax=1) + assert "numItermax" in str(w[-1].message) assert len(w) == 1 a[0] = 100 print('Computing {} EMD '.format(2)) |