Added dual variables computations

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,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,<double*> &cost, maxiter)
+    EMD_wrap(n1,n2,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,
+                <double*> alpha.data, <double*> beta.data, <double*> &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,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,<double*> &cost, maxiter)
+    EMD_wrap(n1,n2,<double*> a.data,<double*> b.data,<double*> M.data,<double*> G.data,
+                <double*> alpha.data, <double*> beta.data, <double*> &cost, maxiter)
+
 
-    return cost
+    return cost, alpha, beta