Mercurial > repos > jjohnson > ensembl_cdna_translate
comparison ensembl_cdna_translate.py @ 0:a8218b11216f draft
Uploaded
author | jjohnson |
---|---|
date | Wed, 29 Nov 2017 15:55:59 -0500 |
parents | |
children | b7f2f5e3390c |
comparison
equal
deleted
inserted
replaced
-1:000000000000 | 0:a8218b11216f |
---|---|
1 #!/usr/bin/env python | |
2 """ | |
3 # | |
4 #------------------------------------------------------------------------------ | |
5 # University of Minnesota | |
6 # Copyright 2017, Regents of the University of Minnesota | |
7 #------------------------------------------------------------------------------ | |
8 # Author: | |
9 # | |
10 # James E Johnson | |
11 # | |
12 #------------------------------------------------------------------------------ | |
13 """ | |
14 | |
15 import argparse | |
16 import sys | |
17 from time import sleep | |
18 | |
19 from Bio.Seq import translate | |
20 | |
21 import requests | |
22 | |
23 import digest | |
24 | |
25 | |
26 server = "https://rest.ensembl.org" | |
27 ext = "/info/assembly/homo_sapiens?" | |
28 max_region = 5000000 | |
29 | |
30 | |
31 def ensembl_rest(ext, headers): | |
32 if True: print >> sys.stderr, "%s" % ext | |
33 r = requests.get(server+ext, headers=headers) | |
34 if r.status_code == 429: | |
35 print >> sys.stderr, "response headers: %s\n" % r.headers | |
36 if 'Retry-After' in r.headers: | |
37 sleep(r.headers['Retry-After']) | |
38 r = requests.get(server+ext, headers=headers) | |
39 if not r.ok: | |
40 r.raise_for_status() | |
41 return r | |
42 | |
43 | |
44 def get_species(): | |
45 results = dict() | |
46 ext = "/info/species" | |
47 req_header = {"Content-Type": "application/json"} | |
48 r = ensembl_rest(ext, req_header) | |
49 for species in r.json()['species']: | |
50 results[species['name']] = species | |
51 print >> sys.stdout,\ | |
52 "%s\t%s\t%s\t%s\t%s"\ | |
53 % (species['name'], species['common_name'], | |
54 species['display_name'], | |
55 species['strain'], | |
56 species['taxon_id']) | |
57 return results | |
58 | |
59 | |
60 def get_biotypes(species): | |
61 biotypes = [] | |
62 ext = "/info/biotypes/%s?" % species | |
63 req_header = {"Content-Type": "application/json"} | |
64 r = ensembl_rest(ext, req_header) | |
65 for entry in r.json(): | |
66 if 'biotype' in entry: | |
67 biotypes.append(entry['biotype']) | |
68 return biotypes | |
69 | |
70 | |
71 def get_toplevel(species): | |
72 coord_systems = dict() | |
73 ext = "/info/assembly/%s?" % species | |
74 req_header = {"Content-Type": "application/json"} | |
75 r = ensembl_rest(ext, req_header) | |
76 toplevel = r.json() | |
77 for seq in toplevel['top_level_region']: | |
78 if seq['coord_system'] not in coord_systems: | |
79 coord_systems[seq['coord_system']] = dict() | |
80 coord_system = coord_systems[seq['coord_system']] | |
81 coord_system[seq['name']] = int(seq['length']) | |
82 return coord_systems | |
83 | |
84 | |
85 def get_transcripts_bed(species, refseq, start, length,params=None): | |
86 bed = [] | |
87 param = params if params else '' | |
88 req_header = {"Content-Type": "text/x-bed"} | |
89 regions = range(start, length, max_region) | |
90 if not regions or regions[-1] < length: | |
91 regions.append(length) | |
92 for end in regions[1:]: | |
93 ext = "/overlap/region/%s/%s:%d-%d?feature=transcript;%s"\ | |
94 % (species, refseq, start, end, param) | |
95 start = end + 1 | |
96 r = ensembl_rest(ext, req_header) | |
97 if r.text: | |
98 bed += r.text.splitlines() | |
99 return bed | |
100 | |
101 | |
102 def get_seq(id, seqtype,params=None): | |
103 param = params if params else '' | |
104 ext = "/sequence/id/%s?type=%s;%s" % (id, seqtype,param) | |
105 req_header = {"Content-Type": "text/plain"} | |
106 r = ensembl_rest(ext, req_header) | |
107 return r.text | |
108 | |
109 | |
110 def get_cdna(id,params=None): | |
111 return get_seq(id, 'cdna',params=params) | |
112 | |
113 | |
114 def get_cds(id,params=None): | |
115 return get_seq(id, 'cds',params=params) | |
116 | |
117 | |
118 def bed_from_line(line): | |
119 fields = line.rstrip('\r\n').split('\t') | |
120 (chrom, chromStart, chromEnd, name, score, strand, | |
121 thickStart, thickEnd, itemRgb, | |
122 blockCount, blockSizes, blockStarts) = fields[0:12] | |
123 bed_entry = BedEntry(chrom=chrom, chromStart=chromStart, chromEnd=chromEnd, | |
124 name=name, score=score, strand=strand, | |
125 thickStart=thickStart, thickEnd=thickEnd, | |
126 itemRgb=itemRgb, | |
127 blockCount=blockCount, | |
128 blockSizes=blockSizes.rstrip(','), | |
129 blockStarts=blockStarts.rstrip(',')) | |
130 return bed_entry | |
131 | |
132 | |
133 class BedEntry(object): | |
134 def __init__(self, chrom=None, chromStart=None, chromEnd=None, | |
135 name=None, score=None, strand=None, | |
136 thickStart=None, thickEnd=None, itemRgb=None, | |
137 blockCount=None, blockSizes=None, blockStarts=None): | |
138 self.chrom = chrom | |
139 self.chromStart = int(chromStart) | |
140 self.chromEnd = int(chromEnd) | |
141 self.name = name | |
142 self.score = int(score) if score is not None else 0 | |
143 self.strand = '-' if str(strand).startswith('-') else '+' | |
144 self.thickStart = int(thickStart) if thickStart else self.chromStart | |
145 self.thickEnd = int(thickEnd) if thickEnd else self.chromEnd | |
146 self.itemRgb = str(itemRgb) if itemRgb is not None else r'100,100,100' | |
147 self.blockCount = int(blockCount) | |
148 if isinstance(blockSizes, str) or isinstance(blockSizes, unicode): | |
149 self.blockSizes = [int(x) for x in blockSizes.split(',')] | |
150 elif isinstance(blockSizes, list): | |
151 self.blockSizes = [int(x) for x in blockSizes] | |
152 else: | |
153 self.blockSizes = blockSizes | |
154 if isinstance(blockStarts, str) or isinstance(blockSizes, unicode): | |
155 self.blockStarts = [int(x) for x in blockStarts.split(',')] | |
156 elif isinstance(blockStarts, list): | |
157 self.blockStarts = [int(x) for x in blockStarts] | |
158 else: | |
159 self.blockStarts = blockStarts | |
160 self.seq = None | |
161 self.pep = None | |
162 | |
163 def __str__(self): | |
164 return '%s\t%d\t%d\t%s\t%d\t%s\t%d\t%d\t%s\t%d\t%s\t%s' % ( | |
165 self.chrom, self.chromStart, self.chromEnd, | |
166 self.name, self.score, self.strand, | |
167 self.thickStart, self.thickEnd, str(self.itemRgb), self.blockCount, | |
168 ','.join([str(x) for x in self.blockSizes]), | |
169 ','.join([str(x) for x in self.blockStarts])) | |
170 | |
171 # (start, end) | |
172 def get_subrange(self, tstart, tstop, debug=False): | |
173 chromStart = self.chromStart | |
174 chromEnd = self.chromEnd | |
175 if debug: | |
176 print >> sys.stderr, "%s" % (str(self)) | |
177 r = range(self.blockCount) | |
178 if self.strand == '-': | |
179 r.reverse() | |
180 bStart = 0 | |
181 bEnd = 0 | |
182 for x in r: | |
183 bEnd = bStart + self.blockSizes[x] | |
184 if bStart <= tstart < bEnd: | |
185 if self.strand == '+': | |
186 chromStart = self.chromStart + self.blockStarts[x] +\ | |
187 (tstart - bStart) | |
188 else: | |
189 chromEnd = self.chromStart + self.blockStarts[x] +\ | |
190 self.blockSizes[x] - (tstart - bStart) | |
191 if bStart <= tstop < bEnd: | |
192 if self.strand == '+': | |
193 chromEnd = self.chromStart + self.blockStarts[x] +\ | |
194 (tstop - bStart) | |
195 else: | |
196 chromStart = self.chromStart + self.blockStarts[x] +\ | |
197 self.blockSizes[x] - (tstop - bStart) | |
198 if debug: | |
199 print >> sys.stderr,\ | |
200 "%3d %s\t%d\t%d\t%d\t%d\t%d\t%d"\ | |
201 % (x, self.strand, bStart, bEnd, | |
202 tstart, tstop, chromStart, chromEnd) | |
203 bStart += self.blockSizes[x] | |
204 return(chromStart, chromEnd) | |
205 | |
206 # get the blocks for sub range | |
207 def get_blocks(self, chromStart, chromEnd): | |
208 tblockCount = 0 | |
209 tblockSizes = [] | |
210 tblockStarts = [] | |
211 for x in range(self.blockCount): | |
212 bStart = self.chromStart + self.blockStarts[x] | |
213 bEnd = bStart + self.blockSizes[x] | |
214 if bStart > chromEnd: | |
215 break | |
216 if bEnd < chromStart: | |
217 continue | |
218 cStart = max(chromStart, bStart) | |
219 tblockStarts.append(cStart - chromStart) | |
220 tblockSizes.append(min(chromEnd, bEnd) - cStart) | |
221 tblockCount += 1 | |
222 return (tblockCount, tblockSizes, tblockStarts) | |
223 | |
224 def trim(self, tstart, tstop, debug=False): | |
225 (tchromStart, tchromEnd) =\ | |
226 self.get_subrange(tstart, tstop, debug=debug) | |
227 (tblockCount, tblockSizes, tblockStarts) =\ | |
228 self.get_blocks(tchromStart, tchromEnd) | |
229 tbed = BedEntry( | |
230 chrom=self.chrom, chromStart=tchromStart, chromEnd=tchromEnd, | |
231 name=self.name, score=self.score, strand=self.strand, | |
232 thickStart=tchromStart, thickEnd=tchromEnd, itemRgb=self.itemRgb, | |
233 blockCount=tblockCount, | |
234 blockSizes=tblockSizes, blockStarts=tblockStarts) | |
235 if self.seq: | |
236 ts = tchromStart-self.chromStart | |
237 te = tchromEnd - tchromStart + ts | |
238 tbed.seq = self.seq[ts:te] | |
239 return tbed | |
240 | |
241 | |
242 def __main__(): | |
243 parser = argparse.ArgumentParser( | |
244 description='Retrieve Ensembl cDNAs and three frame translate') | |
245 parser.add_argument( | |
246 '-s', '--species', default='human', | |
247 help='Ensembl Species to retrieve') | |
248 parser.add_argument( | |
249 '-B', '--biotypes', action='append', default=[], | |
250 help='Restrict Ensembl biotypes to retrieve') | |
251 parser.add_argument( | |
252 '-i', '--input', default=None, | |
253 help='Use BED instead of retrieving cDNA from ensembl (-) for stdin') | |
254 parser.add_argument( | |
255 '-t', '--transcripts', default=None, | |
256 help='Path to output cDNA transcripts.bed (-) for stdout') | |
257 parser.add_argument( | |
258 '-r', '--raw', action='store_true', | |
259 help='Report transcript exacty as returned from Ensembl') | |
260 parser.add_argument( | |
261 '-f', '--fasta', default=None, | |
262 help='Path to output translations.fasta') | |
263 parser.add_argument( | |
264 '-b', '--bed', default=None, | |
265 help='Path to output translations.bed') | |
266 parser.add_argument( | |
267 '-m', '--min_length', type=int, default=7, | |
268 help='Minimum length of protein translation to report') | |
269 parser.add_argument( | |
270 '-e', '--enzyme', default=None, | |
271 help='Digest translation with enzyme') | |
272 parser.add_argument( | |
273 '-a', '--all', action='store_true', | |
274 help='Include reference protein translations') | |
275 parser.add_argument('-v', '--verbose', action='store_true', help='Verbose') | |
276 parser.add_argument('-d', '--debug', action='store_true', help='Debug') | |
277 args = parser.parse_args() | |
278 # print >> sys.stderr, "args: %s" % args | |
279 species = args.species | |
280 input_rdr = None | |
281 if args.input is not None: | |
282 input_rdr = open(args.input, 'r') if args.input != '-' else sys.stdin | |
283 tx_wtr = None | |
284 if args.transcripts is not None: | |
285 tx_wtr = open(args.transcripts, 'w')\ | |
286 if args.transcripts != '-' else sys.stdout | |
287 fa_wtr = open(args.fasta, 'w') if args.fasta is not None else None | |
288 bed_wtr = open(args.bed, 'w') if args.bed is not None else None | |
289 | |
290 enzyme = digest.expasy_rules.get(args.enzyme,args.enzyme) | |
291 | |
292 # print >> sys.stderr, "args biotypes: %s" % args.biotypes | |
293 biotypea = ['biotype=%s' % bt.strip() for biotype in args.biotypes for bt in biotype.split(',')] | |
294 # print >> sys.stderr, "args biotypes: %s" % biotypea | |
295 biotypes = ';'.join(['biotype=%s' % bt.strip() for biotype in args.biotypes for bt in biotype.split(',') if bt.strip()]) | |
296 # print >> sys.stderr, "biotypes: %s" % biotypes | |
297 | |
298 translations = dict() # start : end : seq | |
299 | |
300 def unique_prot(tbed, seq): | |
301 if tbed.chromStart not in translations: | |
302 translations[tbed.chromStart] = dict() | |
303 translations[tbed.chromStart][tbed.chromEnd] = [] | |
304 translations[tbed.chromStart][tbed.chromEnd].append(seq) | |
305 elif tbed.chromEnd not in translations[tbed.chromStart]: | |
306 translations[tbed.chromStart][tbed.chromEnd] = [] | |
307 translations[tbed.chromStart][tbed.chromEnd].append(seq) | |
308 elif seq not in translations[tbed.chromStart][tbed.chromEnd]: | |
309 translations[tbed.chromStart][tbed.chromEnd].append(seq) | |
310 else: | |
311 return False | |
312 return True | |
313 | |
314 def translate_bed(bed): | |
315 translate_count = 0 | |
316 if any([fa_wtr, bed_wtr]): | |
317 transcript_id = bed.name | |
318 refprot = None | |
319 if not args.all: | |
320 try: | |
321 cds = get_cds(transcript_id) | |
322 if len(cds) % 3 != 0: | |
323 cds = cds[:-(len(cds) % 3)] | |
324 refprot = translate(cds) if cds else None | |
325 except: | |
326 refprot = None | |
327 cdna = get_cdna(transcript_id) | |
328 cdna_len = len(cdna) | |
329 for offset in range(3): | |
330 seqend = cdna_len - (cdna_len - offset) % 3 | |
331 aaseq = translate(cdna[offset:seqend]) | |
332 aa_start = 0 | |
333 while aa_start < len(aaseq): | |
334 aa_end = aaseq.find('*', aa_start) | |
335 if aa_end < 0: | |
336 aa_end = len(aaseq) | |
337 prot = aaseq[aa_start:aa_end] | |
338 if enzyme and refprot: | |
339 frags = digest._cleave(prot,enzyme) | |
340 for frag in reversed(frags): | |
341 if frag in refprot: | |
342 prot = prot[:prot.rfind(frag)] | |
343 else: | |
344 break | |
345 if len(prot) < args.min_length: | |
346 pass | |
347 elif refprot and prot in refprot: | |
348 pass | |
349 else: | |
350 tstart = aa_start*3+offset | |
351 tend = aa_end*3+offset | |
352 prot_acc = "%s_%d_%d" % (transcript_id, tstart, tend) | |
353 tbed = bed.trim(tstart, tend) | |
354 if args.all or unique_prot(tbed, prot): | |
355 translate_count += 1 | |
356 tbed.name = prot_acc | |
357 bed_wtr.write("%s\t%s\n" % (str(tbed), prot)) | |
358 bed_wtr.flush() | |
359 fa_id = ">%s\n" % (prot_acc) | |
360 fa_wtr.write(fa_id) | |
361 fa_wtr.write(prot) | |
362 fa_wtr.write("\n") | |
363 fa_wtr.flush() | |
364 aa_start = aa_end + 1 | |
365 return translate_count | |
366 | |
367 if input_rdr: | |
368 translation_count = 0 | |
369 for i, bedline in enumerate(input_rdr): | |
370 try: | |
371 bed = bed_from_line(bedline) | |
372 translation_count += translate_bed(bed) | |
373 except: | |
374 print >> sys.stderr, "BED format error: %s\n" % bedline | |
375 if args.debug or (args.verbose and any([fa_wtr, bed_wtr])): | |
376 print >> sys.stderr,\ | |
377 "%s\tcDNA translations:%d" % (species, translation_count) | |
378 else: | |
379 coord_systems = get_toplevel(species) | |
380 if 'chromosome' in coord_systems: | |
381 for ref in sorted(coord_systems['chromosome'].keys()): | |
382 length = coord_systems['chromosome'][ref] | |
383 if not any([tx_wtr, fa_wtr, bed_wtr]): | |
384 print >> sys.stderr,\ | |
385 "%s\t%s\tlength: %d" % (species, ref, length) | |
386 continue | |
387 if args.debug: | |
388 print >> sys.stderr,\ | |
389 "Retrieving transcripts: %s\t%s\tlength: %d"\ | |
390 % (species, ref, length) | |
391 translation_count = 0 | |
392 start = 0 | |
393 regions = range(start, length, max_region) | |
394 if not regions or regions[-1] < length: | |
395 regions.append(length) | |
396 for end in regions[1:]: | |
397 bedlines = get_transcripts_bed(species, ref, start, end, params=biotypes) | |
398 if args.verbose or args.debug: | |
399 print >> sys.stderr,\ | |
400 "%s\t%s\tstart: %d\tend: %d\tcDNA transcripts:%d"\ | |
401 % (species, ref, start, end, len(bedlines)) | |
402 # start, end, seq | |
403 for i, bedline in enumerate(bedlines): | |
404 try: | |
405 bed = bed_from_line(bedline)\ | |
406 if any([not args.raw, fa_wtr, bed_wtr])\ | |
407 else None | |
408 if tx_wtr: | |
409 tx_wtr.write(bedline if args.raw else str(bed)) | |
410 tx_wtr.write("\n") | |
411 tx_wtr.flush() | |
412 if bed: | |
413 translation_count += translate_bed(bed) | |
414 except Exception as e: | |
415 print >> sys.stderr,\ | |
416 "BED error (%s) : %s\n" % (e, bedline) | |
417 start = end + 1 | |
418 | |
419 if args.debug or (args.verbose and any([fa_wtr, bed_wtr])): | |
420 print >> sys.stderr,\ | |
421 "%s\t%s\tlength: %d\tcDNA translations:%d"\ | |
422 % (species, ref, length, translation_count) | |
423 | |
424 | |
425 if __name__ == "__main__": | |
426 __main__() |