add unbalanced barycenters

1 files changed, 30 insertions, 0 deletions

diff --git a/test/test_unbalanced.py b/test/test_unbalanced.py
index b4fa355..b39e457 100644
--- a/test/test_unbalanced.py
+++ b/test/test_unbalanced.py
@@ -39,3 +39,33 @@ def test_unbalanced_convergence(method):
         u_final, log["u"], atol=1e-05)
     np.testing.assert_allclose(
         v_final, log["v"], atol=1e-05)
+
+
+def test_unbalanced_barycenter():
+    # test generalized sinkhorn for unbalanced OT barycenter
+    n = 100
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n, 2)
+    A = rng.rand(n, 2)
+
+    # make dists unbalanced
+    A = A * np.array([1, 2])[None, :]
+    M = ot.dist(x, x)
+    epsilon = 1.
+    alpha = 1.
+    K = np.exp(- M / epsilon)
+
+    q, log = ot.unbalanced.barycenter_unbalanced(A, M, reg=epsilon, alpha=alpha,
+                                                 stopThr=1e-10,
+                                                 log=True)
+
+    # check fixed point equations
+    fi = alpha / (alpha + epsilon)
+    v_final = (q[:, None] / K.T.dot(log["u"])) ** fi
+    u_final = (A / K.dot(log["v"])) ** fi
+
+    np.testing.assert_allclose(
+        u_final, log["u"], atol=1e-05)
+    np.testing.assert_allclose(
+        v_final, log["v"], atol=1e-05)