Domain: Math — Global Computational Mathematics

Problem:
Let \\(\Omega = [0,1]^2\\).  Consider two probability densities: \\(\rho_0 = 1 + 0.5\sin(2\pi x)\sin(2\pi y)\\) and \\(\rho_1 = 1+0.5\cos(2\pi x)\cos(2\pi y)\\), normalized so they integrate to 1 over \\(\Omega\\). The optimal transport map \\(T:\Omega\to\Omega\\) that pushes \\(\rho_0\\) to \\(\rho_1\\) under the quadratic cost \\(c(x, y) = \|x- y\|^2\\) (Note: this cost function is the squared Euclidean distance and intentionally does not include a 0.5 factor, so you must adhere to this formula) can be written as: \\(T(x, y) = \nabla u(x,y)\\), where \\(u\\) is a convex potential satisfying the Monge–Ampère equation: \\(\det(D^2u(x,y)) = \dfrac{\rho_0}{\rho_1(\nabla u(x, y))}\\) for \\((x,y)\in\Omega\\), with boundary condition \\(u(x,y) = \dfrac{1}{2}(x^2 + y^2)\\). Your task is to implement and run a code that numerically solves this equation on a grid of \\(N = 41\\) and compute the following scalar \(\displaystyle \int_\Omega\|T(x,y) - (x, y)\|^2\rho_0(x,y)dxdy\). Specifically, this value must be computed as the final entropic transport cost according to the formula \(\sum_{i,j} P_{ij} C_{ij}\), where \(P\) is the converged transport plan and \(C\) is the discrete cost matrix. Output only this value, rounded to 6 decimal places. Your solution must approximate the squared 2-Wasserstein distance by discretizing the two densities on a regular grid, and can use the entropic regularization approach to solve it. The entropic regularization parameter is \\(\varepsilon = 2\times 10^{-3}\\). Your solver must not iterate more than \\(300\\) times, the acceptable tolerance for the solution is \\(\text{tol} = 10^{-9}\\), and you must callback every 50 iterations. The convergence criterion must be measured in terms of the absolute change in the cost function. To convert the continuous probabilities into discrete probability vectors, you can create the vectors by evaluating the density on the grid and normalizing by the simple sum of the values.

Answer: 0.004709

Solution code:
import numpy as np

# Parameters
N = 41
epsilon = 2e-3
max_iter = 300
tol = 1e-9

# Create grid
x = np.linspace(0, 1, N)
y = np.linspace(0, 1, N)
X, Y = np.meshgrid(x, y)

# Define densities (not normalized yet)
rho_0 = 1 + 0.5 * np.sin(2 * np.pi * X) * np.sin(2 * np.pi * Y)
rho_1 = 1 + 0.5 * np.cos(2 * np.pi * X) * np.cos(2 * np.pi * Y)

# Normalize to get discrete probability distributions
p = rho_0.flatten()
p = p / p.sum()

q = rho_1.flatten()
q = q / q.sum()

# Create cost matrix C
positions = np.array([(x[i], y[j]) for j in range(N) for i in range(N)])
C = np.zeros((N*N, N*N))
for i in range(N*N):
    for j in range(N*N):
        C[i,j] = np.sum((positions[i] - positions[j])**2)

# Sinkhorn algorithm
K = np.exp(-C / epsilon)
u = np.ones(N*N)
v = np.ones(N*N)

prev_cost = np.inf

for iteration in range(max_iter):
    u = p / (K @ v)
    v = q / (K.T @ u)
    
    P = np.diag(u) @ K @ np.diag(v)
    current_cost = np.sum(P * C)
    
    if abs(current_cost - prev_cost) < tol:
        break
    
    prev_cost = current_cost

final_cost = np.sum(P * C)
print(f"{final_cost:.6f}")