[Bio] / Babel / bin / source2ach.py Repository:
ViewVC logotype

Diff of /Babel/bin/source2ach.py

Parent Directory Parent Directory | Revision Log Revision Log | View Patch Patch

revision 1.9, Wed Dec 1 15:45:57 2010 UTC revision 1.10, Tue Mar 1 15:21:48 2011 UTC
# Line 2  Line 2 
2    
3  import os, sys, re, hashlib  import os, sys, re, hashlib
4  from optparse import OptionParser  from optparse import OptionParser
5    from multiprocessing import Pool
6  from Bio import SeqIO  from Bio import SeqIO
7  from Bio.Seq import Seq  from Bio.Seq import Seq
8  from Bio.SeqRecord import SeqRecord  from Bio.SeqRecord import SeqRecord
# Line 25  Line 26 
26    
27  Values with * are optional, only genbank files contain that data."""  Values with * are optional, only genbank files contain that data."""
28    
29    class Info:
30        def __init__(self):
31            self.source  = ''
32            self.format  = 'fasta'
33            self.verbose = False
34            self.outdir  = ''
35            self.orghead = False
36            self.getont  = False
37            self.getctg  = False
38            self.gettax  = False
39            self.amap    = {}
40    
41  desc_re  = re.compile('(Rec|Sub)Name: Full=(.*?);')  desc_re  = re.compile('(Rec|Sub)Name: Full=(.*?);')
42  entry_re = re.compile('^ENTRY\s+(\S+)\s+CDS\s+(.*)$', re.MULTILINE)  entry_re = re.compile('^ENTRY\s+(\S+)\s+CDS\s+(.*)$', re.MULTILINE)
43  name_re  = re.compile('^NAME\s+(.*)$', re.MULTILINE)  name_re  = re.compile('^NAME\s+(.*)$', re.MULTILINE)
# Line 37  Line 50 
50  gb_go_rm = re.compile('( \[Evidence [A-Z]+\])')  gb_go_rm = re.compile('( \[Evidence [A-Z]+\])')
51  ko_re    = re.compile('^(K\d+)\s+(.*)')  ko_re    = re.compile('^(K\d+)\s+(.*)')
52  ogid_re  = re.compile('^(\S+?OG)\d+')  ogid_re  = re.compile('^(\S+?OG)\d+')
53    nr_re    = re.compile('^gi\|(\d+)\|(\w+)\|(\S+)\|\S*?\s+(.*)\]$')
54    nr_types = {'ref': 'RefSeq', 'gb': 'GenBank', 'emb': 'EMBL', 'dbj': 'DDBJ', 'pir': 'PIR', 'prf': 'PRF', 'pdb': 'PDB'}
55  go_types = ['GO_function', 'GO_process', 'GO_component']  go_types = ['GO_function', 'GO_process', 'GO_component']
56  file_ext = [".md52id2func", ".md52id2ont", ".md52seq", ".id2xref", ".id2tax"]  file_ext = [".md52id2func", ".md52id2ont", ".md52seq", ".id2xref", ".id2tax"]
57    params   = Info()
58    
59    def parse_nr_header(line):
60        items = []
61        hdr_set = nr_re.match(line)
62        if hdr_set:
63            (gi_id, nrdb, nrdb_id, func_org) = hdr_set.groups()
64            if nrdb in params.amap:
65                org_rev = ''
66                rev_fo  = reverse_string(func_org)
67                end_brk = 1
68                beg_brk = 0
69                for i, c in enumerate(rev_fo):
70                    if c == ']': end_brk += 1
71                    if c == '[': beg_brk += 1
72                    if beg_brk == end_brk: break
73                    org_rev += c
74                org_txt  = reverse_string( org_rev )
75                func_txt = reverse_string( rev_fo[i+1:] )
76                items = [ gi_id, params.amap[nrdb], nrdb_id, func_txt.strip(), org_txt.strip() ]
77        return items
78    
79  def get_eggnog_map(filename, verb):  def reverse_string(string):
80        return string[::-1]
81    
82    def get_eggnog_map(filename):
83      emap = {}      emap = {}
84      if filename is None:      if filename is None:
85          return emap          return emap
86      if verb: sys.stdout.write("Parsing file %s ... " %filename)      if params.verbose: sys.stdout.write("Parsing file %s ... " %filename)
87      fhdl = open(filename, 'r')      fhdl = open(filename, 'r')
88      try:      try:
89          for line in fhdl:          for line in fhdl:
# Line 56  Line 95 
95                  emap[ parts[0] ] = [ [parts[3], parts[4]] ]                  emap[ parts[0] ] = [ [parts[3], parts[4]] ]
96      finally:      finally:
97          fhdl.close()          fhdl.close()
98      if verb: sys.stdout.write("Done\n")      if params.verbose: sys.stdout.write("Done\n")
99      return emap      return emap
100    
101  def get_kegg_map(filename, verb):  def get_kegg_map(filename):
102      kmap = {}      kmap = {}
103      if filename is None:      if filename is None:
104          return kmap          return kmap
105      if verb: sys.stdout.write("Parsing file %s ... " %filename)      if params.verbose: sys.stdout.write("Parsing file %s ... " %filename)
106      for text in kegg_iter(filename):      for text in kegg_iter(filename):
107          rec = get_kegg_rec(text)          rec = get_kegg_rec(text)
108          if (not rec.name) or (not rec.description): continue          if (not rec.name) or (not rec.description): continue
109          names = rec.name.split(', ')          names = rec.name.split(', ')
110          if len(names) > 1:          if len(names) > 1:
111              kmap[ names[1] ] = rec.description              kmap[ names[1] ] = rec.description
112      if verb: sys.stdout.write("Done\n")      if params.verbose: sys.stdout.write("Done\n")
113      return kmap      return kmap
114    
115  def kegg_iter(filename):  def kegg_iter(filename):
# Line 110  Line 149 
149          record.seq = Seq( re.sub('\s+', '', seq_s.group(1)), IUPAC.protein )          record.seq = Seq( re.sub('\s+', '', seq_s.group(1)), IUPAC.protein )
150      return record      return record
151    
152  def format_factory(out_files, source, form, org_desc, get_ont, get_ctg, get_tax):  def format_factory(out_files):
153      (prot_f, ont_f, seq_f, ref_f, tax_f) = out_files      (prot_f, ont_f, seq_f, ref_f, tax_f) = out_files
154        source   = params.source
155        form     = params.format
156        org_desc = params.orghead
157        get_ont  = params.getont
158        get_ctg  = params.getctg
159        get_tax  = params.gettax
160        amap     = params.amap
161      if form == 'genbank':      if form == 'genbank':
162          def parse_genbank(rec, amap):          def parse_genbank(rec):
163              if (not rec) or (not rec.id) or (not rec.seq): return              if (not rec) or (not rec.id) or (not rec.seq): return
164              cid   = rec.id              cid   = rec.id
165              cdesc = rec.description.rstrip('.')              cdesc = rec.description.rstrip('.')
166              clen  = len(rec.seq)              seq   = str(rec.seq).upper()
167                clen  = len(seq)
168              org   = rec.annotations['organism']              org   = rec.annotations['organism']
169              # only one entry on rec (first feat is source and is not a feature), use first real feature, no contig data              # only one entry on rec (first feat is source and is not a feature), use first real feature, no contig data
170              if len(rec.features) == 2:              if len(rec.features) == 2:
171                  md5  = hashlib.md5(str(rec.seq)).hexdigest()                  md5  = hashlib.md5(seq).hexdigest()
172                  func = ''                  func = ''
173                  if 'product' in rec.features[1].qualifiers:                  if 'product' in rec.features[1].qualifiers:
174                      func = rec.features[1].qualifiers['product'][0]                      func = rec.features[1].qualifiers['product'][0]
175                  seq_f.write("%s\t%s\n" %(md5, rec.seq))                  seq_f.write("%s\t%s\n" %(md5, seq))
176                  prot_f.write("\t".join([md5, rec.name, func, org, source]) + "\n")                  prot_f.write("\t".join([md5, rec.name, func, org, source]) + "\n")
177                  if get_tax and ('taxonomy' in rec.annotations):                  if get_tax and ('taxonomy' in rec.annotations):
178                      if rec.annotations['taxonomy'][0] == 'Root':                      if rec.annotations['taxonomy'][0] == 'Root':
# Line 166  Line 213 
213                                      go_m = gb_go_re.match(goid)                                      go_m = gb_go_re.match(goid)
214                                      if go_m:                                      if go_m:
215                                          desc = gb_go_rm.sub('', go_m.group(2))                                          desc = gb_go_rm.sub('', go_m.group(2))
216                                          ont_f.write("\t".join([md5, go_m.group(1), desc, gtype.replace('_',':')]) + "\n")                                          ont_f.write("\t".join([md5, go_m.group(1), desc, 'GO']) + "\n")
217                      seq_f.write("%s\t%s\n" %(md5, seq))                      seq_f.write("%s\t%s\n" %(md5, seq))
218                      if get_ctg:                      if get_ctg:
219                          prot_f.write("\t".join([md5, fid, func, org, source, str(beg), str(end), str(strd), cid, cdesc, str(clen)]) + "\n")                          prot_f.write("\t".join([md5, fid, func, org, source, str(beg), str(end), str(strd), cid, cdesc, str(clen)]) + "\n")
# Line 177  Line 224 
224          return parse_genbank          return parse_genbank
225    
226      elif form == 'swiss':      elif form == 'swiss':
227          def parse_swiss(rec, amap):          def parse_swiss(rec):
228              if (not rec) or (not rec.id) or (not rec.seq): return              if (not rec) or (not rec.id) or (not rec.seq): return
229              md5  = hashlib.md5(str(rec.seq)).hexdigest()              seq  = str(rec.seq).upper()
230                md5  = hashlib.md5(seq).hexdigest()
231              func = rec.description              func = rec.description
232              srch = desc_re.search(func)              srch = desc_re.search(func)
233              if srch: func = srch.group(2)              if srch: func = srch.group(2)
234              seq_f.write("%s\t%s\n" %(md5, str(rec.seq)))              seq_f.write("%s\t%s\n" %(md5, seq))
235              prot_f.write("\t".join([md5, rec.id, func, rec.annotations['organism'], source]) + "\n")              prot_f.write("\t".join([md5, rec.id, func, rec.annotations['organism'], source]) + "\n")
236              if len(rec.dbxrefs) > 0:              if len(rec.dbxrefs) > 0:
237                  ref_str = rec.id                  ref_str = rec.id
# Line 191  Line 239 
239                      go_m = up_go_re.match(ref)                      go_m = up_go_re.match(ref)
240                      # output GO ontology                      # output GO ontology
241                      if go_m and get_ont:                      if go_m and get_ont:
242                          ont_f.write("\t".join([md5, go_m.group(1), '', 'GO']) + "\n")                          ont_f.write("\t".join([md5, go_m.group(1), func, 'GO']) + "\n")
243                      else:                      else:
244                          ref_str += "\t" + ref                          ref_str += "\t" + ref
245                  ref_f.write(ref_str + "\n")                  ref_f.write(ref_str + "\n")
246          return parse_swiss          return parse_swiss
247    
248      elif form == 'fasta':      elif form == 'fasta':
249          def parse_fasta(rec, amap):          def parse_fasta(rec):
250              if (not rec) or (not rec.id) or (not rec.seq): return              if (not rec) or (not rec.id) or (not rec.seq): return
251              desc = ''              desc = ''
252              org  = ''              org  = ''
253              func = ''              func = ''
254              hdrs = rec.description.split()              hdrs = rec.description.split()
255              md5  = hashlib.md5(str(rec.seq)).hexdigest()              seq  = str(rec.seq).upper()
256              seq_f.write("%s\t%s\n" %(md5, str(rec.seq)))              md5  = hashlib.md5(seq).hexdigest()
257                seq_f.write("%s\t%s\n" %(md5, seq))
258              if len(hdrs) > 1:              if len(hdrs) > 1:
259                  hdrs.pop(0)                  hdrs.pop(0)
260                  desc = " ".join(hdrs)                  desc = " ".join(hdrs)
# Line 223  Line 272 
272                  prot_f.write("\t".join([md5, rec.id, func, org, source]) + "\n")                  prot_f.write("\t".join([md5, rec.id, func, org, source]) + "\n")
273          return parse_fasta          return parse_fasta
274    
275        elif form == 'nr':
276            def parse_nr(rec):
277                if (not rec) or (not rec.description) or (not rec.seq): return
278                get_seq = False
279                hdrs = rec.description.split('\x01')
280                seq  = str(rec.seq).upper()
281                md5  = hashlib.md5(seq).hexdigest()
282                if len(hdrs) > 0:
283                    for h in hdrs:
284                        info = parse_nr_header(h)
285                        if len(info) > 0:
286                            (gi_id, nrdb, nrdb_id, func, org) = info
287                            get_seq = True
288                            prot_f.write("\t".join([md5, nrdb_id, func, org, nrdb]) + "\n")
289                            ref_f.write("%s\tGI:%s\n"%(nrdb_id, gi_id))
290                    if get_seq:
291                        seq_f.write("%s\t%s\n" %(md5, seq))
292            return parse_nr
293    
294      elif form == 'kegg':      elif form == 'kegg':
295          def parse_kegg(text, amap):          def parse_kegg(text):
296              rec = get_kegg_rec(text)              rec = get_kegg_rec(text)
297              if (not rec.id) or (not rec.description) or (not rec.seq): return              if (not rec.id) or (not rec.description) or (not rec.seq): return
298              md5  = hashlib.md5(str(rec.seq)).hexdigest()              seq  = str(rec.seq).upper()
299                md5  = hashlib.md5(seq).hexdigest()
300              org  = rec.annotations['organism']              org  = rec.annotations['organism']
301              orth = []              orth = []
302              if org in amap:              if org in amap:
# Line 237  Line 306 
306                  orth = rec.annotations['orthology']                  orth = rec.annotations['orthology']
307                  for oid, odesc in orth.iteritems():                  for oid, odesc in orth.iteritems():
308                      ont_f.write("\t".join([md5, oid, odesc, 'KO']) + "\n")                      ont_f.write("\t".join([md5, oid, odesc, 'KO']) + "\n")
309              seq_f.write("%s\t%s\n" %(md5, str(rec.seq)))              seq_f.write("%s\t%s\n" %(md5, seq))
310              prot_f.write("\t".join([md5, rec.id, rec.description, org, source]) + "\n")              prot_f.write("\t".join([md5, rec.id, rec.description, org, source]) + "\n")
311              if len(rec.dbxrefs) > 0:              if len(rec.dbxrefs) > 0:
312                  ref_f.write(rec.id + "\t" + "\t".join(rec.dbxrefs) + "\n")                  ref_f.write(rec.id + "\t" + "\t".join(rec.dbxrefs) + "\n")
# Line 246  Line 315 
315      else:      else:
316          return None          return None
317    
318    def process_file(infile):
319        o_files = []
320        o_hdls  = []
321        fformat = params.format
322        fname   = os.path.basename(infile)
323        for e in file_ext: o_files.append( os.path.join(params.outdir, fname + e) )
324        for f in o_files:  o_hdls.append( open(f, 'w') )
325    
326        parse_format = format_factory( o_hdls )
327        if not parse_format:
328            sys.stderr.write("Invalid format %s\n"%fformat)
329            sys.exit(1)
330    
331        if params.verbose: sys.stdout.write("Parsing file %s ...\n" %infile)
332        if fformat == 'kegg':
333            for rec in kegg_iter(infile):
334                parse_format(rec)
335        else:
336            if fformat == 'nr': fformat = 'fasta'
337            for rec in SeqIO.parse(infile, fformat):
338                parse_format(rec)
339    
340        for h in o_hdls: h.close()
341        return os.path.join(params.outdir, fname)
342    
343    
344  usage = "usage: %prog [options] source input_file1 [input_file2 input_file3 ...]" + __doc__  usage = "usage: %prog [options] source input_file1 [input_file2 input_file3 ...]" + __doc__
345    
346  def main(args):  def main(args):
347        global params
348      parser = OptionParser(usage=usage)      parser = OptionParser(usage=usage)
349      parser.add_option("-f", "--format", dest="format", metavar="FORMAT", default="fasta",      parser.add_option("-f", "--format", dest="format", metavar="FORMAT", default="fasta",
350                        help="FORMAT inputed file is in. Must be one of following: genbank, kegg, swiss, fasta [default is fasta]")                        help="FORMAT inputed file is in. Must be one of following: genbank, kegg, swiss, nr, fasta [default is 'fasta']")
351        parser.add_option("-n", "--nr_dbs", dest="nrdbs", metavar="DBs", default="ref,gb",
352                          help="Comma seperated list of databases to parse out of ncbi nr, from following: ref, gb, emb, dbj, pir, prf, pdb [default is 'ref,gb']")
353        parser.add_option("-p", "--processes", dest="processes", metavar="NUM_PROCESSES", type="int", default=4, help="Number of processes to use [default '4']")
354      parser.add_option("-k", "--kegg_map", dest="keggmap", metavar="FILE", default=None, help="Optional KEGG genome FILE for org name mapping")      parser.add_option("-k", "--kegg_map", dest="keggmap", metavar="FILE", default=None, help="Optional KEGG genome FILE for org name mapping")
355      parser.add_option("-e", "--eggnog_map", dest="eggnogmap", metavar="FILE", default=None, help="Optional eggNOG orthgroups tabed FILE for func mapping")      parser.add_option("-e", "--eggnog_map", dest="eggnogmap", metavar="FILE", default=None, help="Optional eggNOG orthgroups tabed FILE for func mapping")
356      parser.add_option("-d", "--out_dir", dest="outdir", metavar="DIR", default=None, help="DIR to write output [default is current dir]")      parser.add_option("-d", "--out_dir", dest="outdir", metavar="DIR", default="", help="DIR to write output [default is current dir]")
357      parser.add_option("-o", "--org_header", dest="orgheader", action="store_true", default=False, help="For fasta files, header description is organism [default is function]")      parser.add_option("-o", "--org_header", dest="orgheader", action="store_true", default=False, help="For fasta files, header description is organism [default is function]")
358      parser.add_option("-g", "--get_ontology", dest="getont", action="store_true", default=False, help="Output ontology (id, type) for proteins with mapped ontology [default is off]")      parser.add_option("-g", "--get_ontology", dest="getont", action="store_true", default=False, help="Output ontology (id, type) for proteins with mapped ontology [default is off]")
359      parser.add_option("-c", "--get_contig", dest="getcontig", action="store_true", default=False, help="Output contig info for organism genbank files [default is off]")      parser.add_option("-c", "--get_contig", dest="getcontig", action="store_true", default=False, help="Output contig info for organism genbank files [default is off]")
# Line 265  Line 363 
363      (opts, args) = parser.parse_args()      (opts, args) = parser.parse_args()
364      if len(args) < 2:      if len(args) < 2:
365          parser.error("Incorrect number of arguments")          parser.error("Incorrect number of arguments")
366      source = args[0]      sfiles = filter(lambda x: os.path.isfile(x), args[1:])
367      files  = args[1:]      scount = len(sfiles)
     outdir = ''  
     if opts.outdir:  
         outdir = opts.outdir  
