From 3e7eb430ddf59cc2a6bbd2ada29ac22c9d31cd05 Mon Sep 17 00:00:00 2001
From: Nicolas Courty <ncourty@irisa.fr>
Date: Fri, 7 Sep 2018 15:13:44 +0200
Subject: test wasserstein barycenter

---
 test/test_bregman.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

(limited to 'test/test_bregman.py')

diff --git a/test/test_bregman.py b/test/test_bregman.py
index c8e9179..01ec655 100644
--- a/test/test_bregman.py
+++ b/test/test_bregman.py
@@ -105,6 +105,30 @@ def test_bary():
     ot.bregman.barycenter(A, M, reg, log=True, verbose=True)
 
 
+def test_wassersteinbary():
+
+    size = 100  # size of a square image
+    a1 = np.random.randn(size, size)
+    a1 += a1.min()
+    a1 = a1 / np.sum(a1)
+    a2 = np.random.randn(size, size)
+    a2 += a2.min()
+    a2 = a2 / np.sum(a2)
+    # creating matrix A containing all distributions
+    A = np.zeros((2, 100, 100))
+    A[0, :, :] = a1
+    A[1, :, :] = a2
+
+    # wasserstein
+    reg = 1e-3
+    bary_wass = ot.bregman.convolutional_barycenter2d(A, reg)
+
+    np.testing.assert_allclose(1, np.sum(bary_wass))
+
+    # help in checking if log and verbose do not bug the function
+    ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
+
+
 def test_unmix():
 
     n_bins = 50  # nb bins
-- 
cgit v1.2.3