pydot + Graphviz で相関行列を可視化するサンプル

相関行列 (のような対称行列) を pydot + Graphviz で可視化する機会があったので、そのときにやったことのメモ。

サンプルデータ

今回は scikit-learn に付属している「ボストン市の住宅価格」データセットを利用する。

from sklearn.datasets import load_boston
import pandas as pd

boston = load_boston()
df = pd.DataFrame(boston.data, columns=boston.feature_names)
df.head()

相関行列を求める

Pandas DataFrame は corr メソッドで相関行列を求められる。

df.corr()

ヒートマップを描画する

グラフを描画する前に、相関行列を可視化する際によくやるヒートマップを描画してみる。

import seaborn as sns

sns.heatmap(df.corr(), center=0)

これでもいいが、すこしわかりづらい気がする。

グラフを描画する

pydot + Graphviz を使ってグラフを描画してみる。

import itertools
import pydot

def create_graph(df, threshold=0.5):
    """相関係数行列をグラフに変換する関数"""
    assert set(df.index) == set(df.columns)

    # グラフを作成
    graph = pydot.Dot(graph_type='graph')  # graph = 無向グラフ

    # ノードを追加
    for col in df.columns:
        node = pydot.Node(f'"{col}"', label=col)
        graph.add_node(node)

    # エッジを追加
    for a, b in itertools.combinations(df.columns, 2):
        val = df.loc[a, b]
        # 相関係数の絶対値が threshold 未満の場合はエッジを追加しない
        if abs(val) < threshold:
            continue
        edge = pydot.Edge(graph.get_node(f'"{a}"')[0], graph.get_node(f'"{b}"')[0])
        edge.set_label(f'{val:.2f}')
        edge.set_penwidth(4 * (abs(val) - threshold) / (1 - threshold))  # 太さをいい感じにする
        h, s, v = 0.0 if val > 0 else 2 / 3, abs(val), 1.0  # 色を設定する
        edge.set_color(' '.join([str(a) for a in (h, s, v)]))
        graph.add_edge(edge)

    return graph

相関行列は対称行列で、かつ対角成分を可視化する必要がないので、 itertools.combinations(df.columns, 2) でループを回している。

ノードを扱う際にカラム名をダブルクオートで囲んでいる理由については こちら を参照。

from IPython.display import Image

graph = create_graph(df.corr())
display(Image(graph.create(format='png')))

正の相関が赤、負の相関が青、相関の強さはエッジの太さと色の濃さで表現するようにした。

場合によってはこのような可視化の手法もわかりやすいのではないかと思う。