networkxでグラフ信号の可視化 カラーバー付き

はじめに

グラフのノード信号を可視化するとき，カラーバーがあると便利です．
ググっても良い情報が見当たらなかったので，備忘録的に書きます．

方法

通常はnx.draw()でグラフをプロットすることが多いですが，今回は matplotlib.collections.PathCollection オブジェクトをもらうため，意図的に nx.draw_networkx_nodes() を用います．
このオブジェクトを plt.colorbar() に渡すことで，カラーバーを生成します．

import networkx as nx
import numpy as np
import matplotlib.pyplot as plt

G = nx.balanced_tree(r=3, h=2)
pos = nx.spring_layout(G, seed=0)
signal = np.linspace(0, 1, num=G.number_of_nodes())

nodes = nx.draw_networkx_nodes(
    G,
    pos,
    node_color=signal,
    cmap=plt.cm.viridis,
)
edges = nx.draw_networkx_edges(G, pos) # edge も可視化

ax = plt.gca()
ax.set_axis_off()
plt.colorbar(nodes, label='Signal Value')
plt.tight_layout()
plt.show()

上記のコードを実行すると，下図が得られます．

なお，edgeの色をカラーバーで変更したい場合，networkxのドキュメント を参考にしてください．nx.draw_networkx_nodes() と nx.draw_networkx_edges() の返り値が異なるクラスのオブジェクトなので，単にplt.colorbar(edges) と変更するだけでは動作しません．

参考文献

• edgeの色をカラーバーで出力している公式ドキュメント https://networkx.org/documentation/stable/auto_examples/drawing/plot_directed.html
• https://stackoverflow.com/questions/26739248/how-to-add-a-simple-colorbar-to-a-network-graph-plot-in-python
• https://stackoverflow.com/questions/40680910/create-own-colorbar-for-networkx-plot