Function verify_gluing_validity

test/test_ot.py:686–725

Test the validity of the North-West gluing.

(gamma, J, w, pi_list)

def verify_gluing_validity(gamma, J, w, pi_list):
    """
    Test the validity of the North-West gluing.
    """
    nx = get_backend(gamma)
    K = len(pi_list)
    n = pi_list[0].shape[0]
    nk_list = [pi.shape[1] for pi in pi_list]

    # Check first marginal
    a = nx.sum(gamma, axis=tuple(range(1, K + 1)))
    assert nx.allclose(a, nx.sum(pi_list[0], axis=1))

    # Check other marginals
    for k in range(K):
        b_k = nx.sum(gamma, axis=tuple(i for i in range(K + 1) if i != k + 1))
        assert nx.allclose(b_k, nx.sum(pi_list[k], axis=0))

    # Check bi-marginals
    for k in range(K):
        gamma_0k = nx.sum(gamma, axis=tuple(i for i in range(1, K + 1) if i != k + 1))
        assert nx.allclose(gamma_0k, pi_list[k])

    # Check that N <= n + sum_k n_k - K
    N = J.shape[0]
    n_k_sum = sum(nk_list)
    assert N <= n + n_k_sum - K, f"N={N}, n={n}, sum(n_k)={n_k_sum}, K={K}"

    # Check that w is on the simplex
    w_sum = nx.sum(w)
    assert nx.allclose(w_sum, 1), f"Sum of weights w is not 1: {w_sum}"

    # Check that gamma_1...K and (J, w) are consistent
    rho = nx.zeros(nk_list, type_as=gamma)
    for i in range(N):
        jj = J[i]
        rho[tuple(jj)] += w[i]

    gamma_1toK = nx.sum(gamma, axis=0)
    assert nx.allclose(rho, gamma_1toK), "rho and gamma_1...K are not consistent"


def test_north_west_mm_gluing():