EMD 1d without doc

1 files changed, 35 insertions, 0 deletions

diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx
index 83ee6aa..a3d189d 100644
--- a/ot/lp/emd_wrap.pyx
+++ b/ot/lp/emd_wrap.pyx
@@ -93,3 +93,38 @@ def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mod
     cdef int 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
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def emd_1d_sorted(np.ndarray[double, ndim=1, mode="c"] u_weights,
+                  np.ndarray[double, ndim=1, mode="c"] v_weights,
+                  np.ndarray[double, ndim=2, mode="c"] M):
+    r"""
+    Roro's stuff
+    """
+    cdef double cost = 0.
+    cdef int n = u_weights.shape[0]
+    cdef int m = v_weights.shape[0]
+
+    cdef int i = 0
+    cdef double w_i = u_weights[0]
+    cdef int j = 0
+    cdef double w_j = v_weights[0]
+
+    cdef np.ndarray[double, ndim=2, mode="c"] G = np.zeros((n, m),
+                                                           dtype=np.float64)
+    while i < n and j < m:
+        if w_i < w_j or j == m - 1:
+            cost += M[i, j] * w_i
+            G[i, j] = w_i
+            i += 1
+            w_j -= w_i
+            w_i = u_weights[i]
+        else:
+            cost += M[i, j] * w_j
+            G[i, j] = w_j
+            j += 1
+            w_i -= w_j
+            w_j = v_weights[j]
+    return G, cost
\ No newline at end of file