Question about using UCOOT

[ttsesm]

Hi @6Ulm,

Congrats for your work. I am trying to understand whether your work could be useful for my task.

I am using the following objective function to compute the optimal transport for my input data:

def objective_function(params, Cs_3d, ms_3d, Ct_2d, mt_2d):
    """Objective function to minimize."""
    P = params[:6].reshape(2, 3)
    t = params[6:]
    C = cost_matrix2(Cs_3d, ms_3d, Ct_2d, mt_2d, P, t)

    # Compute optimal transport
    a = ot.unif(len(Cs_3d), type_as=C)
    b = ot.unif(len(Ct_2d), type_as=C)
    pi = ot.sinkhorn(a, b, C, reg=0.1)

    return torch.sum(pi * C)

This works relatively well as you can see below from the printed transport plans:

which it seems to catch the undergone correlation between my samples. My only issue is that practically due to the noise/outliers that are also apparent in the graphs my accuracy on trying to find the correspondences is not that high since I am getting matches outside the diagonal line which should be the correct one. Thus, I was thinking whether I could use your UCOOT implementation to smoothen a bit more my transport plan and eliminate as much as possible my outliers. Do you think that this would be a valid approach? Also I saw you have also opened a pull request in POT for integrating the algorithm in the library which would be nice to use it directly from there I guess.

Thanks, and I would be interested to hear your opinion.

[6Ulm]

Hi @ttsesm

Sorry for the late response. Definitely you can try the unbalanced version UCOOT if you believe that outliers present in your data. I'm working on the PR, so that both UCOOT and FUGW can be released asap.

[ttsesm]

Hi @6Ulm,

No worries, it sounds good.

Meanwhile I've tried to run the version in this repository but without great success. I would appreciate if you could have a look in the following code snippet and see whether I am doing something wrong.

My input data are some 3D and 2D gaussian distribution, so as you can understand they are in a different space. I've tried to use both the UCOOT and FUGW but I couldn't get any reasonable output.

import numpy as np
import ot
import torch
from scipy.linalg import sqrtm
from scipy.optimize import minimize

from ucoot.megawass import MegaWass

import matplotlib
from matplotlib import pyplot as plt
import seaborn as sns
from tqdm.auto import tqdm
matplotlib.use("Qt5Agg")

# For an example of comparing 2D gaussians without noise as 1-to-1 with the 3D gaussians uncomment the following input data

# 2D Gaussians without noise
# Covariances, 2x2 matrices
ct = np.array([[[0.00049392, 0.00018833],
    [0.00018833, 0.00007765]],

   [[0.00037171, 0.00021303],
    [0.00021303, 0.00012497]],

   [[0.00004997, 0.00002954],
    [0.00002954, 0.00012977]],

   [[0.00005111, 0.00000863],
    [0.00000863, 0.00002055]],

   [[0.00022878, 0.00014253],
    [0.00014253, 0.00023158]],

   [[0.00010538, 0.00007804],
    [0.00007804, 0.00030141]],

   [[0.00001159, 0.00000829],
    [0.00000829, 0.00000615]],

   [[0.00007754, 0.00002211],
    [0.00002211, 0.00011025]],

   [[0.00070302, 0.00060664],
    [0.00060664, 0.00052477]],

   [[0.00109267, 0.00122162],
    [0.00122162, 0.00137384]]], dtype=np.float32)

# Means, 1x2 vectors
mt = np.array([[-0.34575066, -0.05927338],
   [-0.16067392, -0.15656751],
   [ 0.0884882 , 0.7348885 ],
   [-0.09218572, 0.05595953],
   [ 0.21521865, 0.38351792],
   [-0.22240582, 0.30064186],
   [-0.0520223 , -0.18196055],
   [-0.4030493 , 0.22060394],
   [-0.17406262, -0.09412329],
   [ 0.3939966 , 0.21849754]], dtype=np.float32)


