eaglelandsonce commited on
Commit
f1aa2ca
1 Parent(s): 12440bd

Create scripts/viz.py

Browse files
Files changed (1) hide show
  1. scripts/viz.py +119 -0
scripts/viz.py ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from IPython.display import IFrame
2
+ import json
3
+ import uuid
4
+
5
+ def vis_network(nodes, edges, physics=False):
6
+ html = """
7
+ <html>
8
+ <head>
9
+ <script type="text/javascript" src="../lib/vis/dist/vis.js"></script>
10
+ <link href="../lib/vis/dist/vis.css" rel="stylesheet" type="text/css">
11
+ </head>
12
+ <body>
13
+
14
+ <div id="{id}"></div>
15
+
16
+ <script type="text/javascript">
17
+ var nodes = {nodes};
18
+ var edges = {edges};
19
+
20
+ var container = document.getElementById("{id}");
21
+
22
+ var data = {{
23
+ nodes: nodes,
24
+ edges: edges
25
+ }};
26
+
27
+ var options = {{
28
+ nodes: {{
29
+ shape: 'dot',
30
+ size: 25,
31
+ font: {{
32
+ size: 14
33
+ }}
34
+ }},
35
+ edges: {{
36
+ font: {{
37
+ size: 14,
38
+ align: 'middle'
39
+ }},
40
+ color: 'gray',
41
+ arrows: {{
42
+ to: {{enabled: true, scaleFactor: 0.5}}
43
+ }},
44
+ smooth: {{enabled: false}}
45
+ }},
46
+ physics: {{
47
+ enabled: {physics}
48
+ }}
49
+ }};
50
+
51
+ var network = new vis.Network(container, data, options);
52
+
53
+ </script>
54
+ </body>
55
+ </html>
56
+ """
57
+
58
+ unique_id = str(uuid.uuid4())
59
+ html = html.format(id=unique_id, nodes=json.dumps(nodes), edges=json.dumps(edges), physics=json.dumps(physics))
60
+
61
+ filename = "figure/graph-{}.html".format(unique_id)
62
+
63
+ file = open(filename, "w")
64
+ file.write(html)
65
+ file.close()
66
+
67
+ return IFrame(filename, width="100%", height="400")
68
+
69
+ def draw(graph, options, physics=False, limit=100):
70
+ # The options argument should be a dictionary of node labels and property keys; it determines which property
71
+ # is displayed for the node label. For example, in the movie graph, options = {"Movie": "title", "Person": "name"}.
72
+ # Omitting a node label from the options dict will leave the node unlabeled in the visualization.
73
+ # Setting physics = True makes the nodes bounce around when you touch them!
74
+ query = """
75
+ MATCH (n)
76
+ WITH n, rand() AS random
77
+ ORDER BY random
78
+ LIMIT {limit}
79
+ OPTIONAL MATCH (n)-[r]->(m)
80
+ RETURN n AS source_node,
81
+ id(n) AS source_id,
82
+ r,
83
+ m AS target_node,
84
+ id(m) AS target_id
85
+ """
86
+
87
+ data = graph.run(query, limit=limit)
88
+
89
+ nodes = []
90
+ edges = []
91
+
92
+ def get_vis_info(node, id):
93
+ node_label = list(node.labels())[0]
94
+ prop_key = options.get(node_label)
95
+ vis_label = node.properties.get(prop_key, "")
96
+
97
+ return {"id": id, "label": vis_label, "group": node_label, "title": repr(node.properties)}
98
+
99
+ for row in data:
100
+ source_node = row[0]
101
+ source_id = row[1]
102
+ rel = row[2]
103
+ target_node = row[3]
104
+ target_id = row[4]
105
+
106
+ source_info = get_vis_info(source_node, source_id)
107
+
108
+ if source_info not in nodes:
109
+ nodes.append(source_info)
110
+
111
+ if rel is not None:
112
+ target_info = get_vis_info(target_node, target_id)
113
+
114
+ if target_info not in nodes:
115
+ nodes.append(target_info)
116
+
117
+ edges.append({"from": source_info["id"], "to": target_info["id"], "label": rel.type()})
118
+
119
+ return vis_network(nodes, edges, physics=physics)