368    
369      keggmap   = get_kegg_map(opts.keggmap, opts.verbose)      params.source  = args[0]
370      eggnogmap = get_eggnog_map(opts.eggnogmap, opts.verbose)      params.format  = opts.format
371        params.verbose = opts.verbose
372        params.outdir  = opts.outdir
373        params.orghead = opts.orgheader
374        params.getont  = opts.getont
375        params.getctg  = opts.getcontig
376        params.gettax  = opts.gettax
377    
378        if (params.format == 'nr') and opts.nrdbs:
379            for d in opts.nrdbs.split(','):
380                if d in nr_types:
381                    params.amap[d] = nr_types[d]
382        elif (params.format == 'kegg') and opts.keggmap and os.path.isfile(opts.keggmap):
383            params.amap = get_kegg_map(opts.keggmap)
384        elif (params.format == 'fasta') and opts.eggnogmap and os.path.isfile(opts.eggnogmap):
385            params.amap = get_eggnog_map(opts.eggnogmap)
386    
387      o_files = []      if scount < opts.processes:
388      o_hdls  = []          min_proc = scount
389        else:
390      for e in file_ext:          min_proc = opts.processes
         o_files.append( os.path.join(outdir, source + e) )  
     for f in o_files:  
         o_hdls.append( open(f, 'w') )  
   
     parse_format = format_factory( o_hdls, source, opts.format, opts.orgheader, opts.getont, opts.getcontig, opts.gettax )  
     if not parse_format:  
         parser.error("Incorrect format entered")  
