dmlc--dgl
KounianhuaDu
9aac93ff21
* gnn-explainer * gnn-explainer * gnn-explainer * gnn-explainer * fix * fix * fix * readme * readme Co-authored-by: zhjwy9343 <6593865@qq.com>
103 行
3.4 KiB
Python
103 行
3.4 KiB
Python
# Utility file for graph queries
|
|
|
|
import tkinter
|
|
import matplotlib
|
|
matplotlib.use('TkAgg')
|
|
|
|
import networkx as nx
|
|
import matplotlib.pylab as plt
|
|
import torch as th
|
|
import dgl
|
|
from dgl.sampling import sample_neighbors
|
|
|
|
def extract_subgraph(graph, seed_nodes, hops=2):
|
|
"""
|
|
For the explainability, extract the subgraph of a seed node with the hops specified.
|
|
|
|
Parameters
|
|
----------
|
|
graph: DGLGraph, the full graph to extract from. This time, assume it is a homograph
|
|
seed_nodes: Tensor, index of a node in the graph
|
|
hops: Integer, the number of hops to extract
|
|
|
|
Returns
|
|
-------
|
|
sub_graph: DGLGraph, a sub graph
|
|
origin_nodes: List, list of node ids in the origin graph, sorted from small to large, whose order is the new id. e.g
|
|
[2, 51, 53, 79] means in the new sug_graph, their new node id is [0,1,2,3], the mapping is 2<>0, 51<>1, 53<>2,
|
|
and 79 <> 3.
|
|
new_seed_node: Scalar, the node index of seed_nodes
|
|
"""
|
|
seeds=seed_nodes
|
|
for i in range(hops):
|
|
i_hop = sample_neighbors(graph, seeds, -1)
|
|
seeds = th.cat([seeds, i_hop.edges()[0]])
|
|
|
|
ori_src, ori_dst = i_hop.edges()
|
|
edge_all = th.cat([ori_src, ori_dst])
|
|
origin_nodes, new_edges_all = th.unique(edge_all, return_inverse=True)
|
|
|
|
n = int(new_edges_all.shape[0] / 2)
|
|
new_src = new_edges_all[:n]
|
|
new_dst = new_edges_all[n:]
|
|
|
|
sub_graph = dgl.DGLGraph((new_src, new_dst))
|
|
new_seed_node = th.nonzero(origin_nodes==seed_nodes, as_tuple=True)[0][0]
|
|
|
|
return sub_graph, origin_nodes, new_seed_node
|
|
|
|
|
|
def visualize_sub_graph(sub_graph, edge_weights=None, origin_nodes=None, center_node=None):
|
|
"""
|
|
Use networkx to visualize the sub_graph and,
|
|
if edge weights are given, set edges with different fading of blue.
|
|
|
|
Parameters
|
|
----------
|
|
sub_graph: DGLGraph, the sub_graph to be visualized.
|
|
edge_weights: Tensor, the same number of edges. Values are (0,1), default is None
|
|
origin_nodes: List, list of node ids that will be used to replace the node ids in the subgraph in visualization
|
|
center_node: Tensor, the node id in origin node list to be highlighted with different color
|
|
|
|
Returns
|
|
show the sub_graph
|
|
-------
|
|
|
|
"""
|
|
# Extract original idx and map to the new networkx graph
|
|
# Convert to networkx graph
|
|
g = dgl.to_networkx(sub_graph)
|
|
nx_edges = g.edges(data=True)
|
|
|
|
if not (origin_nodes is None):
|
|
n_mapping = {new_id: old_id for new_id, old_id in enumerate(origin_nodes.tolist())}
|
|
g = nx.relabel_nodes(g, mapping=n_mapping)
|
|
|
|
pos = nx.spring_layout(g)
|
|
|
|
if edge_weights is None:
|
|
options = {"node_size": 1000,
|
|
"alpha": 0.9,
|
|
"font_size":24,
|
|
"width": 4,
|
|
}
|
|
else:
|
|
|
|
ec = [edge_weights[e[2]['id']][0] for e in nx_edges]
|
|
options = {"node_size": 1000,
|
|
"alpha": 0.3,
|
|
"font_size": 12,
|
|
"edge_color": ec,
|
|
"width": 4,
|
|
"edge_cmap": plt.cm.Reds,
|
|
"edge_vmin": 0,
|
|
"edge_vmax": 1,
|
|
"connectionstyle":"arc3,rad=0.1"}
|
|
|
|
nx.draw(g, pos, with_labels=True, node_color='b', **options)
|
|
if not (center_node is None):
|
|
nx.draw(g, pos, nodelist=center_node.tolist(), with_labels=True, node_color='r', **options)
|
|
|
|
plt.show()
|
|
|