Python Sinkhorn Knopp アルゴリズムを実装した．

非負行列の行和，列和を揃える操作を行うSinkhorn Knoppアルゴリズムをpythonで実装しました．Julia版はこっち． 入力行列Aのi行目の行和がr[i], j列目の列和がc[j]になるようにするアルゴリズムが以下です．

Sinkhorn Knoppアルゴリズムは，エントロピー拘束付き最適化と関りがあります． 最適輸送の文脈での最重要アルゴリズムの１つです．

import numpy as np

def main(A, r, c, eps = 1.0E-7):
    """
    Input  : Matrix A, vectors r and c
             A can be a rectangular matrix
    Output : P = diag(u) A diag(v)
    Sum_i p_1i = c_i / sum(c) 
    Sum_j p_j1 = r_j / sum(r) 
    <=> sum(P.T) = c
        sum(P) = r 
    """

    assert A.ndim == 2, "A is not matrix"
    assert r.ndim == 1 and c.ndim == 1, "r and c should be vectors."
    assert A.shape[0] == c.shape[0], "dim of A does not correspond to r ones."
    assert A.shape[1] == r.shape[0], "dim of A does not correspond to c ones." 

    # vector r and c should be normalized.
    r = r / np.sum(r)
    c = c / np.sum(c)
    
    # Initial Step
    v = np.copy(r) # same as v = r / ( A.T @ c)   
    u = np.copy(c) # same as u = c / ( A @ v)

    cnt = 0
    while True:
        # Sinkhorn-knopp Step
        v_new = r / ( A.T @ u)   
        u_new = c / ( A @ v_new)

        if quick_convergence_check(v, v_new, u, u_new, eps) == True:
            break

        if cnt > 25000:
            print("calculation did not converg")
            exit()
        v = v_new
        u = u_new
        cnt += 1

    P = np.diag(u) @ A @ np.diag(v)
    return P

def quick_convergence_check(v, v_new, u, u_new, eps):
    if np.linalg.norm(abs(v - v_new)) < eps:
        if np.linalg.norm(abs(u - u_new)) < eps:
            return True 
    return False 

if __name__ == "__main__":
    A = np.array([[1,5,3,6],[9,5,5,6],[1,7,10,9],[2,3,1,4]], dtype = "float64")
    c = np.array([7,7,4,2], dtype="float64")
    r = np.array([1,2,3,4], dtype="float64")

    P = main(A, r, c)
    print(sum(P.T))
    print(sum(P))

ちなみに，rとcを全て１からなる行列にすると求まる行列は二重確率行列と呼ばれるのですが，二重確率行列のいろいろな性質についてはこの論文が参考になります．

www.sciencedirect.com