Function test_geomloss_solver

test/test_bregman.py:1166–1194

(nx, metric)

@pytest.skip_backend("jax")
@pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean"])
def test_geomloss_solver(nx, metric):
    # test sinkhorn
    n = 10
    a = ot.unif(n)
    b = ot.unif(n)

    X_s = np.reshape(1.0 * np.arange(n), (n, 1))
    X_t = np.reshape(1.0 * np.arange(0, n), (n, 1))

    ab, bb, X_sb, X_tb = nx.from_numpy(a, b, X_s, X_t)

    G_sqe = nx.to_numpy(ot.bregman.empirical_sinkhorn(X_sb, X_tb, 1, metric=metric))

    value, log = ot.bregman.empirical_sinkhorn2_geomloss(
        X_sb, X_tb, 1, metric=metric, log=True
    )
    G_geomloss = nx.to_numpy(log["lazy_plan"][:])

    print(value)

    # call with log = False
    ot.bregman.empirical_sinkhorn2_geomloss(X_sb, X_tb, 1, metric=metric)

    # check equality of plans
    np.testing.assert_allclose(G_sqe, G_geomloss, atol=1e-03)  # metric sqeuclidian

    # check error on wrong metric
    with pytest.raises(ValueError):
        ot.bregman.empirical_sinkhorn2_geomloss(X_sb, X_tb, 1, metric="wrong_metric")


def test_lazy_empirical_sinkhorn(nx):