diff options
Diffstat (limited to 'test/test_ot.py')
| -rw-r--r-- | test/test_ot.py | 67 |
1 files changed, 60 insertions, 7 deletions
diff --git a/test/test_ot.py b/test/test_ot.py index dacae0a..b7306f6 100644 --- a/test/test_ot.py +++ b/test/test_ot.py @@ -7,11 +7,27 @@ import warnings import numpy as np +import pytest from scipy.stats import wasserstein_distance import ot from ot.datasets import make_1D_gauss as gauss -import pytest + + +def test_emd_dimension_mismatch(): + # test emd and emd2 for dimension mismatch + n_samples = 100 + n_features = 2 + rng = np.random.RandomState(0) + + x = rng.randn(n_samples, n_features) + a = ot.utils.unif(n_samples + 1) + + M = ot.dist(x, x) + + np.testing.assert_raises(AssertionError, ot.emd, a, a, M) + + np.testing.assert_raises(AssertionError, ot.emd2, a, a, M) def test_emd_emd2(): @@ -59,12 +75,12 @@ def test_emd_1d_emd2_1d(): np.testing.assert_allclose(wass, wass1d_emd2) # check loss is similar to scipy's implementation for Euclidean metric - wass_sp = wasserstein_distance(u.reshape((-1, )), v.reshape((-1, ))) + wass_sp = wasserstein_distance(u.reshape((-1,)), v.reshape((-1,))) np.testing.assert_allclose(wass_sp, wass1d_euc) # check constraints - np.testing.assert_allclose(np.ones((n, )) / n, G.sum(1)) - np.testing.assert_allclose(np.ones((m, )) / m, G.sum(0)) + np.testing.assert_allclose(np.ones((n,)) / n, G.sum(1)) + np.testing.assert_allclose(np.ones((m,)) / m, G.sum(0)) # check G is similar np.testing.assert_allclose(G, G_1d) @@ -76,6 +92,42 @@ def test_emd_1d_emd2_1d(): ot.emd_1d(u, v, [], []) +def test_emd_1d_emd2_1d_with_weights(): + # test emd1d gives similar results as emd + n = 20 + m = 30 + rng = np.random.RandomState(0) + u = rng.randn(n, 1) + v = rng.randn(m, 1) + + w_u = rng.uniform(0., 1., n) + w_u = w_u / w_u.sum() + + w_v = rng.uniform(0., 1., m) + w_v = w_v / w_v.sum() + + M = ot.dist(u, v, metric='sqeuclidean') + + G, log = ot.emd(w_u, w_v, M, log=True) + wass = log["cost"] + G_1d, log = ot.emd_1d(u, v, w_u, w_v, metric='sqeuclidean', log=True) + wass1d = log["cost"] + wass1d_emd2 = ot.emd2_1d(u, v, w_u, w_v, metric='sqeuclidean', log=False) + wass1d_euc = ot.emd2_1d(u, v, w_u, w_v, metric='euclidean', log=False) + + # check loss is similar + np.testing.assert_allclose(wass, wass1d) + np.testing.assert_allclose(wass, wass1d_emd2) + + # check loss is similar to scipy's implementation for Euclidean metric + wass_sp = wasserstein_distance(u.reshape((-1,)), v.reshape((-1,)), w_u, w_v) + np.testing.assert_allclose(wass_sp, wass1d_euc) + + # check constraints + np.testing.assert_allclose(w_u, G.sum(1)) + np.testing.assert_allclose(w_v, G.sum(0)) + + def test_wass_1d(): # test emd1d gives similar results as emd n = 20 @@ -168,7 +220,6 @@ def test_emd2_multi(): def test_lp_barycenter(): - a1 = np.array([1.0, 0, 0])[:, None] a2 = np.array([0, 0, 1.0])[:, None] @@ -185,7 +236,6 @@ def test_lp_barycenter(): 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.])] @@ -201,7 +251,6 @@ def test_free_support_barycenter(): @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] @@ -295,6 +344,10 @@ def test_dual_variables(): np.testing.assert_almost_equal(cost1, log['cost']) check_duality_gap(a, b, M, G, log['u'], log['v'], log['cost']) + constraint_violation = log['u'][:, None] + log['v'][None, :] - M + + assert constraint_violation.max() < 1e-8 + def check_duality_gap(a, b, M, G, u, v, cost): cost_dual = np.vdot(a, u) + np.vdot(b, v) |