How is transport plan (couplings) retrieved from ot.da.LinearGWTransport()?

[ttsesm]

Hi,

I am trying to use the ot.da.LinearGWTransport() function as described in the POT documentation as follows:

def match_gaussians(Cs, Ms, ct, mt, verbose=False):
    """Match 3D Gaussians to 2D Gaussians using Linear Gromov-Wasserstein Transport."""
    # Compute cost matrices
    C3D, C2D = compute_cost_matrices(Cs, Ms, ct, mt)

    # Initialize and fit the LinearGWTransport model
    gw = ot.da.LinearGWTransport(log=verbose)
    gw.fit(Xs=Ms, Xt=mt, ys=C3D, yt=C2D)

    # Get the transport plan
    transport_plan = gw.coupling # <------------------------------------- This doesn't exist

    # Get the linear operator (projection matrix)
    projection_matrix = gw.L # <------------------------------------- This doesn't exist

    return projection_matrix, transport_plan

where Cs, Ms and ct, mt are my 3D and 2D gaussians respectively (i.e. covariance and mean values).

However, I am not sure how to retrieve back the transport plan and the projection matrix from the fitted distribution. From what I've noticed I can get back the A and B matrices but I am not sure how these are related to the transport plan and projection matrix.

I would appreciated if someone has an idea and/or provide some feedback.

Thanks.

[rflamary]

Linear Transport methods assumes that the left and right distributions are multivariate normal distributions. In practice their parameters such as mean and covariance are estimated from the data.

So they do not provide a transport plan between samples but a continuous affine mapping that is the Monge mapping. You can interpret the transport by looking at the displacement of the samples but there is not OT plan giving you correspondences for those methods (a coarse approximation would be to do a closets sample assignment after mapping but this will not be a transport plan).

[ttsesm]

Wait a moment, because I think I got confused now, so the ot.da.LinearGWTransport() function works on discrete samples and not in distributions? Because in the documentation it mentions about "empirical distributions" and refers to the paper also related to "Gaussian distributions". Then as it is mentioned the fit() builds a coupling matrix from source and target sets of samples.

Sorry for my poor understanding/questions but I am still new both in OT and how to use it through POT.

With @eloitanguy we had a similar discussion in #599 but as far as I remember ot.da.LinearGWTransport() has not been implemented back then. I guess, OT between GMM might fit better my problem?

[eloitanguy]

Hi, GMM OT is promising, unfortunately it only makes sense between GMMs on the same space...
You would have to look into (GMM-) generalised barycenters, which would be an extension of https://arxiv.org/abs/2105.09755 or of https://arxiv.org/abs/2407.13445 , based on GMM-OT https://arxiv.org/abs/1907.05254 .
Another idea would be Gromov for GMMs https://arxiv.org/abs/2104.07970
It all seems feasible, but I expect it to take some effort for a numerically functional method

[ttsesm]

Hi @eloitanguy, thanks for the feedback. Well the majority of the literature you posted I was aware and I am trying to build my solution based on these, I was missing only your recent publication which looks quite nice and interesting as well.

Regarding the GMM OT, and making sense to use it only between GMMs of the same space. Could be possibly a solution to just treat the 2-d Gaussians as degenerate 3-d gaussians and calculate it that way? The issue I see is that it requires full rank covariances, right? But do you think that it might be able to leverage the fact that the Gaussians are orthogonally projected onto the 2-d plane to get around the full rank requirement... then I guess it should give a unique 3rd eigenvector for the (2-d) covariances. Not sure though how to implement this in practice... 🤔

For the generalized barycenters, if I am not wrong it works on discrete samples, i.e. points, right? Plus, I would need to have the projection matrix, something that I do not since I am looking for that. An idea would be to use the mean values, and points on the surfaces of the 3D gaussians (ellipsoids) and 2D gaussians (ellipses) as my discrete points but I would still need the projection matrix.

Finally, Gromov for GMMs is what I was trying to use here by utilizing ot.da.LinearGWTransport() but I am not sure whether I am using it correctly, as @rflamary was mentioning that the mapping should match the distributions. However, didn't really understood what that means or how to evaluate it in practice. I have a self contained script bellow. I am not sure whether you can and have the time to look at it but I would appreciate it.

import numpy as np
import ot
from scipy.linalg import sqrtm

# 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

