HardCalculusca-225

Sinkhorn One Step

Problem

One Sinkhorn iteration alternates scaling to match the marginals. With K = exp(-cost / eps) and initial row scaling u0:

v_j = c_j / sum_i(u0_i * K_ij), then u_i = r_i / sum_j(K_ij * v_j)

The transport plan is P_ij = u_i * K_ij * v_j. Return the plan for one full iteration.

Hint

Entropic regularization turns optimal transport into alternating diagonal scalings.

case 1

in: [[[1,2],[3,4]],[0.5,0.5],[0.5,0.5],1,[1,1]]

expected: [[0.25,0.25],[0.25,0.25]]

case 2

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

expected: [[0.44039853898894127,0.05960146101105879],[0.05960146101105879,0.44039853898894127]]

case 3

in: [[[1]],[1],[1],1,[1]]

expected: [[1]]

case 4

in: [[[1,2],[3,4]],[0.4,0.6],[0.3,0.7],2,[0.5,0.5]]

expected: [[0.12000000000000001,0.27999999999999997],[0.18,0.42000000000000004]]

Starter codePython

def sinkhorn_one_step(cost, r, c, eps, u0):
    # Your code here
    pass

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