# # For an example of comparing 2D gaussians with noise use the following input data
#
# # 2D Gaussians with noise
# # Covariances, 2x2 matrices
# ct = np.array([[[0.00049392, 0.00018833],
#         [0.00018833, 0.00007765]],
#
#        [[0.00037171, 0.00021303],
#         [0.00021303, 0.00012497]],
#
#        [[0.00004997, 0.00002954],
#         [0.00002954, 0.00012977]],
#
#        [[0.00005111, 0.00000863],
#         [0.00000863, 0.00002055]],
#
#        [[0.00022878, 0.00014253],
#         [0.00014253, 0.00023158]],
#
#        [[0.00010538, 0.00007804],
#         [0.00007804, 0.00030141]],
#
#        [[0.00001159, 0.00000829],
#         [0.00000829, 0.00000615]],
#
#        [[0.00007754, 0.00002211],
#         [0.00002211, 0.00011025]],
#
#        [[0.00070302, 0.00060664],
#         [0.00060664, 0.00052477]],
#
#        [[0.00109267, 0.00122162],
#         [0.00122162, 0.00137384]],
#
#        [[0.00034284, 0.00019008],
#         [0.00019008, 0.00011378]],
#
#        [[0.00002524, 0.00003232],
#         [0.00003232, 0.00006822]],
#
#        [[0.00007186, 0.00005018],
#         [0.00005018, 0.00004073]],
#
#        [[0.0001403 , 0.00004997],
#         [0.00004997, 0.00037811]],
#
#        [[0.00005589, 0.00002819],
#         [0.00002819, 0.00002271]]], dtype=float32)
# # Means, 1x2 vectors
# mt = np.array([[-0.34575066, -0.05927338],
#        [-0.16067392, -0.15656751],
#        [ 0.0884882 ,  0.7348885 ],
#        [-0.09218572,  0.05595953],
#        [ 0.21521865,  0.38351792],
#        [-0.22240582,  0.30064186],
#        [-0.0520223 , -0.18196055],
#        [-0.4030493 ,  0.22060394],
#        [-0.17406262, -0.09412329],
#        [ 0.3939966 ,  0.21849754],
#        [ 0.1957326 ,  0.2494914 ],
#        [-0.09074225,  0.04654403],
#        [ 0.44695464,  0.21314013],
#        [-0.11307922, -0.11961696],
#        [-0.12591828, -0.17904231]], dtype=float32)

##########################################################################################################################################################################################

# 3D Gaussians
# Covariances, 3x3 matrices
Cs = np.array([[[ 0.00055704, 0.0002181 , 0.00020697],
    [ 0.0002181 , 0.00026268, 0.00017932],
    [ 0.00020697, 0.00017932, 0.00013152]],

   [[ 0.00029736, 0.00036497, -0.00025849],
    [ 0.00036497, 0.0005222 , -0.00034746],
    [-0.00025849, -0.00034746, 0.00028704]],

   [[ 0.00013889, -0.00002427, 0.00009434],
    [-0.00002427, 0.000106 , -0.00010998],
    [ 0.00009434, -0.00010998, 0.00020017]],

   [[ 0.00013997, 0.0000322 , -0.00003177],
    [ 0.0000322 , 0.00013866, 0.0000159 ],
    [-0.00003177, 0.0000159 , 0.00008226]],

   [[ 0.00015876, 0.00015987, 0.00013226],
    [ 0.00015987, 0.00046588, -0.0000928 ],
    [ 0.00013226, -0.0000928 , 0.00045325]],

   [[ 0.00014214, -0.00009106, 0.00025811],
    [-0.00009106, 0.00042213, 0.00002785],
    [ 0.00025811, 0.00002785, 0.00058036]],

   [[ 0.00003409, -0.00000352, 0.00000672],
    [-0.00000352, 0.00000058, -0.00000025],
    [ 0.00000672, -0.00000025, 0.00000226]],

   [[ 0.00025886, -0.00007174, -0.00000447],
    [-0.00007174, 0.00020796, 0.00000071],
    [-0.00000447, 0.00000071, 0.00017473]],

   [[ 0.00000316, -0.00008337, -0.00002749],
    [-0.00008337, 0.00292826, 0.00008949],
    [-0.00002749, 0.00008949, 0.00106939]],

   [[ 0.00080745, -0.00143929, -0.00031194],
    [-0.00143929, 0.00750104, 0.00179461],
    [-0.00031194, 0.00179461, 0.0016364 ]]], dtype=np.float32)

# Means, 1x3 vectors
Ms = np.array([[ 0.17411666, 0.5035763 , 0.42941484],
   [ 0.5251834 , -0.27657598, -0.18892948],
   [-0.5073861 , 0.46957424, -0.90752286],
   [ 0.56422734, -0.4835634 , -0.91173404],
   [-0.5370846 , 0.22988465, -0.2466668 ],
   [ 0.04097709, 0.48167574, -0.23595962],
   [ 0.49995327, -0.5364788 , -0.3129182 ],
   [ 0.4078352 , 0.46962914, -0.3338474 ],
   [ 0.5331596 , -0.24272342, -0.30946448],
   [-0.57427835, -0.20426048, -0.28931147]], dtype=np.float32)


