forked from aschweickart/CompBioSummer2015
/
DrawDTL.py
134 lines (124 loc) · 4.61 KB
/
DrawDTL.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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#DrawDTL.py
#July 2015
#Annalise and Carter
#File uses turtle graphics to render an image of the Total reconciliations graph
import turtle
import math
NodeLocations = {}
#tells a turtle to draw a line from start to end and returns to its position before the function call
def connect(Turtle,Start,End,rad):
"""Takes as input a turtle, starting coordinates, Start, and ending coordinates, End,
as well as the radius of the nodes, and connects the two locations with a line and arrow"""
Turtle.speed(0)
Turtle.pen(pencolor = "black")
loc = Turtle.pos()
hed = Turtle.heading()
movr = Start[0] < End[0]
movu = Start[1] < End[1]
lesserY = min(Start[1],End[1])
biggerY = max(Start[1],End[1])
lesserX = min(Start[0],End[0])
biggerX = max(Start[0],End[0])
theta = 0
#find angle of the arrow
try:
theta = (math.atan(float(biggerY - lesserY)/float(biggerX - lesserX)))
if not movr and movu:
theta = math.pi - theta
elif not movr and not movu:
theta = math.pi + theta
elif movr and not movu:
theta = (2 * math.pi) - theta
except:
theta = (math.pi/2)
if not movu:
theta += math.pi
#adjust start and ending points so as to not draw inside nodes
Start = (Start[0] + (rad * math.cos(theta)), Start[1])
Start = (Start[0], Start[1] + (rad * math.sin(theta)))
End = (End[0] - (rad * math.cos(theta)), End[1])
End = (End[0], End[1] - (rad * math.sin(theta)))
Turtle.penup()
Turtle.radians()
Turtle.seth(theta)
Turtle.setpos(Start,None)
Turtle.pendown()
Turtle.goto(End,None)
Turtle.degrees(360)
Turtle.stamp()
Turtle.penup()
Turtle.goto(loc,None)
Turtle.setheading(hed)
Turtle.pendown()
Turtle.ht()
def drawNodes(treeMin, eventDict, depth, nodeDict):
"""Takes as input
treeMin - a list of the starting nodes of the best reconciliations
eventDict - the DTL format dictionary
depth - a starting y-coordinate
nodeDict - a dictionary of nodes and their coordinates.
This function recursively draws the nodes of the DTL format dictionary, then
connects them using the connect function aboves"""
numTips = 0
for key in eventDict.keys():
if eventDict[key][0][0] == "C":
numTips+=1
width = numTips * 200
DISPLACE = width/2
dip = 15
if len(eventDict)<25:
radius = 30
else:
radius = 13
dip = 15
width = width/2
DISPLACE = DISPLACE/2
if treeMin == []:
for key in nodeDict:
for item in range(len(nodeDict[key][1:])):
connect(turtle.Turtle(), nodeDict[key][0], nodeDict[key][item+1], radius)
for thing in eventDict[key][item][1:-1]:
if thing !=(None, None):
connect(turtle.Turtle(), nodeDict[key][item+1], nodeDict[thing][0], radius)
return
difference = ((len(eventDict))*2)/numTips
numSols = len(treeMin)
turtle.speed(0)
turtle.pen(pencolor = "black")
eventList = []
newtreeMin = []
for x in range(len(treeMin)):
if not treeMin[x] in nodeDict:
nodeDict[treeMin[x]] = [((x+1)*width/(numSols + 1)-DISPLACE, depth + radius)]
turtle.penup()
turtle.setpos((x+1)*width/(numSols+1)-DISPLACE, depth)
turtle.pendown()
turtle.circle(radius)
turtle.left(130)
turtle.penup()
turtle.forward(radius)
turtle.pendown()
turtle.right(130)
turtle.write(treeMin[x], font = ("arial", 12, "normal"))
for y in eventDict[treeMin[x]]:
if type(y)== list:
eventList.append((y[0], treeMin[x]))
if y[1] !=(None, None) and not y[1] in newtreeMin:
newtreeMin.append(y[1])
if y[2] !=(None, None) and not y[2] in newtreeMin:
newtreeMin.append(y[2])
numEvents = len(eventList)
for event in range(len(eventList)):
turtle.penup()
nodeDict[eventList[event][1]].append(((event+1)*width/(numEvents+1)-DISPLACE, depth -(difference-radius)))
turtle.setpos(((event+1)*width/(numEvents+1))-DISPLACE, depth - difference)
turtle.pendown()
turtle.circle(radius)
turtle.left(95)
turtle.penup()
turtle.forward(radius)
turtle.pendown()
turtle.right(95)
turtle.write(eventList[event][0], font = ("arial", 12, "normal"))
turtle.ht()
drawNodes(newtreeMin, eventDict, depth - 2*difference, nodeDict)