def match_gaussians(Cs, Ms, ct, mt, verbose=False):
    """Match 3D Gaussians to 2D Gaussians using Linear Gromov-Wasserstein Transport."""
    # Compute cost matrices
    C3D, C2D = compute_cost_matrices(Cs, Ms, ct, mt)

    # Create uniform weights
    p = ot.unif(len(Cs))
    q = ot.unif(len(ct))

    # Initialize and fit the LinearGWTransport model
    gw = ot.da.LinearGWTransport(log=verbose)
    gw.fit(Xs=Ms, Xt=mt, ys=C3D, yt=C2D)

    # Get the transport plan
    transport_plan = gw.coupling_ # <------------------------------------- This doesn't exist

    # Get the linear operator (projection matrix)
    projection_matrix = gw.L_ # <------------------------------------- This doesn't exist

    return projection_matrix, transport_plan

def project_and_validate(Cs, Ms, ct, mt, projection_matrix):
    """Project 3D Gaussians to 2D and validate results."""
    for i, (C3D, M3D) in enumerate(zip(Cs, Ms)):
        # Project 3D Gaussian to 2D
        C2D_projected = projection_matrix @ C3D @ projection_matrix.T
        M2D_projected = projection_matrix @ M3D

        print(f"\nGaussian {i+1}:")
        print("Original 3D mean:")
        print(M3D)
        print("Projected 2D mean:")
        print(M2D_projected)
        print("Original 3D covariance:")
        print(C3D)
        print("Projected 2D covariance:")
        print(C2D_projected)


# Match Gaussians
projection_matrix, transport_plan = match_gaussians(Cs, Ms, ct, mt, verbose=True)

print("Estimated projection matrix:")
print(projection_matrix)

print("\nTransport plan:")
print(transport_plan)

# Project and validate results
project_and_validate(Cs, Ms, ct, mt, projection_matrix)

[ttsesm]

Btw, I've also played a bit with Sinkhorn during the weekend. Please find attached another self contained script as well. Basically, I've tried using the Sinkhorn algorithm for optimal transport. Here's a breakdown of the key components:

1. bures_wasserstein_distance: Computes the distance between covariance matrices using the Bures-Wasserstein metric, which is more appropriate for comparing covariance matrices than the e.g. Frobenius norm.
2. cost_matrix: Computes the cost matrix for optimal transport, considering both covariance and mean differences.
3. objective_function: Defines the objective to minimize, which is the total cost of optimal transport between projected 3D ellipsoids and 2D ellipses.
4. estimate_projection_ot: The main function that estimates the projection matrix and translation. It uses multiple random restarts to avoid local minima.
5. Sinkhorn algorithm: Used for computing the optimal transport plan, which in principle can handle noise and inexact correspondences better than hard assignment methods.

This solution seems to work quite well when we have a 1-to-1 correlation between my two sets, the 2D gaussians (ellipses) and the 3D gaussians (ellipsoids) as you can see in the following graphs for 15, 65 and 115 samples that I've tested respectively:

However, when we add noise into the number of ellipses (15 ellipsoids vs. 25 ellipses), which is gonna be my real case scenario, Sinkhorn fails as well:

I've tried to add extra "noisy" ellipsoids as well so that we can make the calculation of the transport plan in a squared format but this didn't help much. This solution looks promising but I need to resolve the noisy cases. If any of you has any idea how to make Sinkhorn more robust to noisy scenarios, I would be happy to test it.

Thanks.

import numpy as np
import ot
from scipy.linalg import sqrtm
from scipy.optimize import minimize
import matplotlib
from matplotlib import pyplot as plt

# 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 bures_wasserstein_distance(C1, C2):
    """Compute the Bures-Wasserstein distance between two covariance matrices."""
    sqrt_C1 = sqrtm(C1)
    return np.trace(C1 + C2 - 2 * sqrtm(sqrt_C1 @ C2 @ sqrt_C1))

def project_3d_to_2d(Sigma3D, M3D, P, t):
    """Project 3D covariance and mean to 2D using projection matrix P and translation t."""
    Sigma2D = P @ Sigma3D @ P.T
    M2D = P @ M3D + t
    return Sigma2D, M2D

def cost_matrix(Cs, Ms, ct, mt, P, t):
    """Compute the cost matrix for optimal transport."""
    n_3d = len(Cs)
    n_2d = len(ct)
    C = np.zeros((n_3d, n_2d))

    for i in range(n_3d):
        Sigma2D, M2D = project_3d_to_2d(Cs[i], Ms[i], P, t)
        for j in range(n_2d):
            cov_dist = bures_wasserstein_distance(Sigma2D, ct[j])
            mean_dist = np.linalg.norm(M2D - mt[j])
            C[i, j] = cov_dist + mean_dist

    return C

