(nx)
| 56 | |
| 57 | |
| 58 | def test_wasserstein_1d(nx): |
| 59 | rng = np.random.RandomState(0) |
| 60 | |
| 61 | n = 100 |
| 62 | x = np.linspace(0, 5, n) |
| 63 | rho_u = np.abs(rng.randn(n)) |
| 64 | rho_u /= rho_u.sum() |
| 65 | rho_v = np.abs(rng.randn(n)) |
| 66 | rho_v /= rho_v.sum() |
| 67 | |
| 68 | xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v) |
| 69 | |
| 70 | # test 1 : wasserstein_1d should be close to scipy W_1 implementation |
| 71 | np.testing.assert_almost_equal( |
| 72 | wasserstein_1d(xb, xb, rho_ub, rho_vb, p=1), |
| 73 | wasserstein_distance(x, x, rho_u, rho_v), |
| 74 | ) |
| 75 | |
| 76 | # test 2 : wasserstein_1d should be close to one when only translating the support |
| 77 | np.testing.assert_almost_equal(wasserstein_1d(xb, xb + 1, p=2), 1.0) |
| 78 | |
| 79 | # test 3 : arrays test |
| 80 | X = np.stack((np.linspace(0, 5, n), np.linspace(0, 5, n) * 10), -1) |
| 81 | Xb = nx.from_numpy(X) |
| 82 | res = wasserstein_1d(Xb, Xb, rho_ub, rho_vb, p=2) |
| 83 | np.testing.assert_almost_equal(100 * res[0], res[1], decimal=4) |
| 84 | |
| 85 | |
| 86 | def test_wasserstein_1d_type_devices(nx): |
nothing calls this directly
no test coverage detected