Function test_solve_sample_geomloss

test/test_solvers.py:610–667

(nx, metric)

@pytest.skip_backend("jax")
@pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean"])
def test_solve_sample_geomloss(nx, metric):
    # test solve_sample when is_Lazy = False
    n_samples_s = 13
    n_samples_t = 7
    n_features = 2
    rng = np.random.RandomState(0)

    x = rng.randn(n_samples_s, n_features)
    y = rng.randn(n_samples_t, n_features)
    a = ot.utils.unif(n_samples_s)
    b = ot.utils.unif(n_samples_t)

    xb, yb, ab, bb = nx.from_numpy(x, y, a, b)

    sol0 = ot.solve_sample(xb, yb, ab, bb, reg=1)

    # solve signe weights
    sol = ot.solve_sample(xb, yb, ab, bb, reg=1, method="geomloss")
    assert_allclose_sol(sol0, sol)

    sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=False, method="geomloss")
    assert_allclose_sol(sol0, sol)

    sol1 = ot.solve_sample(
        xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_tensorized"
    )
    np.testing.assert_allclose(
        nx.to_numpy(sol1.lazy_plan[:]),
        nx.to_numpy(sol.lazy_plan[:]),
        rtol=1e-5,
        atol=1e-5,
    )

    sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_online")
    np.testing.assert_allclose(
        nx.to_numpy(sol1.lazy_plan[:]),
        nx.to_numpy(sol.lazy_plan[:]),
        rtol=1e-5,
        atol=1e-5,
    )

    sol1 = ot.solve_sample(
        xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_multiscale"
    )
    np.testing.assert_allclose(
        nx.to_numpy(sol1.lazy_plan[:]),
        nx.to_numpy(sol.lazy_plan[:]),
        rtol=1e-5,
        atol=1e-5,
    )

    sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss")
    np.testing.assert_allclose(
        nx.to_numpy(sol1.lazy_plan[:]),
        nx.to_numpy(sol.lazy_plan[:]),
        rtol=1e-5,
        atol=1e-5,
    )