diff options
author | Hicham Janati <hicham.janati@inria.fr> | 2019-06-12 15:50:25 +0200 |
---|---|---|
committer | Hicham Janati <hicham.janati@inria.fr> | 2019-06-12 15:50:25 +0200 |
commit | 28b549ef3ef93c01462cd811d6e55c36ae5a76a2 (patch) | |
tree | 75605c0ca1e501d1e5950ad09c6efd73f5e1c85a /test | |
parent | 3c53834d46f093f5770ec76748beb5667bebb6fa (diff) |
add test and example of UOT
Diffstat (limited to 'test')
-rw-r--r-- | test/test_unbalanced.py | 36 |
1 files changed, 36 insertions, 0 deletions
diff --git a/test/test_unbalanced.py b/test/test_unbalanced.py new file mode 100644 index 0000000..863b6f3 --- /dev/null +++ b/test/test_unbalanced.py @@ -0,0 +1,36 @@ +"""Tests for module Unbalanced OT with entropy regularization""" + +# Author: Hicham Janati <hicham.janati@inria.fr> +# +# License: MIT License + +import numpy as np +import ot + + +def test_unbalanced(): + # test generalized sinkhorn for unbalanced OT + n = 100 + rng = np.random.RandomState(42) + + x = rng.randn(n, 2) + a = ot.utils.unif(n) + 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, log=True) + + # 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) |