diff options
author | Hicham Janati <hicham.janati@inria.fr> | 2019-06-18 16:40:06 +0200 |
---|---|---|
committer | Hicham Janati <hicham.janati@inria.fr> | 2019-06-18 16:40:06 +0200 |
commit | 897982718a5fd81a9a591d80a7d50839399fc088 (patch) | |
tree | 48189493c09cda25ee19dfd0b7ef59c2f6819ba7 /test | |
parent | 50bc90058940645a13e2f3e41129bdc97161dc63 (diff) |
fix func names + add more tests
Diffstat (limited to 'test')
-rw-r--r-- | test/test_unbalanced.py | 79 |
1 files changed, 77 insertions, 2 deletions
diff --git a/test/test_unbalanced.py b/test/test_unbalanced.py index b39e457..1395fe1 100644 --- a/test/test_unbalanced.py +++ b/test/test_unbalanced.py @@ -29,7 +29,8 @@ def test_unbalanced_convergence(method): 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 @@ -40,6 +41,44 @@ def test_unbalanced_convergence(method): 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 @@ -59,7 +98,6 @@ def test_unbalanced_barycenter(): 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 @@ -69,3 +107,40 @@ def test_unbalanced_barycenter(): 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) |