example_plot.py

# -*- coding: utf-8 -*-
"""
Created on Thu Jun 29 15:13:12 2023

@author: SomgBird
"""

import torch
from pathlib import Path
import networkx as nx
import torch_geometric
import matplotlib.pyplot as plt


from models import ChebNetStacked
from data_generator import GNNDataset


path = Path('K:/Dev/PyTorch projects/data/rusautocon2023_data/w3-10_n1.0_s30/')
torch.set_printoptions(sci_mode=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = GNNDataset(root=path/'data', transform=None)


in_channels=1
hidden_channels=128
out_channels=2
K=6
cheb_depth=4
linear_depth=1
linear_channels=512

model = ChebNetStacked(
    in_channels=in_channels,
    hidden_channels=hidden_channels,
    out_channels=out_channels,
    K=K,
    cheb_depth=cheb_depth,
    linear_depth=linear_depth,
    linear_channels=linear_channels,
    normalization=None, 
    bias=True,
    droprate=0.3).to(device)

model.load_state_dict(torch.load(path/'models/double_chebfc_hc128_hl4_K6_lc512_ll1.pth'))
model.eval()

data = dataset[450].to(device)
y = model(
    data.x,
    data.edge_index,
    data.edge_weight,
    None,
    4)
pred = y.argmax(dim=1)
data.output = y
data.pred = pred

G = torch_geometric.utils.to_networkx(data, node_attrs=['x', 'output', 'pred'], edge_attrs=['edge_weight'], to_undirected=True)

pos_nodes=nx.spring_layout(G, seed=131232)
pos_labels = {}

for k, v in pos_nodes.items():
    pos_labels[k] = (v[0]+0.01, v[1]+0.1)


x_values, y_values = zip(*pos_labels.values())
x_max = max(x_values)
x_min = min(x_values)
x_margin = (x_max - x_min) * 0.4
x_min = x_min - x_margin/7
x_max = x_max + x_margin

y_max = max(y_values)
y_min = min(y_values)
y_margin = (y_max - y_min) * 0.1
y_min = y_min - y_margin
y_max = y_max + y_margin


def draw(G, pos_nodes, node_labels, node_colors, node_size):
    fig, ax = plt.subplots(figsize=(7,7))
    nx.draw_networkx(G, pos_nodes, with_labels=False, node_size=node_size, ax=ax, node_color=node_colors, cmap=plt.cm.cool)
    nx.draw_networkx_labels(G, pos_labels, labels=node_labels, horizontalalignment='left', ax=ax, font_size=20)
    plt.box(False)
    plt.xlim(x_min, x_max)
    plt.ylim(y_min, y_max)
    plt.show()


# Input data    
node_labels = nx.get_node_attributes(G, 'x')
node_colors = []

for key in node_labels.keys():
    node_colors.append(node_labels[key])
    if node_labels[key] == -1.0:
        node_labels[key] = "-1"
    else:
        node_labels[key] = f"{node_labels[key]:.4f}"

draw(G, pos_nodes, node_labels, node_colors, 300)

# Output data
node_labels = nx.get_node_attributes(G, 'output')
node_colors = []

for key in node_labels.keys():
    node_colors.append(node_labels[key][0])
    node_labels[key] = f"[{node_labels[key][0]:.4f},\n {node_labels[key][1]:.4f}]"


draw(G, pos_nodes, node_labels, node_colors, 300)

# Postprocessed
node_labels = nx.get_node_attributes(G, 'pred')
node_colors = []

for key in node_labels.keys():
    node_colors.append(1/(1 + node_labels[key]))
    node_labels[key] = f"{node_labels[key]:.0f}"


draw(G, pos_nodes, node_labels, node_colors, 300)