summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorAntoine Rolet <antoine.rolet@gmail.com>2017-09-12 19:58:46 +0900
committerAntoine Rolet <antoine.rolet@gmail.com>2017-09-12 19:58:46 +0900
commitdd6f8260d01ce173ef3fe0c900112f0ed5288950 (patch)
tree22275dd64ff26577dc6c01a6e360afb62b2a3296
parent7c6169222979a7e82a83c118bc7117684258d0de (diff)
Made the return of the matrix optional in emd2
-rw-r--r--ot/lp/__init__.py12
-rw-r--r--test/test_ot.py6
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))