diff options
author | Rémi Flamary <remi.flamary@gmail.com> | 2017-03-14 10:30:45 +0100 |
---|---|---|
committer | Rémi Flamary <remi.flamary@gmail.com> | 2017-03-14 10:30:45 +0100 |
commit | a84f2c3e23edd1fa89975bd77b08672f518d5ca4 (patch) | |
tree | 1aadf05357949e6daec2c332eb900e93346ad465 /ot/lp/emd.pyx | |
parent | 84219d9bd87acd9bbb6d1a832cf4ccaee53fed0b (diff) |
add emd2+ multiproc
Diffstat (limited to 'ot/lp/emd.pyx')
-rw-r--r-- | ot/lp/emd.pyx | 60 |
1 files changed, 60 insertions, 0 deletions
diff --git a/ot/lp/emd.pyx b/ot/lp/emd.pyx index de2d4a9..46794ab 100644 --- a/ot/lp/emd.pyx +++ b/ot/lp/emd.pyx @@ -69,3 +69,63 @@ def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mod EMD_wrap(n1,n2,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,<double*> &cost) return G + +@cython.boundscheck(False) +@cython.wraparound(False) +def emd2_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): + """ + Solves the Earth Movers distance problem and returns the optimal transport loss + + gamm=emd(a,b,M) + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + + s.t. \gamma 1 = a + + \gamma^T 1= b + + \gamma\geq 0 + where : + + - M is the metric cost matrix + - a and b are the sample weights + + Parameters + ---------- + a : (ns,) ndarray, float64 + source histogram + b : (nt,) ndarray, float64 + target histogram + M : (ns,nt) ndarray, float64 + loss matrix + + + Returns + ------- + gamma: (ns x nt) ndarray + Optimal transportation matrix for the given parameters + + """ + cdef int n1= M.shape[0] + cdef int n2= M.shape[1] + + cdef float cost=0 + cdef np.ndarray[double, ndim=2, mode="c"] G=np.zeros([n1, n2]) + + if not len(a): + a=np.ones((n1,))/n1 + + if not len(b): + b=np.ones((n2,))/n2 + + # calling the function + EMD_wrap(n1,n2,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,<double*> &cost) + + cost=0 + for i in range(n1): + for j in range(n2): + cost+=G[i,j]*M[i,j] + + return cost + |