Function test_init_plans

test/test_ucoot.py:165–293

(nx, unbalanced_solver, divergence, eps)

    ),
)
def test_init_plans(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)
    G0_samp = px_s[:, None] * py_s[None, :]
    G0_feat = px_f[:, None] * py_f[None, :]

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

    reg_m = (1, 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=None,
        M_feat=None,
        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,
        wx_feat=px_f_nx,
        wy_samp=py_s_nx,
        wy_feat=py_f_nx,
        reg_marginals=reg_m,
        epsilon=eps,