Check that solve_gromov_batch gives the same results as solve for each instance in the batch.
()
| 17 | |
| 18 | |
| 19 | def test_solve_gromov_batch(): |
| 20 | """Check that solve_gromov_batch gives the same results as solve for each instance in the batch.""" |
| 21 | b = 2 |
| 22 | n = 8 |
| 23 | d = 2 |
| 24 | reg = 0.01 |
| 25 | max_iter = 1000 |
| 26 | max_iter_inner = 10000 |
| 27 | tol = 1e-5 |
| 28 | tol_inner = 1e-5 |
| 29 | alpha = 0.5 |
| 30 | |
| 31 | rng = np.random.RandomState(0) |
| 32 | |
| 33 | X1 = rng.randn(b, n, d).astype("float32") |
| 34 | C1 = rng.randn(b, n, n).astype("float32") |
| 35 | |
| 36 | permutation = np.random.permutation(n) |
| 37 | X2 = X1[:, permutation, :] + 0.01 * rng.randn(b, n, d).astype("float32") |
| 38 | C2 = C1[:, permutation, :][:, :, permutation] + 0.01 * rng.randn(b, n, n).astype( |
| 39 | "float32" |
| 40 | ) |
| 41 | |
| 42 | M = dist_batch(X1, X2) |
| 43 | |
| 44 | res = solve_gromov_batch( |
| 45 | alpha=alpha, |
| 46 | reg=reg, |
| 47 | M=M, |
| 48 | C1=C1, |
| 49 | C2=C2, |
| 50 | max_iter=max_iter, |
| 51 | tol=tol, |
| 52 | max_iter_inner=max_iter_inner, |
| 53 | tol_inner=tol_inner, |
| 54 | symmetric=False, |
| 55 | ) |
| 56 | |
| 57 | plan_batch = res.plan |
| 58 | values_quadratic_batch = res.value_quad |
| 59 | values_linear_batch = res.value_linear |
| 60 | |
| 61 | for i in range(b): |
| 62 | M_i = M[i] |
| 63 | C1_i = C1[i] |
| 64 | C2_i = C2[i] |
| 65 | res_i = solve_gromov(C1_i, C2_i, M=M_i, alpha=alpha, symmetric=False) |
| 66 | plan_i = res_i.plan |
| 67 | values_quadratic_i = res_i.value_quad |
| 68 | values_linear_i = res_i.value_linear |
| 69 | np.testing.assert_allclose(values_linear_i, values_linear_batch[i], atol=1e-4) |
| 70 | np.testing.assert_allclose( |
| 71 | values_quadratic_i, values_quadratic_batch[i], atol=1e-4 |
| 72 | ) |
| 73 | np.testing.assert_allclose(plan_i, plan_batch[i], atol=1e-05) |
| 74 | |
| 75 | |
| 76 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected