2D Histogram Wasserstein Distance via POT Library

testOt.py

"""
Programmer: Chris Tralie
Purpose: To use the POT library (https://github.com/rflamary/POT)
to compute the Entropic regularized Wasserstein distance 
between points on a 2D grid
"""
import numpy as np
import matplotlib.pyplot as plt
import ot

def testMovingDisc():
    """
    Show optimal transport on a moving disc in a 50x50 grid
    """
    ## Step 1: Setup problem
    pix = np.linspace(-1, 1, 80)
    # Setup grid
    X, Y = np.meshgrid(pix, pix)
    # Compute pariwise distances between points on 2D grid so we know
    # how to score the Wasserstein distance
    coords = np.array([X.flatten(), Y.flatten()]).T
    coordsSqr = np.sum(coords**2, 1)
    M = coordsSqr[:, None] + coordsSqr[None, :] - 2*coords.dot(coords.T)
    M[M < 0] = 0
    M = np.sqrt(M)
    ts = np.linspace(-0.8, 0.8, 100)
    
    ## Step 2: Compute L2 distances and Wasserstein
    Images = []
    radius = 0.2
    L2Dists = [0.0]
    WassDists = [0.0]
    for i, t in enumerate(ts):
        I = 1e-5 + np.array((X-t)**2 + (Y-t)**2 < radius**2, dtype=float)
        I /= np.sum(I)
        Images.append(I)
        if i > 0:
            L2Dists.append(np.sqrt(np.sum((I-Images[0])**2)))
            wass = ot.sinkhorn2(Images[0].flatten(), I.flatten(), M, 1.0)
            print(wass)
            WassDists.append(wass)
    
    ## Step 3: Make Animation
    L2Dist = np.array(L2Dists)
    WassDists = np.array(WassDists)
    I0 = Images[0]
    plt.figure(figsize=(15, 5))
    displacements = np.sqrt(2)*(ts - ts[0])
    for i, I in enumerate(Images):
        plt.clf()
        D = np.concatenate((I0[:, :, None], I[:, :, None], 0*I[:, :, None]), 2)
        D = D*255/np.max(I0)
        D = np.array(D, dtype=np.uint8)
        plt.subplot(131)
        plt.imshow(D, extent = (pix[0], pix[-1], pix[-1], pix[0]))
        plt.subplot(132)
        plt.plot(displacements, L2Dists)
        plt.stem([displacements[i]], [L2Dists[i]])
        plt.xlabel("Displacements")
        plt.ylabel("L2 Dist")
        plt.title("L2 Dist")
        plt.subplot(133)
        plt.plot(displacements, WassDists)
        plt.stem([displacements[i]], [WassDists[i]])
        plt.xlabel("Displacements")
        plt.ylabel("Wasserstein Dist")
        plt.title("Wasserstein Dist")
        plt.savefig("%i.png"%i, bbox_inches='tight')

if __name__ == '__main__':
    testMovingDisc()

Great questions. I unfortunately have not used this library beyond this example (which was to help a student), so I might not be the best person to ask, but I will try

Yes, the nonlinear part is strange, but I wonder if it's just because of the shape and discretization of the discs. It would be worth trying this with two squares that move past each other only in the x direction, for example.

That's strange that it's not equal to zero from I to itself, because it should be

Yes, I was as confused as you about problems when there are zeros, and that is indeed why I added 1e-5

Thank you for your answer. I tryed with small squares of 4 pixels and the problem was already here...
I will continue to investigate !

I finally have the answer to my questions, so I thought I would give them to you too:

The function ot.sinkhorn2 is actually an aproximation of the optimal transport problem. This is why it causes some problems.
The exact value for the Wasserstein distance is obtained by using the ot.emd2 function instead.
It is a bit longer and the number of iterations must be increased, but it works !

Regards

That is great to know, thank you so much for the update!

…

On Mon, Jul 29, 2019, 9:16 AM raphiol ***@***.***> wrote: I finally have the answer to my questions, so I thought I would give them to you too: The function ot.sinkhorn2 is actually an aproximation of the optimal transport problem. This is why it causes some problems. The exact value for the Wasserstein distance is obtained by using the ot.emd2 function instead. It is a bit longer and the number of iterations must be increased, but it works ! Regards — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <https://gist.github.com/66352ae6ab06c009f02c705385a446f3?email_source=notifications&email_token=AAJWDZXAFPXUNIS62VXKFRTQB3UUNA5CNFSM4IE2C44KYY3PNVWWK3TUL52HS4DFVNDWS43UINXW23LFNZ2KUY3PNVWWK3TUL5UWJTQAFWDZQ#gistcomment-2983832>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AAJWDZSTILXVWH6V62C5NMLQB3UUNANCNFSM4IE2C44A> .

This gist is perfect for describing earth movers distance. Thank you!

Yaaay so glad it helped! Usually when I get these notifications there's a bug, so I'm happy to see a fun one for a change

…

On Mon, Mar 23, 2020 at 5:40 PM Justin ***@***.***> wrote: ***@***.**** commented on this gist. ------------------------------ This gist is perfect for describing earth movers distance. Thank you! — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <https://gist.github.com/66352ae6ab06c009f02c705385a446f3#gistcomment-3224778>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJWDZXKDMCRTLPKGZNYYBLRI7JLXANCNFSM4IE2C44A> .