numpyで相互情報量(連続値)の計算

モチベーション

連続な確率変数$X$と$Y$の相互情報量$I(X;Y)$をPythonで計算したい．
$$
I(X;Y) = \int_Y \int_X p(x, y) \log \frac{p(x,y)}{p(x)p(y)} dx dy
$$

コード

import numpy

def mutual_information(X, Y, bins=10):
    # 同時確率分布p(x,y)の計算
    p_xy, xedges, yedges = np.histogram2d(X, Y, bins=bins, density=True)
    
    # p(x)p(y)の計算
    p_x, _ = np.histogram(X, bins=xedges, density=True)
    p_y, _ = np.histogram(Y, bins=yedges, density=True)
    p_x_y = p_x[:, np.newaxis] * p_y
    
    # dx と dy
    dx = xedges[1] - xedges[0]
    dy = yedges[1] - yedges[0]
    
    # 積分の要素
    elem = p_xy * np.ma.log(p_xy / p_x_y)
    # 相互情報量とp(x, y), p(x)p(y)を出力
    return np.sum(elem * dx * dy), p_xy, p_x_y

ポイント

とりあえず相互情報量を計算したい場合は上の関数を使えばOKです．
ついでに実装の上で大事な点をいくつか残しておきます．

np.histogram2dのdensity

density=Trueにすると確率が返ってくると漠然と考えていたらnp.sum(p_xy)が1にならなくて少し焦りました．
注意すべきポイントはp_xyが確率ではなく確率密度という点です．

$X$と$Y$は連続変数なので，ヒストグラムで近似しているのは確率密度です．
したがってビンの幅も考慮して足し合わせると1になります．

np.histogramやnp.histogram2dは確率密度とビン(コード中のedges)を返します．
このビンからdxやdyを計算しておくことが必要になります．

import numpy as np

N = 1000
X = np.random.normal(loc=0, scale=1, size=N)

p_x, edges = np.histogram(X, bins=10, density=True)

# 何も考えずに確率密度の和を取った場合，当然1にはならない
print(np.sum(p_x))  # 出力例: 1.580769264599771

# ビン幅を考慮して和を取ると1になる
dx = edges[1] - edges[0]
print(np.sum(p_x * dx))  # 出力例: 1.0000000000000002

p_x_yの計算

コード中のp_x_yは$p(x)p(y)$を計算しようとしています．
実は最初に以下のコードで計算していて，うまくいきませんでした．

p_x_y = p_x * p_y

正しくは

p_x_y = p_x[:, np.newaxis] * p_y

です．前者はp_x_yが一次配列で，後者はp_x_yが二次配列になります．

実行例

実行例1 (２つのsin波)

両者は独立ではないため，$p(x, y)$と$p(x)p(y)$に差が生じ，相互情報量は大きくなる．

import matplotlib.pyplot as plt

# sin波とcos波
t = np.linspace(-5, 5, num=1000)
X = np.sin(2 * np.pi * t)
Y = np.cos(3 * np.pi * t)

# 相互情報量の計算
mi, p_xy, p_x_y = mutual_information(X, Y, bins=30)

# 結果の出力
plt.figure(dpi=100)
ax1 = plt.subplot(121)
ax2 = plt.subplot(122)
ax1.set_title(r'$P_{XY}(x, y)$')
ax1.imshow(p_xy)
ax2.set_title(r'$P_{X}(x) P_{Y}(y)$')
ax2.imshow(p_x_y)
plt.suptitle('MI = {}'.format(mi))
plt.show()

実行例2 (独立な正規分布)

２つの変数が独立な場合，$p(x, y)$と$p(x)p(y)$が一致し，相互情報量は小さくなる．

import matplotlib.pyplot as plt
# ２つの独立な正規分布
N = 10000
X = np.random.normal(size=N)
Y = np.random.normal(size=N)

# 相互情報量の計算
mi, p_xy, p_x_y = mutual_information(X, Y, bins=30)
実行例
# 結果の出力
plt.figure(dpi=100)
ax1 = plt.subplot(121)
ax2 = plt.subplot(122)
ax1.set_title(r'$P_{XY}(x, y)$')
ax1.imshow(p_xy)
ax2.set_title(r'$P_{X}(x) P_{Y}(y)$')
ax2.imshow(p_x_y)
plt.suptitle('MI = {}'.format(mi))
plt.show()