-
Notifications
You must be signed in to change notification settings - Fork 0
/
rectangle_diagram.py
92 lines (74 loc) · 2.47 KB
/
rectangle_diagram.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
import thompson
from fractions import Fraction
from collections import deque
def vertex_placement(tree: thompson.Vertex):
vertices = deque()
leaves = []
for v in tree.in_order_vertices():
if v.is_leaf():
leaves.append(v)
else:
vertices.append(v)
placements = {L: (Fraction(i), Fraction(0)) for (i, L) in enumerate(leaves)}
while vertices:
v = vertices.popleft()
try:
lx, ly = placements[v.left]
rx, ry = placements[v.right]
except KeyError:
vertices.append(v)
continue
else:
x = (ry - ly + lx + rx) / 2
y = ly + x - lx
placements[v] = (x, y)
return placements, leaves
def get_segments(word1, word2):
tree1, tree2, label_to_vertex = \
thompson.construct_trees(word1, word2)
placements, leaves = vertex_placement(tree1)
placements2, _ = vertex_placement(tree2)
bottom_placements = {v: (x, -y)
for v, (x, y)
in placements2.items()}
placements.update(bottom_placements)
segments = []
for tree in tree1, tree2:
for v in tree.in_order_vertices():
if v.is_leaf():
continue
segments += [
(placements[v], placements[v.left]),
(placements[v], placements[v.right]),
]
return segments, [placements[x] for x in leaves]
def plot_rectangle_diagram(word1, word2=None):
if word2 is None:
word1, word2 = word1.split(' / ')
import matplotlib.pyplot as plt
from matplotlib import collections as mc
segs, leaves = get_segments(word1, word2)
segs = mc.LineCollection(segs, colors='k', linewidths=1)
fig = plt.figure()
ax1 = fig.add_subplot(1, 1, 1)
ax1.add_collection(segs)
for L in leaves:
ax1.plot(*L, 'r.')
ax1.autoscale()
fig.show()
def plot_all(lines, rows=4, cols=4):
import matplotlib.pyplot as plt
from matplotlib import collections as mc
if '\n' in lines:
lines = lines.splitlines()
fig = plt.figure()
for i, line in enumerate(lines, start=1):
word1, word2 = line.split(" / ")
segs, leaves = get_segments(word1, word2)
segs = mc.LineCollection(segs, colors='k', linewidths=1)
ax1 = fig.add_subplot(rows, cols, i)
ax1.add_collection(segs)
for L in leaves:
ax1.plot(*L, 'r.')
ax1.autoscale()
fig.show()