391    
392      for f in files:      if min_proc == 1:
393          if opts.verbose:          if params.verbose: sys.stdout.write("Parsing %d %s files, single threaded\n"%(scount, params.format))
394              sys.stdout.write("Parsing file %s ...\n" %f)          rfiles = []
395          if (opts.format == 'kegg'):          for f in sfiles:
396              for rec in kegg_iter(f):              r = process_file(f)
397                  parse_format(rec, keggmap)              rfiles.append(r)
398          else:          else:
399              for rec in SeqIO.parse(f, opts.format):          if params.verbose: sys.stdout.write("Parsing %d %s files using %d threades\n"%(scount, params.format, min_proc))
400                  parse_format(rec, eggnogmap)          pool   = Pool(processes=min_proc)
401            rfiles = pool.map(process_file, sfiles, 1)
402            pool.close()
403            pool.join()
404        if params.verbose: sys.stdout.write("Done\n")
405    
406      for h in o_hdls:      if params.verbose: sys.stdout.write("Merging %d files ... "%len(rfiles))
407          h.close()      for e in file_ext:
408      for f in o_files:          out_file = os.path.join(params.outdir, params.source + e)
409          if os.path.isfile(f) and (os.path.getsize(f) == 0):          file_set = map( lambda x: x + e, rfiles )
410              if opts.verbose:          if os.path.isfile(out_file):
411                  sys.stdout.write("Removing empty file %s\n" %f)              os.remove(out_file)
412            os.system( "cat %s | sort -u > %s"%(" ".join(file_set), out_file) )
413            for f in file_set:
414              os.remove(f)              os.remove(f)
415            if os.path.isfile(out_file) and (os.path.getsize(out_file) == 0):
416                os.remove(out_file)
417        if params.verbose: sys.stdout.write("Done\n")
418    
419      if opts.verbose:      if params.verbose: sys.stdout.write("All Done\n")
         sys.stdout.write("All Done\n")  
420    
421  if __name__ == "__main__":  if __name__ == "__main__":
422      sys.exit(main(sys.argv))      sys.exit(main(sys.argv))

Legend:
Removed from v.1.9  
changed lines
  Added in v.1.10

MCS Webmaster
ViewVC Help
Powered by ViewVC 1.0.3