diff options
Diffstat (limited to 'ot')
| -rw-r--r-- | ot/da.py | 2 | ||||
| -rw-r--r-- | ot/lp/__init__.py | 31 | ||||
| -rw-r--r-- | ot/lp/emd_wrap.pyx | 18 |
3 files changed, 27 insertions, 24 deletions
@@ -1370,7 +1370,7 @@ class EMDTransport(BaseTransport): # coupling estimation self.coupling_ = emd( - a=self.mu_s, b=self.mu_t, M=self.cost_, numItermax=self.max_iter + a=self.mu_s, b=self.mu_t, M=self.cost_, num_iter_max=self.max_iter ) return self diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index 0f40c19..ab7cb97 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -12,11 +12,11 @@ import multiprocessing import numpy as np # import compiled emd -from .emd_wrap import emd_c, checkResult +from .emd_wrap import emd_c, check_result from ..utils import parmap -def emd(a, b, M, numItermax=100000, log=False): +def emd(a, b, M, num_iter_max=100000, log=False): """Solves the Earth Movers distance problem and returns the OT matrix @@ -41,7 +41,7 @@ def emd(a, b, M, numItermax=100000, log=False): Target histogram (uniform weigth if empty list) M : (ns,nt) ndarray, float64 loss matrix - numItermax : int, optional (default=100000) + 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) @@ -54,7 +54,7 @@ def emd(a, b, M, numItermax=100000, log=False): Optimal transportation matrix for the given parameters log: dict If input log is true, a dictionary containing the cost and dual - variables + variables and exit status Examples @@ -94,20 +94,20 @@ def emd(a, b, M, numItermax=100000, log=False): if len(b) == 0: b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1] - G, cost, u, v, resultCode = emd_c(a, b, M, numItermax) - resultCodeString = checkResult(resultCode) + G, cost, u, v, result_code = emd_c(a, b, M, num_iter_max) + resultCodeString = check_result(result_code) if log: log = {} log['cost'] = cost log['u'] = u log['v'] = v log['warning'] = resultCodeString - log['resultCode'] = resultCode + log['result_code'] = result_code return G, log return G -def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log=False): +def emd2(a, b, M, processes=multiprocessing.cpu_count(), num_iter_max=100000, log=False): """Solves the Earth Movers distance problem and returns the loss .. math:: @@ -131,7 +131,7 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log= Target histogram (uniform weigth if empty list) M : (ns,nt) ndarray, float64 loss matrix - numItermax : int, optional (default=100000) + num_iter_max : int, optional (default=100000) The maximum number of iterations before stopping the optimization algorithm if it has not converged. @@ -139,6 +139,9 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log= ------- gamma: (ns x nt) ndarray Optimal transportation matrix for the given parameters + log: dict + If input log is true, a dictionary containing the cost and dual + variables and exit status Examples @@ -180,19 +183,19 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log= if log: def f(b): - G, cost, u, v, resultCode = emd_c(a, b, M, numItermax) - resultCodeString = checkResult(resultCode) + G, cost, u, v, resultCode = emd_c(a, b, M, num_iter_max) + resultCodeString = check_result(resultCode) log = {} log['G'] = G log['u'] = u log['v'] = v log['warning'] = resultCodeString - log['resultCode'] = resultCode + log['result_code'] = resultCode return [cost, log] else: def f(b): - G, cost, u, v, resultCode = emd_c(a, b, M, numItermax) - checkResult(resultCode) + G, cost, u, v, result_code = emd_c(a, b, M, num_iter_max) + check_result(result_code) return cost if len(b.shape) == 1: diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx index 19bcdd8..7ebdd2a 100644 --- a/ot/lp/emd_wrap.pyx +++ b/ot/lp/emd_wrap.pyx @@ -20,15 +20,15 @@ cdef extern from "EMD.h": cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED -def checkResult(resultCode): - if resultCode == OPTIMAL: +def check_result(result_code): + if result_code == OPTIMAL: return None - if resultCode == INFEASIBLE: + if result_code == INFEASIBLE: message = "Problem infeasible. Check that a and b are in the simplex" - elif resultCode == UNBOUNDED: + elif result_code == UNBOUNDED: message = "Problem unbounded" - elif resultCode == MAX_ITER_REACHED: + elif result_code == MAX_ITER_REACHED: message = "numItermax reached before optimality. Try to increase numItermax." warnings.warn(message) return message @@ -36,7 +36,7 @@ def checkResult(resultCode): @cython.boundscheck(False) @cython.wraparound(False) -def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mode="c"] b,np.ndarray[double, ndim=2, mode="c"] M, int numItermax): +def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mode="c"] b, np.ndarray[double, ndim=2, mode="c"] M, int num_iter_max): """ Solves the Earth Movers distance problem and returns the optimal transport matrix @@ -63,7 +63,7 @@ def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mod target histogram M : (ns,nt) ndarray, float64 loss matrix - numItermax : int + num_iter_max : int The maximum number of iterations before stopping the optimization algorithm if it has not converged. @@ -90,6 +90,6 @@ def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mod b=np.ones((n2,))/n2 # calling the function - cdef int resultCode = EMD_wrap(n1,n2,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, numItermax) + cdef int result_code = EMD_wrap(n1, n2, <double*> a.data, <double*> b.data, <double*> M.data, <double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, num_iter_max) - return G, cost, alpha, beta, resultCode + return G, cost, alpha, beta, result_code |