Wasserstein-1 Dual Check

Problem

The Wasserstein-1 distance has the dual form sup over 1-Lipschitz f of E_mu[f] - E_nu[f].

Given sorted support points xs, potential values f_values, and the weights of mu and nu, compute the dual objective and check whether f is 1-Lipschitz: every secant slope |df/dx| must be at most 1 within 1e-9. Return [objective, feasible].

Hint

Kantorovich-Rubinstein duality turns the optimal transport problem into a Lipschitz-constrained optimization.

Test cases

case 1

in: [[0,1,2],[0,0.5,1],[0.5,0.5,0],[0,0.5,0.5]]

expected: [-0.5,true]

case 2

in: [[0,1],[0,2],[1,0],[0,1]]

expected: [-2,false]

case 3

in: [[0,1],[0,1],[0.5,0.5],[0.5,0.5]]

expected: [0,true]

case 4

in: [[0,1,2],[1,0,1],[0.3333333333333333,0.3333333333333333,0.3333333333333333],[0,0,1]]

expected: [-0.3333333333333334,true]

Starter codePython

def wasserstein1_dual_check(xs, f_values, weights_mu, weights_nu):
    # Return [objective, feasible]
    # Your code here
    pass

4 test cases · run your code against them in the editor