Function test_emd_laplace_class

test/test_da.py:841–908

test_emd_laplace_transport

(nx)

@pytest.skip_backend("jax")
@pytest.skip_backend("tf")
def test_emd_laplace_class(nx):
    """test_emd_laplace_transport"""
    ns = 50
    nt = 50

    Xs, ys = make_data_classif("3gauss", ns)
    Xt, yt = make_data_classif("3gauss2", nt)

    Xs, ys, Xt, yt = nx.from_numpy(Xs, ys, Xt, yt)

    for log in [True, False]:
        otda = ot.da.EMDLaplaceTransport(
            reg_lap=0.01, max_iter=1000, tol=1e-9, verbose=False, log=log
        )

        # test its computed
        otda.fit(Xs=Xs, ys=ys, Xt=Xt)

        assert hasattr(otda, "coupling_")
        assert hasattr(otda, "log_")

        # test dimensions of coupling
        assert_equal(otda.coupling_.shape, ((Xs.shape[0], Xt.shape[0])))

        # test all margin constraints
        mu_s = unif(ns)
        mu_t = unif(nt)

        assert_allclose(
            nx.to_numpy(nx.sum(otda.coupling_, axis=0)), mu_t, rtol=1e-3, atol=1e-3
        )
        assert_allclose(
            nx.to_numpy(nx.sum(otda.coupling_, axis=1)), mu_s, rtol=1e-3, atol=1e-3
        )

        # test transform
        transp_Xs = otda.transform(Xs=Xs)
        [assert_equal(x.shape, y.shape) for x, y in zip(transp_Xs, Xs)]

        Xs_new = nx.from_numpy(make_data_classif("3gauss", ns + 1)[0])
        transp_Xs_new = otda.transform(Xs_new)

        # check that the oos method is working
        assert_equal(transp_Xs_new.shape, Xs_new.shape)

        # test inverse transform
        transp_Xt = otda.inverse_transform(Xt=Xt)
        assert_equal(transp_Xt.shape, Xt.shape)

        Xt_new = nx.from_numpy(make_data_classif("3gauss2", nt + 1)[0])
        transp_Xt_new = otda.inverse_transform(Xt=Xt_new)

        # check that the oos method is working
        assert_equal(transp_Xt_new.shape, Xt_new.shape)

        # test fit_transform
        transp_Xs = otda.fit_transform(Xs=Xs, Xt=Xt)
        assert_equal(transp_Xs.shape, Xs.shape)