diff options
author | Rémi Flamary <remi.flamary@gmail.com> | 2018-05-11 17:24:09 +0200 |
---|---|---|
committer | Rémi Flamary <remi.flamary@gmail.com> | 2018-05-11 17:24:09 +0200 |
commit | fdb2f3af19d04872bafa0d9ec5563732e1d6209b (patch) | |
tree | bc9d94d0d83126e68e633ce3030f801007426fe5 /test/test_ot.py | |
parent | 36f4f7ed2116841d7fe9514ee250bbf16e77b72d (diff) |
add test for barycenter
Diffstat (limited to 'test/test_ot.py')
-rw-r--r-- | test/test_ot.py | 35 |
1 files changed, 34 insertions, 1 deletions
diff --git a/test/test_ot.py b/test/test_ot.py index ea6d9dc..bf23e8c 100644 --- a/test/test_ot.py +++ b/test/test_ot.py @@ -10,7 +10,7 @@ import numpy as np import ot from ot.datasets import get_1D_gauss as gauss - +import pytest def test_doctest(): import doctest @@ -117,6 +117,39 @@ 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 |