-
Notifications
You must be signed in to change notification settings - Fork 0
/
location_tree.py
59 lines (47 loc) · 1.33 KB
/
location_tree.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
from ete2 import NodeStyle, TreeStyle, Tree, TextFace
import sys
import csv
t = Tree(sys.argv[1])
R = t.get_midpoint_outgroup()
t.set_outgroup(R)
colours = {'X' : 'red', 'Y' : 'blue', 'Z': 'green', 'W' : 'yellow'}
#colours = {'CC-A' : 'red', 'CC-B' : 'blue', 'CC-C': 'green', 'CC-D' : 'yellow'}
samples = {}
with open(sys.argv[2]) as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
samples[row['IsolateID']] = row
# Basic tree style
ts = TreeStyle()
ts.show_leaf_name = True
ts.mode = "c"
#ts.arc_start = -180 # 0 degrees = 3 o'clock
#ts.arc_span = 180
nstyle = NodeStyle()
nstyle['shape'] = 'sphere'
nstyle['size'] = 4
nstyle['fgcolor'] = 'black'
for n in t.traverse():
n.set_style(nstyle)
# Create an independent node style for each leaf, which
# specifies the colour given in the locations.csv file
for n in t.get_leaves():
name = n.get_leaf_names()[0]
print name
nstyle = NodeStyle()
try:
nstyle["fgcolor"] = colours[samples[name]['Hospitals']]
#Hospitals']]
except KeyError:
nstyle["fgcolor"] = "grey"
nstyle["size"] = 10
n.set_style(nstyle)
try:
n.add_face(TextFace(samples[name]['Locations']), column=0, position = "branch-right")
except:
pass
try:
n.add_face(TextFace(samples[name]['SamplingPeriod']), column=1, position = "branch-right")
except:
pass
t.render(file_name=sys.argv[3], tree_style=ts)