| 89 | |
| 90 | |
| 91 | def test_partial_wasserstein(nx): |
| 92 | n_samples = 20 # nb samples (gaussian) |
| 93 | n_noise = 20 # nb of samples (noise) |
| 94 | |
| 95 | mu = np.array([0, 0]) |
| 96 | cov = np.array([[1, 0], [0, 2]]) |
| 97 | |
| 98 | rng = np.random.RandomState(42) |
| 99 | xs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov, random_state=rng) |
| 100 | xs = np.append(xs, (rng.rand(n_noise, 2) + 1) * 4).reshape((-1, 2)) |
| 101 | xt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov, random_state=rng) |
| 102 | xt = np.append(xt, (rng.rand(n_noise, 2) + 1) * -3).reshape((-1, 2)) |
| 103 | |
| 104 | M = ot.dist(xs, xt) |
| 105 | |
| 106 | p = ot.unif(n_samples + n_noise) |
| 107 | q = ot.unif(n_samples + n_noise) |
| 108 | |
| 109 | m = 0.5 |
| 110 | |
| 111 | p, q, M = nx.from_numpy(p, q, M) |
| 112 | |
| 113 | w0, log0 = ot.partial.partial_wasserstein(p, q, M, m=m, log=True) |
| 114 | w, log = ot.partial.entropic_partial_wasserstein( |
| 115 | p, q, M, reg=1, m=m, log=True, verbose=True |
| 116 | ) |
| 117 | |
| 118 | # check constraints |
| 119 | np.testing.assert_equal(to_numpy(nx.sum(w0, axis=1) - p) <= 1e-5, [True] * len(p)) |
| 120 | np.testing.assert_equal(to_numpy(nx.sum(w0, axis=0) - q) <= 1e-5, [True] * len(q)) |
| 121 | np.testing.assert_equal(to_numpy(nx.sum(w0, axis=1) - p) <= 1e-5, [True] * len(p)) |
| 122 | np.testing.assert_equal(to_numpy(nx.sum(w0, axis=0) - q) <= 1e-5, [True] * len(q)) |
| 123 | |
| 124 | # check transported mass |
| 125 | np.testing.assert_allclose(np.sum(to_numpy(w0)), m, atol=1e-04) |
| 126 | np.testing.assert_allclose(np.sum(to_numpy(w)), m, atol=1e-04) |
| 127 | |
| 128 | w0, log0 = ot.partial.partial_wasserstein2(p, q, M, m=m, log=True) |
| 129 | w0_val = ot.partial.partial_wasserstein2(p, q, M, m=m, log=False) |
| 130 | |
| 131 | G = log0["T"] |
| 132 | |
| 133 | np.testing.assert_allclose(w0, w0_val, atol=1e-1, rtol=1e-1) |
| 134 | |
| 135 | # check constraints |
| 136 | np.testing.assert_equal(to_numpy(nx.sum(G, axis=1) - p) <= 1e-5, [True] * len(p)) |
| 137 | np.testing.assert_equal(to_numpy(nx.sum(G, axis=0) - q) <= 1e-5, [True] * len(q)) |
| 138 | np.testing.assert_allclose(np.sum(to_numpy(G)), m, atol=1e-04) |
| 139 | |
| 140 | empty_array = nx.zeros(0, type_as=M) |
| 141 | w = ot.partial.partial_wasserstein(empty_array, empty_array, M=M, m=None) |
| 142 | |
| 143 | # check constraints |
| 144 | np.testing.assert_equal(to_numpy(nx.sum(w, axis=1) - p) <= 1e-5, [True] * len(p)) |
| 145 | np.testing.assert_equal(to_numpy(nx.sum(w, axis=0) - q) <= 1e-5, [True] * len(q)) |
| 146 | np.testing.assert_equal(to_numpy(nx.sum(w, axis=1) - p) <= 1e-5, [True] * len(p)) |
| 147 | np.testing.assert_equal(to_numpy(nx.sum(w, axis=0) - q) <= 1e-5, [True] * len(q)) |
| 148 | |