Merge branch 'master' into stochastic_OT

2 files changed, 94 insertions, 0 deletions

diff --git a/test/test_ot.py b/test/test_ot.py
index 399e549..45e777a 100644
--- a/test/test_ot.py
+++ b/test/test_ot.py
@@ -135,6 +135,21 @@ def test_lp_barycenter():
     np.testing.assert_allclose(bary.sum(), 1)
 
 
+def test_free_support_barycenter():
+
+    measures_locations = [np.array([-1.]).reshape((1, 1)), np.array([1.]).reshape((1, 1))]
+    measures_weights = [np.array([1.]), np.array([1.])]
+
+    X_init = np.array([-12.]).reshape((1, 1))
+
+    # obvious barycenter location between two diracs
+    bar_locations = np.array([0.]).reshape((1, 1))
+
+    X = ot.lp.free_support_barycenter(measures_locations, measures_weights, X_init)
+
+    np.testing.assert_allclose(X, bar_locations, rtol=1e-5, atol=1e-7)
+
+
 @pytest.mark.skipif(not ot.lp.cvx.cvxopt, reason="No cvxopt available")
 def test_lp_barycenter_cvxopt():
 
diff --git a/test/test_smooth.py b/test/test_smooth.py
new file mode 100644
index 0000000..2afa4f8
--- /dev/null
+++ b/test/test_smooth.py
@@ -0,0 +1,79 @@
+"""Tests for ot.smooth model """
+
+# Author: Remi Flamary <remi.flamary@unice.fr>
+#
+# License: MIT License
+
+import numpy as np
+import ot
+import pytest
+
+
+def test_smooth_ot_dual():
+
+    # get data
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 2)
+    u = ot.utils.unif(n)
+
+    M = ot.dist(x, x)
+
+    with pytest.raises(NotImplementedError):
+        Gl2, log = ot.smooth.smooth_ot_dual(u, u, M, 1, reg_type='none')
+
+    Gl2, log = ot.smooth.smooth_ot_dual(u, u, M, 1, reg_type='l2', log=True, stopThr=1e-10)
+
+    # check constratints
+    np.testing.assert_allclose(
+        u, Gl2.sum(1), atol=1e-05)  # cf convergence sinkhorn
+    np.testing.assert_allclose(
+        u, Gl2.sum(0), atol=1e-05)  # cf convergence sinkhorn
+
+    # kl regyularisation
+    G = ot.smooth.smooth_ot_dual(u, u, M, 1, reg_type='kl', stopThr=1e-10)
+
+    # check constratints
+    np.testing.assert_allclose(
+        u, G.sum(1), atol=1e-05)  # cf convergence sinkhorn
+    np.testing.assert_allclose(
+        u, G.sum(0), atol=1e-05)  # cf convergence sinkhorn
+
+    G2 = ot.sinkhorn(u, u, M, 1, stopThr=1e-10)
+    np.testing.assert_allclose(G, G2, atol=1e-05)
+
+
+def test_smooth_ot_semi_dual():
+
+    # get data
+    n = 100
+    rng = np.random.RandomState(0)
+
+    x = rng.randn(n, 2)
+    u = ot.utils.unif(n)
+
+    M = ot.dist(x, x)
+
+    with pytest.raises(NotImplementedError):
+        Gl2, log = ot.smooth.smooth_ot_semi_dual(u, u, M, 1, reg_type='none')
+
+    Gl2, log = ot.smooth.smooth_ot_semi_dual(u, u, M, 1, reg_type='l2', log=True, stopThr=1e-10)
+
+    # check constratints
+    np.testing.assert_allclose(
+        u, Gl2.sum(1), atol=1e-05)  # cf convergence sinkhorn
+    np.testing.assert_allclose(
+        u, Gl2.sum(0), atol=1e-05)  # cf convergence sinkhorn
+
+    # kl regyularisation
+    G = ot.smooth.smooth_ot_semi_dual(u, u, M, 1, reg_type='kl', stopThr=1e-10)
+
+    # check constratints
+    np.testing.assert_allclose(
+        u, G.sum(1), atol=1e-05)  # cf convergence sinkhorn
+    np.testing.assert_allclose(
+        u, G.sum(0), atol=1e-05)  # cf convergence sinkhorn
+
+    G2 = ot.sinkhorn(u, u, M, 1, stopThr=1e-10)
+    np.testing.assert_allclose(G, G2, atol=1e-05)