[WIP] OpenMP EMD

.github/workflows/build_tests.yml

@@ -40,7 +40,7 @@ jobs:
         flake8 examples/ ot/ test/ --count --max-line-length=127 --statistics
     - name: Install POT
       run: |
-        pip install -e .
+        pip install --verbose --no-build-isolation -e .
     - name: Run tests
       run: |
         python -m pytest -v test/ ot/ --doctest-modules --ignore ot/gpu/ --cov=ot
@@ -71,37 +71,50 @@ jobs:
         pip install -U "sklearn"
     - name: Install POT
       run: |
-        pip install -e .
+        pip install --verbose -e .
     - name: Run tests
       run: |
         python -m pytest -v test/ ot/ --ignore ot/gpu/
 
 
   macos:
-     runs-on: macOS-latest
-     strategy:
-       max-parallel: 4
-       matrix:
-         python-version: [3.7]
-
-     steps:
-     - uses: actions/checkout@v1
-     - name: Set up Python ${{ matrix.python-version }}
-       uses: actions/setup-python@v1
-       with:
-         python-version: ${{ matrix.python-version }}
-     - name: Install dependencies
-       run: |
-         python -m pip install --upgrade pip
-         pip install -r requirements.txt
-         pip install pytest "pytest-cov<2.6"
-         pip install -U "sklearn"
-     - name: Install POT
-       run: |
-         pip install -e .
-     - name: Run tests
-       run: |
-         python -m pytest -v test/ ot/ --doctest-modules --ignore ot/gpu/ --cov=ot
+    runs-on: macOS-latest
+    strategy:
+      max-parallel: 4
+      matrix:
+        python-version: [3.7]
+
+    steps:
+    - uses: actions/checkout@v1
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        # install OpenMP (via homebrew)
+        /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"
+        brew install libomp
+        # install (necessary) recipes
+        python -m pip install -U pip setuptools
+        pip install -U setuptools 
+        pip install -r requirements.txt
+        pip install pytest "pytest-cov<2.6"
+        pip install -U "sklearn"
+    - name: Set environment variables
+      run: |
+        echo '::set-env name=CC::/usr/bin/g++' # hacky(!!) force c++ compiler
+        echo '::set-env name=CXX::/usr/bin/g++'
+        echo '::set-env name=CPPFLAGS::"-Xpreprocessor -fopenmp"'
+        echo '::set-env name=CFLAGS::"-I/usr/local/opt/libomp/include"'
+        echo '::set-env name=CXXFLAGS::"-I/usr/local/opt/libomp/include"'
+        echo '::set-env name=LDFLAGS::"-Wl,-rpath,/usr/local/opt/libomp/lib -L/usr/local/opt/libomp/lib -lomp"'
+    - name: Install POT
+      run: |
+        pip install --verbose -e .
+    - name: Run tests
+      run: |
+        python -m pytest -v test/ ot/ --doctest-modules --ignore ot/gpu/ --cov=ot
 
 
   windows:
@@ -125,7 +138,7 @@ jobs:
         pip install -U "sklearn"
     - name: Install POT
       run: |
-        pip install -e .
+        pip install --verbose -e .
     - name: Run tests
       run: |
         python -m pytest -v test/ ot/ --doctest-modules --ignore ot/gpu/ --cov=ot

ot/lp/EMD.h

@@ -16,22 +16,22 @@
 #ifndef EMD_H
 #define EMD_H
 
-#include <iostream>
 #include <vector>
 #include "network_simplex_simple.h"
 
 using namespace lemon;
-typedef unsigned int node_id_type;
+
+typedef int64_t arc_id_type; // handle (n1*n2+n1+n2) nodes (I64_MAX=3037000500^2)
+typedef double supply_type; // handle sum of supplies and demand (should be signed)
+typedef double cost_type; // handle number of arcs * maximum cost (should be signed)
 
 enum ProblemType {
     INFEASIBLE,
     OPTIMAL,
     UNBOUNDED,
-	MAX_ITER_REACHED
 };
 
-int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, int maxIter);
-
-
+int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G,
+             double* alpha, double* beta, double *cost, int maxIter);
 