def objective_function(params, Cs, Ms, ct, mt):
    """Objective function to minimize."""
    P = params[:6].reshape(2, 3)
    t = params[6:]
    C = cost_matrix(Cs, Ms, ct, mt, P, t)

    # Compute optimal transport
    a = ot.unif(len(Cs))
    b = ot.unif(len(ct))
    pi = ot.sinkhorn(a, b, C, reg=0.1)

    return np.sum(pi * C)

def estimate_projection_ot(Cs, Ms, ct, mt, n_restarts=10):
    """Estimate projection matrix and translation using optimal transport."""
    best_params = None
    best_loss = float('inf')

    for _ in range(n_restarts):
        # Random initialization
        P_init = np.random.randn(2, 3)
        P_init /= np.linalg.norm(P_init, axis=1)[:, None]
        t_init = np.random.randn(2)
        params_init = np.concatenate([P_init.flatten(), t_init])

        # Optimize
        result = minimize(objective_function, params_init, args=(Cs, Ms, ct, mt),
                          method='L-BFGS-B', options={'maxiter': 1000})

        if result.fun < best_loss:
            best_loss = result.fun
            best_params = result.x

    P = best_params[:6].reshape(2, 3)
    t = best_params[6:]
    return P, t

# Use the provided data
ct = np.array(ct, dtype=np.float64)
mt = np.array(mt, dtype=np.float64)
Cs = np.array(Cs, dtype=np.float64)
Ms = np.array(Ms, dtype=np.float64)

# Estimate projection matrix and translation
estimated_P, estimated_t = estimate_projection_ot(Cs, Ms, ct, mt)

print("Estimated projection matrix:")
print(estimated_P)
print("\nEstimated translation:")
print(estimated_t)

# Validate the results
for i in range(len(Cs)):
    projected_cov, projected_mean = project_3d_to_2d(Cs[i], Ms[i], estimated_P, estimated_t)
    print(f"\nEllipsoid {i+1}:")
    print("Original 3D covariance:")
    print(Cs[i])
    print("Projected 2D covariance:")
    print(projected_cov)
    print("Original 3D mean:")
    print(Ms[i])
    print("Projected 2D mean:")
    print(projected_mean)

# Compute final correspondences
C = cost_matrix(Cs, Ms, ct, mt, estimated_P, estimated_t)
a = ot.unif(len(Cs))
b = ot.unif(len(ct))
pi = ot.sinkhorn(a, b, C, reg=0.1)

print("\nFinal correspondences (3D index, 2D index, probability):")
for i in range(len(Cs)):
    j = np.argmax(pi[i])
    print(f"({i}, {j}, {pi[i, j]:.4f})")

fig = plt.figure()
ax = fig.add_subplot(111)
cax = ax.matshow(pi)
fig.colorbar(cax)

[ttsesm]

@rflamary, @eloitanguy any idea on the examples I've attached above and how I could possibly improve my output especially for the noisy case?

[eloitanguy]

Hi @ttsesm, thanks for the interesting questions!
It seems to us that you are asking open research questions, which, while certainly valuable for the OT community, are not well suited for the POT discussion page. If you are looking for a research collaboration, please approach privately people whose reasearch you find the most relevant. I'm sorry we couldn't invest a lot of time to your open questions here, as you can imagine we are both quite swamped.

PS I cannot speak for everyone, but for safety reasons I would never open a stranger's zip file ;)

[ttsesm]

Sure, I understand @eloitanguy 👍

In any case, I've removed the zip files and replaced them by the actual code snippet just in case someone wants to have a look at it or test it.

[ttsesm]

@eloitanguy based on the above code and experimenting a bit more, I've managed to have it running also for the non equal samples test as you can see below:

However, another question now is whether there is a way to smoothen even more the transport plan so that to reduce the noise (the appearing outliers). For example from the graph it is clear that there is correlation between the first 250 samples of both sets (visualized by the diagonal). However, due to the outliers I am not getting my matches to be always on the diagonal of these first 250 samples (actually on 300 samples experiment I am getting approx. a 46% accuracy, while the correlation through the diagonal clear). I've tried increasing a bit the reg= value of the ot.sinkhorn() parameters but didn't really improve a lot. Is there another way, that I could consider or that you might be aware?