annotate cpo_galaxy_tree.py @ 1:fea89c4d5227 draft

Uploaded
author jjjjia
date Thu, 16 Aug 2018 19:27:05 -0400
parents
children cabceaa239e4
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
1 #!/home/jjjjia/.conda/envs/py36/bin/python
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
2
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
3 #$ -S /home/jjjjia/.conda/envs/py36/bin/python
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
4 #$ -V # Pass environment variables to the job
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
5 #$ -N CPO_pipeline # Replace with a more specific job name
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
6 #$ -wd /home/jjjjia/testCases # Use the current working dir
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
7 #$ -pe smp 8 # Parallel Environment (how many cores)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
8 #$ -l h_vmem=11G # Memory (RAM) allocation *per core*
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
9 #$ -e ./logs/$JOB_ID.err
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
10 #$ -o ./logs/$JOB_ID.log
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
11 #$ -m ea
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
12 #$ -M bja20@sfu.ca
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
13
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
14 #~/scripts/pipeline.py -i BC11-Kpn005_S2 -f /data/jjjjia/R1/BC11-Kpn005_S2_L001_R1_001.fastq.gz -r /data/jjjjia/R2/BC11-Kpn005_S2_L001_R2_001.fastq.gz -o pipelineResultsQsub -e "Klebsiella pneumoniae"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
15
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
16 import subprocess
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
17 import pandas
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
18 import optparse
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
19 import os
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
20 import datetime
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
21 import sys
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
22 import time
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
23 import urllib.request
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
24 import gzip
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
25 import collections
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
26 import json
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
27 import ete3
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
28 import numpy
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
29
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
30 #parses some parameters
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
31 parser = optparse.OptionParser("Usage: %prog [options] arg1 arg2 ...")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
32 parser.add_option("-t", "--tree", dest="treePath", type="string", default="./pipelineTest/tree.txt", help="identifier of the isolate")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
33 parser.add_option("-d", "--distance", dest="distancePath", type="string", default="./pipelineTest/distance.tab", help="absolute file path forward read (R1)")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
34 parser.add_option("-m", "--metadata", dest="metadataPath", type="string", default="./pipelineTest/metadata.tsv",help="absolute file path to reverse read (R2)")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
35 (options,args) = parser.parse_args()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
36 treePath = str(options.treePath).lstrip().rstrip()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
37 distancePath = str(options.distancePath).lstrip().rstrip()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
38 metadataPath = str(options.metadataPath).lstrip().rstrip()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
39
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
40
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
41 #region result objects
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
42 #define some objects to store values from results
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
43 #//TODO this is not the proper way of get/set private object variables. every value has manually assigned defaults intead of specified in init(). Also, use property(def getVar, def setVar).
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
44 class workflowResult(object):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
45 def __init__(self):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
46 self.new = False
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
47 self.ID = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
48 self.ExpectedSpecies = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
49 self.MLSTSpecies = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
50 self.SequenceType = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
51 self.MLSTScheme = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
52 self.CarbapenemResistanceGenes =""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
53 self.OtherAMRGenes=""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
54 self.TotalPlasmids = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
55 self.plasmids = []
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
56 self.DefinitelyPlasmidContigs =""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
57 self.LikelyPlasmidContigs=""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
58 self.row = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
59 class plasmidObj(object):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
60 def __init__(self):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
61 self.PlasmidsID = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
62 self.Num_Contigs = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
63 self.PlasmidLength = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
64 self.PlasmidRepType = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
65 self.PlasmidMobility = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
66 self.NearestReference = ""
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
67
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
68 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
69
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
70 #region useful functions
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
71 def read(path):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
72 return [line.rstrip('\n') for line in open(path)]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
73 def execute(command):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
74 process = subprocess.Popen(command, shell=False, cwd=curDir, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
75
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
76 # Poll process for new output until finished
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
77 while True:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
78 nextline = process.stdout.readline()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
79 if nextline == '' and process.poll() is not None:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
80 break
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
81 sys.stdout.write(nextline)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
82 sys.stdout.flush()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
83
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
84 output = process.communicate()[0]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
85 exitCode = process.returncode
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
86
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
87 if (exitCode == 0):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
88 return output
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
89 else:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
90 raise subprocess.CalledProcessError(exitCode, command)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
91 def httpGetFile(url, filepath=""):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
92 if (filepath == ""):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
93 return urllib.request.urlretrieve(url)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
94 else:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
95 urllib.request.urlretrieve(url, filepath)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
96 return True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
97 def gunzip(inputpath="", outputpath=""):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
98 if (outputpath == ""):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
99 with gzip.open(inputpath, 'rb') as f:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
100 gzContent = f.read()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
101 return gzContent
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
102 else:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
103 with gzip.open(inputpath, 'rb') as f:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
104 gzContent = f.read()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
105 with open(outputpath, 'wb') as out:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
106 out.write(gzContent)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
107 return True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
108 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
109
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
110 #region functions to parse result files
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
111 def ParseWorkflowResults(pathToResult):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
112 _worflowResult = {}
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
113 r = pandas.read_csv(pathToResult, delimiter='\t', header=None) #read the kraken2report.tsv
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
114 r = r.replace(numpy.nan, '', regex=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
115 for i in range(len(r.index)):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
116 _results = workflowResult()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
117 if(str(r.iloc[i,0]).lower() == "new"):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
118 _results.new = True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
119 else:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
120 _results.new = False
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
121 _results.ID = str(r.iloc[i,1])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
122 _results.ExpectedSpecies = str(r.iloc[i,2])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
123 _results.MLSTSpecies = str(r.iloc[i,3])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
124 _results.SequenceType = str(r.iloc[i,4])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
125 _results.MLSTScheme = (str(r.iloc[i,5]))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
126 _results.CarbapenemResistanceGenes = (str(r.iloc[i,6]))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
127 _results.OtherAMRGenes = (str(r.iloc[i,7]))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
128 _results.TotalPlasmids = int(r.iloc[i,8])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
129 for j in range(0,_results.TotalPlasmids):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
130 _plasmid = plasmidObj()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
131 _plasmid.PlasmidsID =(((str(r.iloc[i,9])).split(";"))[j])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
132 _plasmid.Num_Contigs = (((str(r.iloc[i,10])).split(";"))[j])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
133 _plasmid.PlasmidLength = (((str(r.iloc[i,11])).split(";"))[j])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
134 _plasmid.PlasmidRepType = (((str(r.iloc[i,12])).split(";"))[j])
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
135 _plasmid.PlasmidMobility = ((str(r.iloc[i,13])).split(";"))[j]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
136 _plasmid.NearestReference = ((str(r.iloc[i,14])).split(";"))[j]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
137 _results.plasmids.append(_plasmid)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
138 _results.DefinitelyPlasmidContigs = (str(r.iloc[i,15]))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
139 _results.LikelyPlasmidContigs = (str(r.iloc[i,16]))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
140 _results.row = "\t".join(str(x) for x in r.ix[i].tolist())
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
141 _worflowResult[_results.ID] = _results
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
142 return _worflowResult
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
143
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
144 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
145
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
146 def Main():
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
147 metadata = ParseWorkflowResults(metadataPath)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
148 distance = read(distancePath)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
149 treeFile = "".join(read(treePath))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
150
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
151 distanceDict = {}
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
152 for i in range(len(distance)):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
153 temp = distance[i].split("\t")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
154 distanceDict[temp[0]] = temp[1:]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
155 #region step5: tree construction
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
156 t = ete3.Tree(treeFile)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
157 t.set_outgroup(t&"Reference")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
158
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
159 ts = ete3.TreeStyle()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
160 ts.show_leaf_name = True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
161 ts.show_branch_length = True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
162 ts.scale = 2000 #pixel per branch length unit
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
163 ts.branch_vertical_margin = 15 #pixel between branches
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
164 style2 = ete3.NodeStyle()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
165 style2["fgcolor"] = "#000000"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
166 style2["shape"] = "circle"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
167 style2["vt_line_color"] = "#0000aa"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
168 style2["hz_line_color"] = "#0000aa"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
169 style2["vt_line_width"] = 2
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
170 style2["hz_line_width"] = 2
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
171 style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
172 style2["hz_line_type"] = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
173 for n in t.traverse():
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
174 n.set_style(style2)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
175 '''
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
176 #region create detailed tree
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
177
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
178 plasmidCount = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
179 for n in t.traverse():
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
180 if (n.is_leaf() and not n.name == "Reference"):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
181 mData = metadata[n.name.replace(".fa","")]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
182 face = ete3.faces.TextFace(mData.MLSTSpecies,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
183 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
184 face.margin_left = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
185 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
186 n.add_face(face, 0, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
187 face = ete3.faces.TextFace(mData.SequenceType,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
188 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
189 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
190 n.add_face(face, 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
191 face = ete3.faces.TextFace(mData.CarbapenemResistanceGenes,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
192 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
193 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
194 n.add_face(face, 2, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
195 index = 3
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
196 if (mData.TotalPlasmids > plasmidCount):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
197 plasmidCount = mData.TotalPlasmids
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
198 for i in range(0, mData.TotalPlasmids):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
199 face = ete3.faces.TextFace(mData.plasmids[i].PlasmidRepType,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
200 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
201 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
202 n.add_face(face, index, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
203 index+=1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
204 face = ete3.faces.TextFace(mData.plasmids[i].PlasmidMobility,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
205 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
206 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
207 n.add_face(face, index, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
208 index+=1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
209
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
210 face = ete3.faces.TextFace("Species",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
211 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
212 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
213 face.margin_left = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
214 (t&"Reference").add_face(face, 0, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
215 face = ete3.faces.TextFace("Sequence Type",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
216 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
217 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
218 (t&"Reference").add_face(face, 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
219 face = ete3.faces.TextFace("Carbapenamases",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
220 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
221 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
222 (t&"Reference").add_face(face, 2, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
223 index = 3
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
224 for i in range(0, plasmidCount):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
225 face = ete3.faces.TextFace("plasmid " + str(i) + " replicons",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
226 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
227 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
228 (t&"Reference").add_face(face, index, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
229 index+=1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
230 face = ete3.faces.TextFace("plasmid " + str(i) + " mobility",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
231 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
232 face.margin_right = 10
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
233 (t&"Reference").add_face(face, index, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
234 index+=1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
235
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
236 t.render("./pipelineTest/tree.png", w=5000,units="mm", tree_style=ts)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
237
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
238 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
239 '''
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
240 #region create box tree
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
241 #region step5: tree construction
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
242 treeFile = "".join(read("./pipelineTest/tree.txt"))
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
243 t = ete3.Tree(treeFile)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
244 t.set_outgroup(t&"Reference")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
245
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
246 ts = ete3.TreeStyle()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
247 ts.show_leaf_name = True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
248 ts.show_branch_length = True
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
249 ts.scale = 2000 #pixel per branch length unit
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
250 ts.branch_vertical_margin = 15 #pixel between branches
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
251 style2 = ete3.NodeStyle()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
252 style2["fgcolor"] = "#000000"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
253 style2["shape"] = "circle"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
254 style2["vt_line_color"] = "#0000aa"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
255 style2["hz_line_color"] = "#0000aa"
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
256 style2["vt_line_width"] = 2
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
257 style2["hz_line_width"] = 2
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
258 style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
259 style2["hz_line_type"] = 0
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
260 for n in t.traverse():
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
261 n.set_style(style2)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
262 plasmidIncs = {}
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
263 for key in metadata:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
264 for plasmid in metadata[key].plasmids:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
265 for inc in plasmid.PlasmidRepType.split(","):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
266 if (inc.lower().find("inc") > -1):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
267 if not (inc in plasmidIncs):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
268 plasmidIncs[inc] = [metadata[key].ID]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
269 else:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
270 if metadata[key].ID not in plasmidIncs[inc]:
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
271 plasmidIncs[inc].append(metadata[key].ID)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
272 #plasmidIncs = sorted(plasmidIncs)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
273 for n in t.traverse():
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
274 if (n.is_leaf() and n.name == "Reference"):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
275 face = ete3.faces.TextFace("New?",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
276 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
277 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
278 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
279 (t&"Reference").add_face(face, 0, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
280 for i in range(len(plasmidIncs)): #this loop adds the columns (aka the incs) to the reference node
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
281 face = ete3.faces.TextFace(list(plasmidIncs.keys())[i],fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
282 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
283 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
284 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
285 (t&"Reference").add_face(face, i + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
286 face = ete3.faces.TextFace("MLSTScheme",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
287 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
288 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
289 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
290 (t&"Reference").add_face(face, len(plasmidIncs) + 0 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
291 face = ete3.faces.TextFace("Sequence Type",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
292 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
293 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
294 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
295 (t&"Reference").add_face(face, len(plasmidIncs) + 1 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
296 face = ete3.faces.TextFace("Carbapenamases",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
297 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
298 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
299 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
300 (t&"Reference").add_face(face, len(plasmidIncs) + 2 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
301 for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds the columns (aka the incs) to the reference node
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
302 face = ete3.faces.TextFace(distanceDict[list(distanceDict.keys())[0]][i],fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
303 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
304 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
305 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
306 (t&"Reference").add_face(face, len(plasmidIncs) + 2 + i + 1 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
307 elif (n.is_leaf() and not n.name == "Reference"):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
308 if (metadata[n.name.replace(".fa","")].new == True):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
309 face = ete3.faces.RectFace(30,30,"green","green") # TextFace("Y",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
310 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
311 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
312 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
313 face.vt_align = 1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
314 face.ht_align = 1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
315 n.add_face(face, 0, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
316 for incs in plasmidIncs: #this loop adds presence/absence to the sample nodes
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
317 if (n.name.replace(".fa","") in plasmidIncs[incs]):
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
318 face = ete3.faces.RectFace(30,30,"black","black") # TextFace("Y",fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
319 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
320 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
321 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
322 face.vt_align = 1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
323 face.ht_align = 1
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
324 n.add_face(face, list(plasmidIncs.keys()).index(incs) + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
325 mData = metadata[n.name.replace(".fa","")]
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
326 face = ete3.faces.TextFace(mData.MLSTSpecies,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
327 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
328 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
329 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
330 n.add_face(face, len(plasmidIncs) + 0 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
331 face = ete3.faces.TextFace(mData.SequenceType,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
332 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
333 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
334 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
335 n.add_face(face, len(plasmidIncs) + 1 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
336 face = ete3.faces.TextFace(mData.CarbapenemResistanceGenes,fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
337 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
338 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
339 n.add_face(face, len(plasmidIncs) + 2 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
340 for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds distance matrix
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
341 face = ete3.faces.TextFace(list(distanceDict[n.name])[i],fsize=10,tight_text=True)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
342 face.border.margin = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
343 face.margin_right = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
344 face.margin_left = 5
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
345 n.add_face(face, len(plasmidIncs) + 2 + i + 1 + 1, "aligned")
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
346
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
347 t.render("./tree.png", w=5000,units="mm", tree_style=ts)
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
348
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
349 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
350 #endregion
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
351
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
352
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
353 start = time.time()#time the analysis
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
354
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
355 #analysis time
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
356 Main()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
357
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
358 end = time.time()
fea89c4d5227 Uploaded
jjjjia
parents:
diff changeset
359 print("Finished!\nThe analysis used: " + str(end-start) + " seconds")