非負行列の行和，列和を揃える操作を行う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."
r = r / np.sum(r)
c = c / np.sum(c)
v = np.copy(r)
u = np.copy(c)
cnt = 0
while True:
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
