From 88c62c39a9623e8b58ebb776a9deddc96b43b4a0 Mon Sep 17 00:00:00 2001 From: arolet Date: Fri, 21 Jul 2017 12:12:48 +0900 Subject: Added dual variables computations --- ot/lp/EMD.h | 3 ++- ot/lp/EMD_wrapper.cpp | 6 ++++-- ot/lp/__init__.py | 11 +++++++---- ot/lp/emd_wrap.pyx | 22 ++++++++++++++++------ 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/ot/lp/EMD.h b/ot/lp/EMD.h index 59a5af8..fb7feca 100644 --- a/ot/lp/EMD.h +++ b/ot/lp/EMD.h @@ -24,6 +24,7 @@ using namespace lemon; typedef unsigned int node_id_type; -void EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double *cost, int max_iter); +void EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, + double* alpha, double* beta, double *cost, int max_iter); #endif diff --git a/ot/lp/EMD_wrapper.cpp b/ot/lp/EMD_wrapper.cpp index cc13230..6bda6a7 100644 --- a/ot/lp/EMD_wrapper.cpp +++ b/ot/lp/EMD_wrapper.cpp @@ -15,8 +15,8 @@ #include "EMD.h" -void EMD_wrap(int n1, int n2, double *X, double *Y, - double *D, double *G, double *cost, int max_iter) { +void EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G, + double* alpha, double* beta, double *cost, int max_iter) { // beware M and C anre strored in row major C style!!! int n, m, i,cur; @@ -99,6 +99,8 @@ void EMD_wrap(int n1, int n2, double *X, double *Y, int i = di.source(a); int j = di.target(a); *(G+indI[i]*n2+indJ[j-n]) = net.flow(a); + *(alpha + indI[i]) = net.potential(i); + *(beta + indJ[j-n]) = net.potential(j); } }; diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index 673242d..915a18c 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -11,7 +11,7 @@ import multiprocessing -def emd(a, b, M, max_iter=-1): +def emd(a, b, M, dual_variables=False, max_iter=-1): """Solves the Earth Movers distance problem and returns the OT matrix @@ -80,7 +80,10 @@ def emd(a, b, M, max_iter=-1): if len(b) == 0: b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1] - return emd_c(a, b, M, max_iter) + G, alpha, beta = emd_c(a, b, M, max_iter) + if dual_variables: + return G, alpha, beta + return G def emd2(a, b, M,processes=multiprocessing.cpu_count(), max_iter=-1): """Solves the Earth Movers distance problem and returns the loss @@ -151,12 +154,12 @@ def emd2(a, b, M,processes=multiprocessing.cpu_count(), max_iter=-1): b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1] if len(b.shape)==1: - return emd2_c(a, b, M, max_iter) + return emd2_c(a, b, M, max_iter)[0] 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, max_iter) + return emd2_c(a,b,M, max_iter)[0] res= parmap(f, [b[:,i] for i in range(nb)],processes) return np.array(res) diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx index c4ba125..813596f 100644 --- a/ot/lp/emd_wrap.pyx +++ b/ot/lp/emd_wrap.pyx @@ -12,7 +12,8 @@ cimport cython cdef extern from "EMD.h": - void EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double *cost, int max_iter) + void EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double *cost, + double* alpha, double* beta, int max_iter) @@ -58,6 +59,8 @@ def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mod cdef float cost=0 cdef np.ndarray[double, ndim=2, mode="c"] G=np.zeros([n1, n2]) + cdef np.ndarray[double, ndim=1, mode="c"] alpha=np.zeros(n1) + cdef np.ndarray[double, ndim=1, mode="c"] beta=np.zeros(n2) if not len(a): a=np.ones((n1,))/n1 @@ -65,10 +68,13 @@ def emd_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mod if not len(b): b=np.ones((n2,))/n2 + print alpha.size + print beta.size # calling the function - EMD_wrap(n1,n2, a.data, b.data, M.data, G.data, &cost, maxiter) + EMD_wrap(n1,n2, a.data, b.data, M.data, G.data, + alpha.data, beta.data, &cost, maxiter) - return G + return G, alpha, beta @cython.boundscheck(False) @cython.wraparound(False) @@ -112,15 +118,19 @@ def emd2_c( np.ndarray[double, ndim=1, mode="c"] a,np.ndarray[double, ndim=1, mo cdef float cost=0 cdef np.ndarray[double, ndim=2, mode="c"] G=np.zeros([n1, n2]) + + cdef np.ndarray[double, ndim = 1, mode = "c"] alpha = np.zeros([n1]) + cdef np.ndarray[double, ndim = 1, mode = "c"] beta = np.zeros([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, a.data, b.data, M.data, G.data, &cost, maxiter) + EMD_wrap(n1,n2, a.data, b.data, M.data, G.data, + alpha.data, beta.data, &cost, maxiter) + - return cost + return cost, alpha, beta -- cgit v1.2.3