Wasserstein距離

[WGAN](https://arxiv.org/abs/1701.07875) などで使われるWasserstein距離は確率分布の間の距離のひとつです。本稿では、離散型確率分布に対するWasserstein距離の定義と、Pythonによる計算の例を示します。

## Wasserstein距離の定義
ふたつのカテゴリカル分布$a$、$b$それぞれの次元を$m$、$n$とします。このふたつの分布の各次元の間で、確率の移動のコスト行列 $C \in \mathbb{R}^{n \times m}$ が与えられたとき、
ふたつの分布を周辺分布にもつような結合確率分布の空間

```math
U(a,b) := \{P \in \mathbb{R}^{n \times m} : P \mathbb{1}_m = a \ {\rm and} \ P^T \mathbb{1}_n = b\}
```

の中で、最小の総移動コストを求める問題

```math
L_C(a,b) := \min_{P\in U(a,b)} \langle C, P\rangle
```

をKantorovichの最適輸送問題と呼びます。ここで、 $\mathbb{1}_m$ は、$m$次元のすべての要素が$1$のベクトル、$ \langle C, P\rangle:= \sum\_{i,j} C\_{i,j} P\_{i,j}  $ です。 

コスト行列$C$が、ある距離行列$D$、パラメーター$p>1$に対して、$C = D^p$で与えられたとき、分布$a$、$b$間のp-Wassserstein距離$W_p(a,b)$は、次にように定義されます。

```math
W_p(a,b) := L_{D^p}(a,b)^{1/p}
```

特に、$p=1$のとき、Earth movers距離と呼びます。

## Wasserstein距離の計算
Wasserstein距離の計算、すなわちKantrovichの最適輸送問題は次のように線形計画問題として解くことができます。

まず、$p_i \in R^n$、$c_i \in R^n$で

```math
P = \left(\begin{array}{cccc} p_1 & p_2 & \cdots & p_m \end{array}\right)
```

```math
C = \left(\begin{array}{cccc} c_1 & c_2 & \cdots & c_m \end{array}\right)
```

と表すと、コスト関数は、

```math
L_C(a,b) := \min_{P\in U(a,b)} \left(\begin{array}{cccc}{}^tc_1 &{}^tc_2 &\cdots & {}^tc_m \end{array}\right)  \left(\begin{array}{c}p_1 \\ p_2 \\ \vdots \\ p_m\end{array}\right)
```

と書き直せます。また、$\phi_{i} \in \mathbb{R}^n$を$i$番目の要素が$1$、それ以外の要素が$0$のベクトルとすると、結合確率についての制約は

```math
\begin{array}{c}
\left(\begin{array}{ccc}
\phi_{1} & \cdots  & \phi_1 \\
\phi_2 & \cdots  & \phi_2 \\
\vdots & & \vdots \\
\phi_n & \cdots  & \phi_n \\
\end{array}\right)\\
\left(\begin{array}{ccc}
\mathbb{1}_n & 0 & \cdots & 0\\
0 & \mathbb{1}_n  & \cdots & 0\\
\vdots & & & \vdots \\
0  & \cdots & 0 & \mathbb{1}_n \\
\end{array}\right)
\end{array} 
\left(\begin{array}{c} p_1 \\ p_2 \\ \vdots \\ p_m\end{array}\right) 
=\left(\begin{array}{c}a \\ b\end{array}\right)
```
と書き直せます。
これで線形な制約下で線形なコスト関数を最小化する $p_i$ を求める線形計画問題として解くことができます。以下でPythonを使った計算例を示します。


```python
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
```

確率分布 $a$、$b$

```python
n = 10
np.random.seed(0)
a = np.abs(np.random.randn(n))
a = a / np.sum(a)
b = np.abs(np.random.randn(n))
b = b / np.sum(b)

plt.figure()
plt.bar(np.arange(n),a, width=1.0)
plt.figure()
plt.bar(np.arange(n),b, width=1.0)
```
![a.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/1639e7a6-497f-bbf0-7490-db578aa43a92.png)
![b.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/a56a59f1-2637-b122-4c12-f40a2ecbdee5.png)



結合確率についての制約式

```python
cond_b = np.concatenate((a,b),axis=0)
cond_a = np.zeros((2*n,n*n))
for i in range(n):
    ind_vector_a = np.zeros(n)
    ind_vector_a[i] = 1
    cond_a[i,:] = np.tile(ind_vector_a, n)
    ind_vector_b = np.ones(n)
    cond_a[n+i,(n*i):(n*(i+1))] =  ind_vector_b
```

距離行列 $D$、コスト行列$C$

```python
from matplotlib import cm
d_mat = np.zeros([n,n])
for i in range(n):
    d_mat[i,(i+1):n] = np.arange(1,n-(i))
i_lower = np.tril_indices(n, -1)
d_mat[i_lower] = d_mat.T[i_lower]
p = 1
c_mat = d_mat**p
c_mat_reshape = np.squeeze(np.asarray(c_mat.reshape((1, n*n))))

plt.imshow(c_mat, cmap=cm.gist_heat)
plt.axis('off')
plt.colorbar()
```
![c.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/d3161582-fd93-b134-caa7-a010f5dc8fc6.png)



線形計画問題を解く

```python
from scipy.optimize import linprog
sol = linprog(c_mat_reshape, A_eq=cond_a, b_eq=cond_b)
```

1-Wasserstein距離は

```python
sol.fun # -> 1.3601368298912992!
```

結合確率分布は

```python
from matplotlib import cm
p_mat = np.reshape(sol.x, (n,n))
plt.imshow(p_mat, cmap=cm.gist_heat)
plt.axis('off')
plt.colorbar()
```
![p.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/78c2c0c4-4a97-9274-d82d-c3abf233d1fb.png)


結合確率分布の周辺分布がもとの確率分布$a$、$b$と一致していることの確認

```python
p_mat_margin_col = np.sum(p_mat, axis=1)
p_mat_margin_row = np.sum(p_mat, axis=0)


plt.figure()
plt.bar(np.arange(n),a, width=1.0, label='$a$')
plt.bar(np.arange(n), p_mat_margin_row, label='$P\mathbb{1}$')
plt.legend()
plt.figure()
plt.bar(np.arange(n),b, width=1.0, label='$b$')
plt.bar(np.arange(n), p_mat_margin_col, label='$P^T \mathbb{1}$')
plt.legend()
```
![a_result.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/44f69ac5-292c-24a9-6182-5897dc470a71.png)
![b_result.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/35175/eb068183-e844-8b2f-8b5f-86bc314274ca.png)

## 参考文献
* [Gabriel Peyré and Marco Cuturi, Computational Optimal Transport, ArXiv:1803.00567, 2018.](https://arxiv.org/abs/1803.00567)
* https://en.wikipedia.org/wiki/Wasserstein_metric
* https://vincentherrmann.github.io/blog/wasserstein/