-#endif
+#endif // EMD_H

ot/lp/EMD_wrapper.cpp

@@ -9,100 +9,97 @@
  *
  * Please give relevant credit to the original author (Nicolas Bonneel) if
  * you use this code for a publication.
- *
  */
 
+#include <iostream>
 #include "EMD.h"
 
 
 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* alpha, double* beta, double *cost, int maxIter) {
+    // Beware M and C are stored in row major C style!!!
+    int n; // number of nonzero coordinates (row)
+    int m; // number of nonzero coordinates (column)
+    int cur;
 
     typedef FullBipartiteDigraph Digraph;
     DIGRAPH_TYPEDEFS(FullBipartiteDigraph);
 
-    // Get the number of non zero coordinates for r and c
-    n=0;
+    // Count number of non zero coordinates for source (row) and target (column)
+    n = 0;
     for (int i=0; i<n1; i++) {
-        double val=*(X+i);
+        double val = *(X+i);
         if (val>0) {
             n++;
-        }else if(val<0){
+        } else if (val<0) {
 			return INFEASIBLE;
 		}
     }
-    m=0;
+    m = 0;
     for (int i=0; i<n2; i++) {
-        double val=*(Y+i);
-        if (val>0) {
+        double val = *(Y+i);
+        if (val > 0) {
             m++;
-        }else if(val<0){
+        } else if (val < 0) {
 			return INFEASIBLE;
 		}
     }
 
     // Define the graph
-
-    std::vector<int> indI(n), indJ(m);
-    std::vector<double> weights1(n), weights2(m);
+    std::vector<int> indI(n);
+    std::vector<int> indJ(m);
+    // Works faster with integer weights, however weights can be a density (aka
+    // double).
+    std::vector<supply_type> weights1(n);
+    std::vector<supply_type> weights2(m);
     Digraph di(n, m);
-    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)
+    NetworkSimplexSimple<Digraph, supply_type, cost_type, arc_id_type> net(di, true, n+m, n*m, maxIter);
 
+    // Set supply and demand; don't account for 0 values (faster)
     cur=0;
     for (int i=0; i<n1; i++) {
-        double val=*(X+i);
+        double val = *(X+i);
         if (val>0) {
-            weights1[ cur ] = val;
+            weights1[cur] = val;
             indI[cur++]=i;
         }
     }
-
-    // Demand is actually negative supply...
-
+    // Demand is actually negative supply
     cur=0;
     for (int i=0; i<n2; i++) {
-        double val=*(Y+i);
+        double val = *(Y+i);
         if (val>0) {
-            weights2[ cur ] = -val;
+            weights2[cur] = -val;
             indJ[cur++]=i;
         }
     }
 
-
     net.supplyMap(&weights1[0], n, &weights2[0], m);
 
     // Set the cost of each edge
     for (int i=0; i<n; i++) {
         for (int j=0; j<m; j++) {
-            double val=*(D+indI[i]*n2+indJ[j]);
+            cost_type val = *(D+indI[i]*n2+indJ[j]);
             net.setCost(di.arcFromId(i*m+j), val);
         }
     }
 
-
-    // Solve the problem with the network simplex algorithm
-
-    int ret=net.run();
-    if (ret==(int)net.OPTIMAL || ret==(int)net.MAX_ITER_REACHED) {
-        *cost = 0;
-        Arc a; di.first(a);
+    // Solve the problem using the network simplex algorithm
+    int ret = net.run();
+    *cost = net.totalCost();
+    if (ret == (int)net.OPTIMAL) {
+        // *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]));
+            // *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);
         }
-
     }
-
-
     return ret;
 }
-

ot/lp/__init__.py

@@ -371,10 +371,6 @@ def emd2(a, b, M, processes=multiprocessing.cpu_count(),
     b = np.asarray(b, dtype=np.float64)
     M = np.asarray(M, dtype=np.float64)
 
-    # problem with pikling Forks
-    if sys.platform.endswith('win32'):
-        processes = 1
-
     # if empty array given then use uniform distributions
     if len(a) == 0:
         a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0]
