Function test_linear_part

test/test_ucoot.py:445–578

(nx, unbalanced_solver, divergence, eps)

    ),
)
def test_linear_part(nx, unbalanced_solver, divergence, eps):
    n_samples = 5  # nb samples

    mu_s = np.array([0, 0])
    cov_s = np.array([[1, 0], [0, 1]])

    xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s, random_state=4)
    xt = xs[::-1].copy()

    px_s, px_f = ot.unif(n_samples), ot.unif(2)
    py_s, py_f = ot.unif(n_samples), ot.unif(2)

    xs_nx, xt_nx = nx.from_numpy(xs, xt)
    px_s_nx, px_f_nx, py_s_nx, py_f_nx = nx.from_numpy(px_s, px_f, py_s, py_f)

    # linear part
    M_samp = np.ones((n_samples, n_samples))
    np.fill_diagonal(np.fliplr(M_samp), 0)
    M_feat = np.ones((2, 2))
    np.fill_diagonal(M_feat, 0)
    M_samp_nx, M_feat_nx = nx.from_numpy(M_samp, M_feat)

    reg_m = (10, 5)
    alpha = (0.1, 0.2)
    max_iter_ot = 5
    max_iter = 5
    tol = 1e-7
    tol_ot = 1e-7

    # test couplings
    pi_sample, pi_feature = unbalanced_co_optimal_transport(
        X=xs,
        Y=xt,
        wx_samp=px_s,
        wx_feat=px_f,
        wy_samp=py_s,
        wy_feat=py_f,
        reg_marginals=reg_m,
        epsilon=eps,
        divergence=divergence,
        unbalanced_solver=unbalanced_solver,
        alpha=alpha,
        M_samp=M_samp,
        M_feat=M_feat,
        init_pi=None,
        init_duals=None,
        max_iter=max_iter,
        tol=tol,
        max_iter_ot=max_iter_ot,
        tol_ot=tol_ot,
        log=False,
        verbose=False,
    )

    pi_sample_nx, pi_feature_nx = unbalanced_co_optimal_transport(
        X=xs_nx,
        Y=xt_nx,
        wx_samp=px_s_nx,