Function test_emd_transport_class

test/test_da.py:488–586

test_sinkhorn_transport

(nx)

@pytest.skip_backend("jax")
@pytest.skip_backend("tf")
def test_emd_transport_class(nx):
    """test_sinkhorn_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)

    otda = ot.da.EMDTransport()

    # test its computed
    otda.fit(Xs=Xs, Xt=Xt)
    assert hasattr(otda, "cost_")
    assert hasattr(otda, "coupling_")

    # test dimensions of coupling
    assert_equal(otda.cost_.shape, ((Xs.shape[0], Xt.shape[0])))
    assert not np.any(np.isnan(nx.to_numpy(otda.cost_))), "cost is finite"
    assert_equal(otda.coupling_.shape, ((Xs.shape[0], Xt.shape[0])))
    assert np.all(np.isfinite(nx.to_numpy(otda.coupling_))), "coupling is finite"

    # test 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(transp_Xs.shape, Xs.shape)

    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)

    # check label propagation
    transp_yt = otda.transform_labels(ys)
    assert_equal(transp_yt.shape[0], yt.shape[0])
    assert_equal(transp_yt.shape[1], len(np.unique(ys)))

    # check inverse label propagation
    transp_ys = otda.inverse_transform_labels(yt)
    assert_equal(transp_ys.shape[0], ys.shape[0])
    assert_equal(transp_ys.shape[1], len(np.unique(yt)))