-
Notifications
You must be signed in to change notification settings - Fork 1
/
plot.py
111 lines (85 loc) · 3.26 KB
/
plot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
from PIL import Image, ImageDraw, ImageFont
import math
from bitset import BitSet, index
from graphviz import Graph
def plot(graph, engine='dot', filename='output/test', vertex_names={}):
"""
Possible engines: dot, neato, fdp, sfdp, twopi, circo
Vertex_names is an optional dict from vertices to strings.
"""
g = Graph(format='png', engine=engine)
def vertex_to_string(v):
return (vertex_names[v] if v in vertex_names else '') + ' ({})'.format(str(index(v)))
for v in graph:
g.node(vertex_to_string(v))
for v, w in graph.edges:
g.edge(vertex_to_string(v), vertex_to_string(w))
g.render(filename)
def plot_circle(graph, im=None, color=128):
"""Draw given graph on given draw object."""
if im == None:
size = (512, 512)
im = Image.new('RGB', size, 'white')
else:
size = im.size
margin = 30
draw = ImageDraw.Draw(im)
radius = min(size) / 2 - margin
vertexcoords = {}
nr_vertices = len(graph.vertices)
for i, v in enumerate(sorted(graph.vertices, key=lambda v: v.identifier)):
r = i * 2 * math.pi / nr_vertices
vertexcoords[BitSet(v)] = (radius * (1 + math.cos(r)) + margin,
radius * (1 + math.sin(r)) + margin)
for b, v in graph.vertices.items():
coords = vertexcoords[b]
draw_vertex(draw, graph.vertices[b], coords, color)
for w in v.neighbors:
draw.line([vertexcoords[b], vertexcoords[w]],
fill=color, width=1)
return im
def plot_bipartite(graph, im=None, color=128):
"""Draw given graph on given draw object."""
if im == None:
size = (512, 512)
im = Image.new('RGB', size, 'white')
else:
size = im.size
margin = 30
draw = ImageDraw.Draw(im)
# For each vertex we assign a vertical interval
# The vertex is drawn in the center of that interval
xcoord = margin
vertexcoords1 = {}
interval_length = size[1] / len(graph.group1)
for i, v in enumerate(graph.group1):
ycoord = i * interval_length + interval_length / 2
vertexcoords1[v] = (xcoord, ycoord)
vertexcoords2 = {}
xcoord = size[0] - margin
interval_length = size[1] / len(graph.group2)
for i, v in enumerate(graph.group2):
ycoord = i * interval_length + interval_length / 2
vertexcoords2[v] = (xcoord, ycoord)
for v in graph.group1:
coords = vertexcoords1[v]
draw_vertex(draw, v, coords, color)
for v in graph.group2:
coords = vertexcoords2[v]
draw_vertex(draw, v, coords, color)
for v in graph.group1:
for w in graph(v):
try:
draw.line([vertexcoords1[v], vertexcoords2[w]], fill=color)
except KeyError:
draw.line([vertexcoords1[w], vertexcoords2[v]], fill=color)
return im
def draw_vertex(draw, vertex, coords, color=128):
"""Draw a single vertex."""
radius = 5
uppercoords = (coords[0] - radius, coords[1] - radius)
lowercoords = (coords[0] + radius, coords[1] + radius)
box = [uppercoords, lowercoords]
draw.ellipse(box, fill=color)
font = ImageFont.truetype("resources/FreeMono.ttf", 30)
draw.text(lowercoords, str(vertex.identifier), fill=color, font=font)