<s> '''<STR_LIT>''' <EOL> import sys <EOL> import CGAT . Experiment as E <EOL> from CGAT . Diamond import * <EOL> import collections <EOL> import CGAT . IOTools as IOTools <EOL> def readCogMap ( cog_map ) : <EOL> '''<STR_LIT>''' <EOL> gene2cog = { } <EOL> for line in IOTools . openFile ( cog_map ) : <EOL> data = line [ : - <NUM_LIT:1> ] . split ( "<STR_LIT:\t>" ) <EOL> gene2cog [ data [ <NUM_LIT:0> ] ] = data [ <NUM_LIT:1> ] <EOL> return gene2cog <EOL> def main ( argv = None ) : <EOL> """<STR_LIT>""" <EOL> if argv is None : <EOL> argv = sys . argv <EOL> parser = E . OptionParser ( version = "<STR_LIT>" , <EOL> usage = globals ( ) [ "<STR_LIT>" ] ) <EOL> parser . add_option ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" , type = "<STR_LIT>" , <EOL> choices = ( "<STR_LIT>" , None ) , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_option ( "<STR_LIT>" , dest = "<STR_LIT>" , action = "<STR_LIT:store_true>" , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_option ( "<STR_LIT>" , <EOL> dest = "<STR_LIT>" , <EOL> action = "<STR_LIT:store_true>" , <EOL> help = """<STR_LIT>""" ) <EOL> parser . add_option ( "<STR_LIT>" , dest = "<STR_LIT>" , type = "<STR_LIT:string>" , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_option ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" , type = "<STR_LIT:int>" , <EOL> help = """<STR_LIT>""" ) <EOL> parser . set_defaults ( method = None , <EOL> sum_cog = False , <EOL> evaluate_cog = False , <EOL> cog_map = None , <EOL> nsamples = <NUM_LIT> ) <EOL> ( options , args ) = E . Start ( parser , argv = argv ) <EOL> if options . evaluate_cog : <EOL> assert options . cog_map , """<STR_LIT>""" <EOL> assert not options . method , """<STR_LIT>""" <EOL> E . info ( "<STR_LIT>" % options . cog_map ) <EOL> gene2cog = readCogMap ( options . cog_map ) <EOL> E . info ( "<STR_LIT>" ) <EOL> E . info ( "<STR_LIT>" ) <EOL> options . stdout . write ( "<STR_LIT>" ) <EOL> c = <NUM_LIT:0> <EOL> for alignments in query_iterator ( alignment_iterator ( options . stdin ) ) : <EOL> c += <NUM_LIT:1> <EOL> scores = [ ] <EOL> if c <= options . nsamples : <EOL> for alignment in alignments : <EOL> scores . append ( alignment . score ) <EOL> best = max ( scores ) <EOL> best_alignments = [ <EOL> x for x in alignments if x . score == best ] <EOL> if len ( best_alignments ) > <NUM_LIT:1> : <EOL> best_alignments = random . sample ( best_alignments , <NUM_LIT:1> ) <EOL> best_alignment = best_alignments [ <NUM_LIT:0> ] <EOL> best_cog = gene2cog [ best_alignment . ref ] <EOL> pbest = float ( len ( <EOL> [ gene2cog [ x . ref ] <EOL> for x in alignments <EOL> if gene2cog [ x . ref ] == best_cog ] ) ) / len ( alignments ) * <NUM_LIT:100> <EOL> nalignments = len ( alignments ) <EOL> options . stdout . write ( <EOL> "<STR_LIT:\t>" . join ( map ( <EOL> str , [ alignments [ <NUM_LIT:0> ] . qid , <EOL> pbest , <EOL> nalignments ] ) ) + "<STR_LIT:\n>" <EOL> ) <EOL> else : <EOL> break <EOL> return <EOL> counts = collections . defaultdict ( int ) <EOL> E . info ( "<STR_LIT>" ) <EOL> assert options . method , "<STR_LIT>" <EOL> if options . method == "<STR_LIT>" : <EOL> if options . sum_cog : <EOL> E . warn ( """<STR_LIT>""" ) <EOL> assert options . cog_map , """<STR_LIT>""" <EOL> E . info ( """<STR_LIT>""" <EOL> % options . cog_map ) <EOL> gene2cog = readCogMap ( options . cog_map ) <EOL> E . info ( "<STR_LIT>" ) <EOL> E . info ( "<STR_LIT>" ) <EOL> query_it = query_iterator ( alignment_iterator ( options . stdin ) ) <EOL> for best in best_alignment_iterator ( query_it ) : <EOL> cog = gene2cog [ best . ref ] <EOL> if cog == "<STR_LIT>" or cog . find ( "<STR_LIT:;>" ) != - <NUM_LIT:1> : <EOL> continue <EOL> counts [ cog ] += <NUM_LIT:1> <EOL> else : <EOL> E . info ( "<STR_LIT>" ) <EOL> query_it = query_iterator ( alignment_iterator ( options . stdin ) ) <EOL> for best in best_alignment_iterator ( query_it ) : <EOL> counts [ best . ref ] += <NUM_LIT:1> <EOL> E . info ( "<STR_LIT>" ) <EOL> E . info ( "<STR_LIT>" ) <EOL> options . stdout . write ( "<STR_LIT>" ) <EOL> for ref , count in counts . iteritems ( ) : <EOL> options . stdout . write ( "<STR_LIT:\t>" . join ( [ ref , str ( count ) ] ) + "<STR_LIT:\n>" ) <EOL> E . Stop ( ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> sys . exit ( main ( sys . argv ) ) </s>
