Function test_jcpot_transport_class

test/test_da.py:722–798

test_jcpot_transport

(nx)

@pytest.skip_backend("jax")
@pytest.skip_backend("tf")
def test_jcpot_transport_class(nx):
    """test_jcpot_transport"""

    ns1 = 50
    ns2 = 50
    nt = 50

    Xs1, ys1 = make_data_classif("3gauss", ns1)
    Xs2, ys2 = make_data_classif("3gauss", ns2)

    Xt, yt = make_data_classif("3gauss2", nt)

    Xs1, ys1, Xs2, ys2, Xt, yt = nx.from_numpy(Xs1, ys1, Xs2, ys2, Xt, yt)

    Xs = [Xs1, Xs2]
    ys = [ys1, ys2]

    for log in [True, False]:
        otda = ot.da.JCPOTTransport(
            reg_e=1, max_iter=10000, tol=1e-9, verbose=True, log=log
        )

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

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

        # test dimensions of coupling
        for i, xs in enumerate(Xs):
            assert_equal(otda.coupling_[i].shape, ((xs.shape[0], Xt.shape[0])))

        # test all margin constraints
        mu_t = unif(nt)

        for i in range(len(Xs)):
            # test margin constraints w.r.t. uniform target weights for each coupling matrix
            assert_allclose(
                nx.to_numpy(nx.sum(otda.coupling_[i], axis=0)),
                mu_t,
                rtol=1e-3,
                atol=1e-3,
            )

            if log:
                # test margin constraints w.r.t. modified source weights for each source domain

                assert_allclose(
                    nx.to_numpy(
                        nx.dot(otda.log_["D1"][i], nx.sum(otda.coupling_[i], axis=1))
                    ),
                    nx.to_numpy(otda.proportions_),
                    rtol=1e-3,
                    atol=1e-3,
                )

        # test transform