| 107 | |
| 108 | |
| 109 | def test_coot_with_linear_terms(nx): |
| 110 | n_samples = 60 # nb samples |
| 111 | |
| 112 | mu_s = np.array([0, 0]) |
| 113 | cov_s = np.array([[1, 0], [0, 1]]) |
| 114 | |
| 115 | xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s, random_state=4) |
| 116 | xt = xs[::-1].copy() |
| 117 | xs_nx = nx.from_numpy(xs) |
| 118 | xt_nx = nx.from_numpy(xt) |
| 119 | |
| 120 | M_samp = np.ones((n_samples, n_samples)) |
| 121 | np.fill_diagonal(np.fliplr(M_samp), 0) |
| 122 | M_feat = np.ones((2, 2)) |
| 123 | np.fill_diagonal(M_feat, 0) |
| 124 | M_samp_nx, M_feat_nx = nx.from_numpy(M_samp), nx.from_numpy(M_feat) |
| 125 | |
| 126 | alpha = (1, 2) |
| 127 | |
| 128 | # test couplings |
| 129 | anti_id_sample = np.flipud(np.eye(n_samples, n_samples)) / n_samples |
| 130 | id_feature = np.eye(2, 2) / 2 |
| 131 | |
| 132 | pi_sample, pi_feature = coot(X=xs, Y=xt, alpha=alpha, M_samp=M_samp, M_feat=M_feat) |
| 133 | pi_sample_nx, pi_feature_nx = coot( |
| 134 | X=xs_nx, Y=xt_nx, alpha=alpha, M_samp=M_samp_nx, M_feat=M_feat_nx |
| 135 | ) |
| 136 | pi_sample_nx = nx.to_numpy(pi_sample_nx) |
| 137 | pi_feature_nx = nx.to_numpy(pi_feature_nx) |
| 138 | |
| 139 | np.testing.assert_allclose(pi_sample, anti_id_sample, atol=1e-04) |
| 140 | np.testing.assert_allclose(pi_sample_nx, anti_id_sample, atol=1e-04) |
| 141 | np.testing.assert_allclose(pi_feature, id_feature, atol=1e-04) |
| 142 | np.testing.assert_allclose(pi_feature_nx, id_feature, atol=1e-04) |
| 143 | |
| 144 | # test marginal distributions |
| 145 | px_s, px_f = ot.unif(n_samples), ot.unif(2) |
| 146 | py_s, py_f = ot.unif(n_samples), ot.unif(2) |
| 147 | |
| 148 | np.testing.assert_allclose(px_s, pi_sample_nx.sum(0), atol=1e-04) |
| 149 | np.testing.assert_allclose(py_s, pi_sample_nx.sum(1), atol=1e-04) |
| 150 | np.testing.assert_allclose(px_f, pi_feature_nx.sum(0), atol=1e-04) |
| 151 | np.testing.assert_allclose(py_f, pi_feature_nx.sum(1), atol=1e-04) |
| 152 | |
| 153 | np.testing.assert_allclose(px_s, pi_sample.sum(0), atol=1e-04) |
| 154 | np.testing.assert_allclose(py_s, pi_sample.sum(1), atol=1e-04) |
| 155 | np.testing.assert_allclose(px_f, pi_feature.sum(0), atol=1e-04) |
| 156 | np.testing.assert_allclose(py_f, pi_feature.sum(1), atol=1e-04) |
| 157 | |
| 158 | # test COOT distance |
| 159 | |
| 160 | coot_np = coot2(X=xs, Y=xt, alpha=alpha, M_samp=M_samp, M_feat=M_feat) |
| 161 | coot_nx = nx.to_numpy( |
| 162 | coot2(X=xs_nx, Y=xt_nx, alpha=alpha, M_samp=M_samp_nx, M_feat=M_feat_nx) |
| 163 | ) |
| 164 | np.testing.assert_allclose(coot_np, 0, atol=1e-08) |
| 165 | np.testing.assert_allclose(coot_nx, 0, atol=1e-08) |
| 166 | |