Merge pull request #29 from arolet/ot_dual_variables

Dual variables in EMD_wrapper

Author: rflamary

README.md

@@ -138,12 +138,12 @@ The contributors to this library are:
 * [Léo Gautheron](https://github.com/aje) (GPU implementation)
 * [Nathalie Gayraud](https://www.linkedin.com/in/nathalie-t-h-gayraud/?ppe=1)
 * [Stanislas Chambon](https://slasnista.github.io/)
+* [Antoine Rolet](https://arolet.github.io/)
 
 This toolbox benefit a lot from open source research and we would like to thank the following persons for providing some code (in various languages):
 
 * [Gabriel Peyré](http://gpeyre.github.io/) (Wasserstein Barycenters in Matlab)
 * [Nicolas Bonneel](http://liris.cnrs.fr/~nbonneel/) ( C++ code for EMD)
-* [Antoine Rolet](https://arolet.github.io/) ( Mex file for EMD )
 * [Marco Cuturi](http://marcocuturi.net/) (Sinkhorn Knopp in Matlab/Cuda)
 
 

ot/lp/EMD.h

@@ -26,9 +26,10 @@ typedef unsigned int node_id_type;
 enum ProblemType {
     INFEASIBLE,
     OPTIMAL,
-    UNBOUNDED
+    UNBOUNDED,
+	MAX_ITER_REACHED
 };
 
-int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double *cost, int max_iter);
+int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, int maxIter);
 
 #endif

ot/lp/EMD_wrapper.cpp

@@ -15,104 +15,92 @@
 #include "EMD.h"
 
 
-int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double *cost, int max_iter)  {
+int EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G,
+                double* alpha, double* beta, double *cost, int maxIter)  {
 // beware M and C anre strored in row major C style!!!
-  int n, m, i,cur;
-  double  max;
+    int n, m, i, cur;
 
     typedef FullBipartiteDigraph Digraph;
   DIGRAPH_TYPEDEFS(FullBipartiteDigraph);
 
   // Get the number of non zero coordinates for r and c
     n=0;
-    for (node_id_type i=0; i<n1; i++) {
+    for (int i=0; i<n1; i++) {
         double val=*(X+i);
         if (val>0) {
             n++;
-        }
+        }else if(val<0){
+			return INFEASIBLE;
+		}
     }
     m=0;
-    for (node_id_type i=0; i<n2; i++) {
+    for (int i=0; i<n2; i++) {
         double val=*(Y+i);
         if (val>0) {
             m++;
-        }
+        }else if(val<0){
+			return INFEASIBLE;
+		}
     }
 
-
     // Define the graph
 
     std::vector<int> indI(n), indJ(m);
     std::vector<double> weights1(n), weights2(m);
     Digraph di(n, m);
-    NetworkSimplexSimple<Digraph,double,double, node_id_type> net(di, true, n+m, n*m, max_iter);
+    NetworkSimplexSimple<Digraph,double,double, node_id_type> net(di, true, n+m, n*m, maxIter);
 
     // Set supply and demand, don't account for 0 values (faster)
 
-    max=0;
     cur=0;
-    for (node_id_type i=0; i<n1; i++) {
+    for (int i=0; i<n1; i++) {
         double val=*(X+i);
         if (val>0) {
-            weights1[ di.nodeFromId(cur) ] = val;
-            max+=val;
+            weights1[ cur ] = val;
             indI[cur++]=i;
         }
     }
 
     // Demand is actually negative supply...
 
-    max=0;
     cur=0;
-    for (node_id_type i=0; i<n2; i++) {
+    for (int i=0; i<n2; i++) {
         double val=*(Y+i);
         if (val>0) {
-            weights2[ di.nodeFromId(cur) ] = -val;
+            weights2[ cur ] = -val;
             indJ[cur++]=i;
-
-            max-=val;
         }
     }
 
 
     net.supplyMap(&weights1[0], n, &weights2[0], m);
 
     // Set the cost of each edge
-    max=0;
-    for (node_id_type i=0; i<n; i++) {
-        for (node_id_type j=0; j<m; j++) {
+    for (int i=0; i<n; i++) {
+        for (int j=0; j<m; j++) {
             double val=*(D+indI[i]*n2+indJ[j]);
             net.setCost(di.arcFromId(i*m+j), val);
-            if (val>max) {
-                max=val;
-            }
         }
     }
 
 
     // Solve the problem with the network simplex algorithm
 
     int ret=net.run();
-    if (ret!=(int)net.OPTIMAL) {
-        if (ret==(int)net.INFEASIBLE) {
-            std::cout << "Infeasible problem";
+    if (ret==(int)net.OPTIMAL || ret==(int)net.MAX_ITER_REACHED) {
+        *cost = 0;
+        Arc a; di.first(a);
+        for (; a != INVALID; di.next(a)) {
+            int i = di.source(a);
+            int j = di.target(a);
+            double flow = net.flow(a);
+            *cost += flow * (*(D+indI[i]*n2+indJ[j-n]));
+            *(G+indI[i]*n2+indJ[j-n]) = flow;
+            *(alpha + indI[i]) = -net.potential(i);
+            *(beta + indJ[j-n]) = net.potential(j);
         }
-        if (ret==(int)net.UNBOUNDED)
-        {
-            std::cout << "Unbounded problem";
-        }
-    } else
-    {
-        for (node_id_type i=0; i<n; i++)
-        {
-            for (node_id_type j=0; j<m; j++)
-            {
-                *(G+indI[i]*n2+indJ[j]) = net.flow(di.arcFromId(i*m+j));
-            }
-        };
-        *cost = net.totalCost();
-
-    };
+
+    }
 
 
     return ret;

ot/lp/__init__.py

@@ -7,14 +7,16 @@
 #
 # License: MIT License
 
+import multiprocessing
+
 import numpy as np
+
 # import compiled emd
-from .emd_wrap import emd_c, emd2_c
+from .emd_wrap import emd_c, check_result
 from ..utils import parmap
-import multiprocessing
 
 
-def emd(a, b, M, numItermax=100000):
+def emd(a, b, M, numItermax=100000, log=False):
     """Solves the Earth Movers distance problem and returns the OT matrix
 
 
@@ -42,11 +44,17 @@ def emd(a, b, M, numItermax=100000):
     numItermax : 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 matrix.
 
     Returns
     -------
     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
@@ -82,14 +90,24 @@ def emd(a, b, M, numItermax=100000):
 
     # if empty array given then use unifor distributions
     if len(a) == 0:
-        a = np.ones((M.shape[0], ), dtype=np.float64)/M.shape[0]
+        a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0]
     if len(b) == 0:
-        b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1]
-
-    return emd_c(a, b, M, numItermax)
-
-
-def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000):
+        b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
+
+    G, cost, u, v, result_code = emd_c(a, b, M, numItermax)
+    result_code_string = check_result(result_code)
+    if log:
+        log = {}
+        log['cost'] = cost
+        log['u'] = u
+        log['v'] = v
+        log['warning'] = result_code_string
+        log['result_code'] = result_code
+        return G, log
+    return G
+
+
+def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000, log=False, return_matrix=False):
     """Solves the Earth Movers distance problem and returns the loss
 
     .. math::
@@ -116,11 +134,19 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000):
     numItermax : 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
     -------
     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
@@ -156,17 +182,31 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(), numItermax=100000):
 
     # if empty array given then use unifor distributions
     if len(a) == 0:
-        a = np.ones((M.shape[0], ), dtype=np.float64)/M.shape[0]
+        a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0]
     if len(b) == 0:
-        b = np.ones((M.shape[1], ), dtype=np.float64)/M.shape[1]
+        b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1]
 
-    if len(b.shape) == 1:
-        return emd2_c(a, b, M, numItermax)
+    if log or return_matrix:
+        def f(b):
+            G, cost, u, v, resultCode = emd_c(a, b, M, numItermax)
+            result_code_string = check_result(resultCode)
+            log = {}
+            if return_matrix:
+                log['G'] = G
+            log['u'] = u
+            log['v'] = v
+            log['warning'] = result_code_string
+            log['result_code'] = resultCode
+            return [cost, log]
     else:
-        nb = b.shape[1]
-        # res = [emd2_c(a, b[:, i].copy(), M, numItermax) for i in range(nb)]
-
         def f(b):
-            return emd2_c(a, b, M, numItermax)
-        res = parmap(f, [b[:, i] for i in range(nb)], processes)
-        return np.array(res)
+            G, cost, u, v, result_code = emd_c(a, b, M, numItermax)
+            check_result(result_code)
+            return cost
+
+    if len(b.shape) == 1:
+        return f(b)
+    nb = b.shape[1]
+
+    res = parmap(f, [b[:, i] for i in range(nb)], processes)
+    return res