diff options
author | Romain Tavenard <romain.tavenard@univ-rennes2.fr> | 2019-06-27 10:54:13 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2019-06-27 10:54:13 +0200 |
commit | bbc56e74bf119b8810c0de7b446bb01b30efc3c2 (patch) | |
tree | 223e26e51da0edd3057fa16820ce7f1882a94f59 /test | |
parent | 0d333e004636f5d25edea6bb195e8e4d9a95ba98 (diff) | |
parent | 2364d56aad650d501753cc93a69ea1b8ddf28b0a (diff) |
Merge branch 'master' into master
Diffstat (limited to 'test')
-rw-r--r-- | test/test_unbalanced.py | 146 |
1 files changed, 146 insertions, 0 deletions
diff --git a/test/test_unbalanced.py b/test/test_unbalanced.py new file mode 100644 index 0000000..1395fe1 --- /dev/null +++ b/test/test_unbalanced.py @@ -0,0 +1,146 @@ +"""Tests for module Unbalanced OT with entropy regularization""" + +# Author: Hicham Janati <hicham.janati@inria.fr> +# +# License: MIT License + +import numpy as np +import ot +import pytest + + +@pytest.mark.parametrize("method", ["sinkhorn"]) +def test_unbalanced_convergence(method): + # test generalized sinkhorn for unbalanced OT + n = 100 + rng = np.random.RandomState(42) + + x = rng.randn(n, 2) + a = ot.utils.unif(n) + + # make dists unbalanced + b = ot.utils.unif(n) * 1.5 + + M = ot.dist(x, x) + epsilon = 1. + alpha = 1. + K = np.exp(- M / epsilon) + + G, log = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg=epsilon, alpha=alpha, + stopThr=1e-10, method=method, + log=True) + loss = ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha, + method=method) + # check fixed point equations + fi = alpha / (alpha + epsilon) + v_final = (b / 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) + + # check if sinkhorn_unbalanced2 returns the correct loss + np.testing.assert_allclose((G * M).sum(), loss, atol=1e-5) + + +@pytest.mark.parametrize("method", ["sinkhorn"]) +def test_unbalanced_multiple_inputs(method): + # test generalized sinkhorn for unbalanced OT + n = 100 + rng = np.random.RandomState(42) + + x = rng.randn(n, 2) + a = ot.utils.unif(n) + + # make dists unbalanced + b = rng.rand(n, 2) + + M = ot.dist(x, x) + epsilon = 1. + alpha = 1. + K = np.exp(- M / epsilon) + + loss, log = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg=epsilon, + alpha=alpha, + stopThr=1e-10, method=method, + log=True) + # check fixed point equations + fi = alpha / (alpha + epsilon) + v_final = (b / K.T.dot(log["u"])) ** fi + + u_final = (a[:, None] / 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) + + assert len(loss) == b.shape[1] + + +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) + + +def test_implemented_methods(): + IMPLEMENTED_METHODS = ['sinkhorn'] + TO_BE_IMPLEMENTED_METHODS = ['sinkhorn_stabilized', + 'sinkhorn_epsilon_scaling'] + NOT_VALID_TOKENS = ['foo'] + # test generalized sinkhorn for unbalanced OT barycenter + n = 3 + rng = np.random.RandomState(42) + + x = rng.randn(n, 2) + a = ot.utils.unif(n) + + # make dists unbalanced + b = ot.utils.unif(n) * 1.5 + + M = ot.dist(x, x) + epsilon = 1. + alpha = 1. + for method in IMPLEMENTED_METHODS: + ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, + method=method) + ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha, + method=method) + with pytest.warns(UserWarning, match='not implemented'): + for method in set(TO_BE_IMPLEMENTED_METHODS): + ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, + method=method) + ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha, + method=method) + with pytest.raises(ValueError): + for method in set(NOT_VALID_TOKENS): + ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, + method=method) + ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha, + method=method) |