그래프 예쁘게 그리기 2탄 : networkx.draw 함수

1탄에서 networkx layout 함수를 사용했다. 

그래프 예쁘게 그리기! networkx layout 함수

node classification 후에 node coloring을 하려고 하는데 제일 간단하게 할 수 있는 방법이 뭘까 찾아봤다. 
현재 karateclub data라서, 0과 33노드 기준으로 2개의 클래스로 나누는 문제이다. 

다음 코드처럼 딱 하나가 추가됐다. 현재 node_labels 는 prediction 결과로 torch.Size([34]) , x 는 torch.Size([34, 34]) 의 텐서이다. 

prediction 은 각 노드의 클래스가 적혀있고, x는 diagonal matrix의 모습이다. 이 상태로 그래프를 그릴 때, nx.draw 함수를 사용하였다. 

nx.draw 함수 
Parameters : 
G : graph
pos
ax : Matplotlib Axes object
node_size : scalar or array
node_color
font_size
등등... option이 다양하다. 자세한 내용 참고

def plot_graph(data, node_labels=None):
    fig, ax = plt.subplots()
    fig.set_tight_layout(False)
    nx_g = to_networkx(data=data).to_undirected()
    pos = nx.kamada_kawai_layout(nx_g)
    
    if node_labels is not None:
        node_color = [[1, 1, 0.5] if x else [1,0.5,0.5] for x in node_labels]
   
    nx.draw(nx_g, pos, with_labels=True, node_size = 200, node_color = node_color, font_size = 8)
    
    plt.show()

for ~ if ~ else 문을 한줄로 작성하기 

a   if   i>5   else   b   for  i   in list

for i in list : 
   if i>5 :
     list[i]=a
   else :
     list[i]=b

그래프 결과