def gaussian_wasserstein_distance(C1, C2, m1, m2):
    """Compute the Wasserstein distance between two Gaussian distributions."""
    diff_means = np.linalg.norm(m1 - m2)**2
    cov_dist = np.trace(C1 + C2 - 2*sqrtm(sqrtm(C1) @ C2 @ sqrtm(C1)))
    return diff_means + cov_dist

def compute_cost_matrices(Cs, Ms, ct, mt):
    """Compute cost matrices for 3D and 2D Gaussians."""
    n_3d = len(Cs)
    n_2d = len(ct)

    C3D = np.zeros((n_3d, n_3d))
    C2D = np.zeros((n_2d, n_2d))

    for i in range(n_3d):
        for j in range(i+1, n_3d):
            C3D[i, j] = C3D[j, i] = gaussian_wasserstein_distance(Cs[i], Cs[j], Ms[i], Ms[j])

    for i in range(n_2d):
        for j in range(i+1, n_2d):
            C2D[i, j] = C2D[j, i] = gaussian_wasserstein_distance(ct[i], ct[j], mt[i], mt[j])

    return C3D, C2D

rho = (0.1, 0.1) # use (float("inf"), float("inf")) if use COOT
eps = (0.1)

# Compute cost matrices
C3D, C2D = compute_cost_matrices(Cs, Ms, ct, mt)

# C3D = []
# for i in range(0, len(Cs)):
#     cov = np.eye(4)
#     center = np.eye(4)
#     center[:3, 3] = Ms[i]
#     cov[:3, :3] = Cs[i]
#     Q = center @ cov @ center.T
#
#     C3D.append(Q)
#
# C2D = []
#
# for i in range(0, len(ct)):
#     cov = np.eye(3);
#     cov[:2, :2] = ct[i]
#     center = np.eye(3)
#     center[:2, 2] = mt[i]
#     q = center @ cov @ center.T
#
#     C2D.append(q)
#
# ellipsoids_torch = torch.from_numpy(np.array(C3D)).to(dtype=torch.float64, device="cuda")
# ellipses_torch = torch.from_numpy(np.array(C2D)).to(dtype=torch.float64, device="cuda")
# megawass = MegaWass(nits_bcd=20, nits_uot=1000, tol_bcd=1e-6, tol_uot=1e-6, eval_bcd=1, eval_uot=20)
# (pi_samp, pi_feat), _, log_cost, log_ent_cost = megawass.solver_fucoot(
#         X=ellipsoids_torch,
#         Y=ellipses_torch,
#         rho=rho,
#         eps=eps,
#         log=True,
#         verbose=True,
#         early_stopping_tol=1e-6
#     )

ellipsoids_torch = torch.from_numpy(C3D).to(dtype=torch.float64, device="cuda")
ellipses_torch = torch.from_numpy(C2D).to(dtype=torch.float64, device="cuda")
megawass = MegaWass(nits_bcd=20, nits_uot=1000, tol_bcd=1e-6, tol_uot=1e-6, eval_bcd=1, eval_uot=20)
(pi_samp, pi_feat), _, log_cost, log_ent_cost = megawass.solver_fugw_simple(
        X=ellipsoids_torch,
        Y=ellipses_torch,
        rho=rho,
        eps=eps,
        log=True,
        verbose=True,
        early_stopping_tol=1e-6
    )

fig = plt.figure()
ax = fig.add_subplot(111)
cax = ax.matshow(pi_feat.detach().cpu().numpy())
fig.colorbar(cax)

print("finished")

The solver_fucoot I couldn't really make it work since the dimensions of my input data seems to be out of range and I am not quite sure how to properly pass it in the existing format. Then the fugw version based on the cost matrices seems to give poor results as you can see on the estimated feature matrix bellow:

Can you confirm that I am using it properly or I am doing something wrong? In theory and in practice I should be getting a diagonal matrix if everything was working fine.

[ttsesm]

The problem that I am mainly facing is how to treat my data as input to the corresponding functions. For example I am having the same problem when I am trying to use the robust Gromov-Wasserstein (which actually cites your work) but again I am having difficulty to set my data in a correct format. The way that I am creating my cost matrices apparently is not the best one... 🤔