| 208 | |
| 209 | |
| 210 | def test_coot_warmstart(nx): |
| 211 | n_samples = 80 # nb samples |
| 212 | |
| 213 | mu_s = np.array([2, 3]) |
| 214 | cov_s = np.array([[1, 0], [0, 1]]) |
| 215 | |
| 216 | xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s, random_state=125) |
| 217 | xt = xs[::-1].copy() |
| 218 | xs_nx = nx.from_numpy(xs) |
| 219 | xt_nx = nx.from_numpy(xt) |
| 220 | |
| 221 | # initialize warmstart |
| 222 | rng = np.random.RandomState(42) |
| 223 | init_pi_sample = rng.rand(n_samples, n_samples) |
| 224 | init_pi_sample = init_pi_sample / np.sum(init_pi_sample) |
| 225 | init_pi_sample_nx = nx.from_numpy(init_pi_sample) |
| 226 | |
| 227 | init_pi_feature = rng.rand(2, 2) |
| 228 | init_pi_feature /= init_pi_feature / np.sum(init_pi_feature) |
| 229 | init_pi_feature_nx = nx.from_numpy(init_pi_feature) |
| 230 | |
| 231 | init_duals_sample = (rng.random(n_samples) * 2 - 1, rng.random(n_samples) * 2 - 1) |
| 232 | init_duals_sample_nx = ( |
| 233 | nx.from_numpy(init_duals_sample[0]), |
| 234 | nx.from_numpy(init_duals_sample[1]), |
| 235 | ) |
| 236 | |
| 237 | init_duals_feature = (rng.random(2) * 2 - 1, rng.random(2) * 2 - 1) |
| 238 | init_duals_feature_nx = ( |
| 239 | nx.from_numpy(init_duals_feature[0]), |
| 240 | nx.from_numpy(init_duals_feature[1]), |
| 241 | ) |
| 242 | |
| 243 | warmstart = { |
| 244 | "pi_sample": init_pi_sample, |
| 245 | "pi_feature": init_pi_feature, |
| 246 | "duals_sample": init_duals_sample, |
| 247 | "duals_feature": init_duals_feature, |
| 248 | } |
| 249 | |
| 250 | warmstart_nx = { |
| 251 | "pi_sample": init_pi_sample_nx, |
| 252 | "pi_feature": init_pi_feature_nx, |
| 253 | "duals_sample": init_duals_sample_nx, |
| 254 | "duals_feature": init_duals_feature_nx, |
| 255 | } |
| 256 | |
| 257 | # test couplings |
| 258 | pi_sample, pi_feature = coot(X=xs, Y=xt, warmstart=warmstart) |
| 259 | pi_sample_nx, pi_feature_nx = coot(X=xs_nx, Y=xt_nx, warmstart=warmstart_nx) |
| 260 | pi_sample_nx = nx.to_numpy(pi_sample_nx) |
| 261 | pi_feature_nx = nx.to_numpy(pi_feature_nx) |
| 262 | |
| 263 | anti_id_sample = np.flipud(np.eye(n_samples, n_samples)) / n_samples |
| 264 | id_feature = np.eye(2, 2) / 2 |
| 265 | |
| 266 | np.testing.assert_allclose(pi_sample, anti_id_sample, atol=1e-04) |
| 267 | np.testing.assert_allclose(pi_sample_nx, anti_id_sample, atol=1e-04) |