@@ -421,16 +417,9 @@ def f(b):
             check_result(result_code)
             return cost
 
-    if len(b.shape) == 1:
+    if b.shape[0] == 1:
         return f(b)
-    nb = b.shape[1]
-
-    if processes > 1:
-        res = parmap(f, [b[:, i] for i in range(nb)], processes)
-    else:
-        res = list(map(f, [b[:, i].copy() for i in range(nb)]))
-
-    return res
+    return list(map(f, [b[:, i].copy() for i in range(b.shape[1])]))
 
 
 def free_support_barycenter(measures_locations, measures_weights, X_init, b=None, weights=None, numItermax=100,

ot/lp/emd_wrap.pyx

@@ -19,8 +19,9 @@ import warnings
 
 
 cdef extern from "EMD.h":
-    int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, int maxIter) nogil
-    cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED
+    int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G,
+                 double* alpha, double* beta, double *cost, int maxIter) nogil
+    cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED
 
 
 def check_result(result_code):
@@ -31,18 +32,18 @@ def check_result(result_code):
         message = "Problem infeasible. Check that a and b are in the simplex"
     elif result_code == UNBOUNDED:
         message = "Problem unbounded"
-    elif result_code == MAX_ITER_REACHED:
-        message = "numItermax reached before optimality. Try to increase numItermax."
     warnings.warn(message)
     return message
 
 @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 max_iter):
+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 max_iter):
     """
-        Solves the Earth Movers distance problem and returns the optimal transport matrix
-
-        gamm=emd(a,b,M)
+    Solves the Earth Movers distance problem and returns the optimal
+    transport matrix :math:`gamma=emd(a,b,M)`.
 
     .. math::
         \gamma = arg\min_\gamma <\gamma,M>_F
@@ -82,7 +83,6 @@ def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mod
     -------
     gamma: (ns x nt) numpy.ndarray
         Optimal transportation matrix for the given parameters
-
     """
     cdef int n1= M.shape[0]
     cdef int n2= M.shape[1]
@@ -109,10 +109,12 @@ def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mod
     # init OT matrix
     G=np.zeros([n1, n2])
 
-    # calling the function
+    # solve lp, allowing multi-thread ops
     with nogil:
-        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, max_iter)
-
+        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, max_iter)
     return G, cost, alpha, beta, result_code
 
 
@@ -126,7 +128,7 @@ def emd_1d_sorted(np.ndarray[double, ndim=1, mode="c"] u_weights,
                   double p=1.):
     r"""
     Solves the Earth Movers distance problem between sorted 1d measures and
-    returns the OT matrix and the associated cost
+    returns the OT matrix and the associated cost.
 
     Parameters
     ----------
@@ -144,7 +146,7 @@ def emd_1d_sorted(np.ndarray[double, ndim=1, mode="c"] u_weights,
         `'sqeuclidean'`, `'minkowski'`, `'cityblock'`,  or `'euclidean'` metrics
         are used.
     p: float, optional (default=1.0)
-         The p-norm to apply for if metric='minkowski'
+        The p-norm to apply for if metric='minkowski'
 
     Returns
     -------
@@ -153,8 +155,8 @@ def emd_1d_sorted(np.ndarray[double, ndim=1, mode="c"] u_weights,
     indices: (n, 2) ndarray, int64
         Indices of the values stored in gamma for the Optimal transportation
         matrix
-    cost
-        cost associated to the optimal transportation
+    cost: float
+        Cost associated to the optimal transportation
     """
     cdef double cost = 0.
     cdef Py_ssize_t n = u_weights.shape[0]
@@ -170,7 +172,7 @@ def emd_1d_sorted(np.ndarray[double, ndim=1, mode="c"] u_weights,
     cdef np.ndarray[double, ndim=1, mode="c"] G = np.zeros((n + m - 1, ),
                                                            dtype=np.float64)
     cdef np.ndarray[long, ndim=2, mode="c"] indices = np.zeros((n + m - 1, 2),
-                                                              dtype=np.int)
+                                                               dtype=np.int)
     cdef Py_ssize_t cur_idx = 0
     while True:
         if metric == 'sqeuclidean':