Merge pull request #47 from rflamary/bary

LP Wasserstein barycenter with scipy linear solver and/or cvxopt

2 files changed, 37 insertions, 1 deletions

diff --git a/test/test_gpu.py b/test/test_gpu.py
index 615c2a7..1e97c45 100644
--- a/test/test_gpu.py
+++ b/test/test_gpu.py
@@ -76,4 +76,4 @@ def test_gpu_sinkhorn_lpl1():
             time3 - time2))
         describe_res(G2)
 
-        np.testing.assert_allclose(G1, G2, rtol=1e-5, atol=1e-5)
+        np.testing.assert_allclose(G1, G2, rtol=1e-3, atol=1e-3)
diff --git a/test/test_ot.py b/test/test_ot.py
index ea6d9dc..cc25bf4 100644
--- a/test/test_ot.py
+++ b/test/test_ot.py
@@ -10,6 +10,7 @@ import numpy as np
 
 import ot
 from ot.datasets import get_1D_gauss as gauss
+import pytest
 
 
 def test_doctest():
@@ -117,6 +118,41 @@ def test_emd2_multi():
     np.testing.assert_allclose(emd1, emdn)
 
 
+def test_lp_barycenter():
+
+    a1 = np.array([1.0, 0, 0])[:, None]
+    a2 = np.array([0, 0, 1.0])[:, None]
+
+    A = np.hstack((a1, a2))
+    M = np.array([[0, 1.0, 4.0], [1.0, 0, 1.0], [4.0, 1.0, 0]])
+
+    # obvious barycenter between two diracs
+    bary0 = np.array([0, 1.0, 0])
+
+    bary = ot.lp.barycenter(A, M, [.5, .5])
+
+    np.testing.assert_allclose(bary, bary0, rtol=1e-5, atol=1e-7)
+    np.testing.assert_allclose(bary.sum(), 1)
+
+
+@pytest.mark.skipif(not ot.lp.cvx.cvxopt, reason="No cvxopt available")
+def test_lp_barycenter_cvxopt():
+
+    a1 = np.array([1.0, 0, 0])[:, None]
+    a2 = np.array([0, 0, 1.0])[:, None]
+
+    A = np.hstack((a1, a2))
+    M = np.array([[0, 1.0, 4.0], [1.0, 0, 1.0], [4.0, 1.0, 0]])
+
+    # obvious barycenter between two diracs
+    bary0 = np.array([0, 1.0, 0])
+
+    bary = ot.lp.barycenter(A, M, [.5, .5], solver=None)
+
+    np.testing.assert_allclose(bary, bary0, rtol=1e-5, atol=1e-7)
+    np.testing.assert_allclose(bary.sum(), 1)
+
+
 def test_warnings():
     n = 100  # nb bins
     m = 100  # nb bins