<s> import pdb <EOL> import os , sys <EOL> import itertools <EOL> from cPickle import * <EOL> from collections import defaultdict , namedtuple <EOL> from pbtools . pbtranscript . Utils import check_ids_unique <EOL> import pbtools . pbtranscript . tofu_wrap as tofu_wrap <EOL> import pbtools . pbtranscript . BioReaders as BioReaders <EOL> import pbtools . pbtranscript . branch . branch_simple2 as branch_simple2 <EOL> import pbtools . pbtranscript . counting . compare_junctions as compare_junctions <EOL> from pbtools . pbtranscript . io . SeqReaders import LazyFastaReader , LazyFastqReader <EOL> from pbcore . io . FastaIO import FastaWriter <EOL> from pbcore . io . FastqIO import FastqWriter <EOL> from bx . intervals . cluster import ClusterTree <EOL> def pick_rep ( fa_fq_filename , sam_filename , gff_filename , group_filename , output_filename , is_fq = False , pick_least_err_instead = False ) : <EOL> """<STR_LIT>""" <EOL> if is_fq : <EOL> fd = LazyFastqReader ( fa_fq_filename ) <EOL> fout = FastqWriter ( output_filename ) <EOL> else : <EOL> fd = LazyFastaReader ( fa_fq_filename ) <EOL> fout = FastaWriter ( output_filename ) <EOL> rep_info = { } <EOL> id_to_rep = { } <EOL> for line in open ( group_filename ) : <EOL> pb_id , members = line . strip ( ) . split ( '<STR_LIT:\t>' ) <EOL> print >> sys . stderr , "<STR_LIT>" , pb_id <EOL> best_id = None <EOL> best_seq = None <EOL> best_qual = None <EOL> best_err = <NUM_LIT> <EOL> err = <NUM_LIT> <EOL> max_len = <NUM_LIT:0> <EOL> for x in members . split ( '<STR_LIT:U+002C>' ) : <EOL> if is_fq and pick_least_err_instead : <EOL> err = sum ( i ** - ( i / <NUM_LIT> ) for i in fd [ x ] . quality ) <EOL> if ( is_fq and pick_least_err_instead and err < best_err ) or ( ( not is_fq or not pick_least_err_instead ) and len ( fd [ x ] . sequence ) >= max_len ) : <EOL> best_id = x <EOL> best_seq = fd [ x ] . sequence <EOL> if is_fq : <EOL> best_qual = fd [ x ] . quality <EOL> best_err = err <EOL> max_len = len ( fd [ x ] . sequence ) <EOL> rep_info [ pb_id ] = ( best_id , best_seq , best_qual ) <EOL> id_to_rep [ best_id ] = pb_id <EOL> f_gff = open ( gff_filename , '<STR_LIT:w>' ) <EOL> coords = { } <EOL> record_storage = { } <EOL> for r in BioReaders . GMAPSAMReader ( sam_filename , True ) : <EOL> if r . qID in id_to_rep : <EOL> pb_id = id_to_rep [ r . qID ] <EOL> best_id , best_seq , best_qual = rep_info [ pb_id ] <EOL> if r . qID not in coords : <EOL> coords [ r . qID ] = "<STR_LIT>" . format ( r . sID , r . sStart , r . sEnd , r . flag . strand ) <EOL> isoform_index = <NUM_LIT:1> <EOL> record_storage [ pb_id ] = r <EOL> else : <EOL> coords [ r . qID ] += "<STR_LIT>" . format ( r . sID , r . sStart , r . sEnd , r . flag . strand ) <EOL> isoform_index = <NUM_LIT:1> <EOL> old_r = record_storage [ pb_id ] <EOL> f_gff . write ( "<STR_LIT>" . format ( chr = old_r . sID , s = old_r . segments [ <NUM_LIT:0> ] . start + <NUM_LIT:1> , e = old_r . segments [ - <NUM_LIT:1> ] . end , pi = pb_id , j = isoform_index , strand = old_r . flag . strand ) ) <EOL> for s in old_r . segments : <EOL> f_gff . write ( "<STR_LIT>" . format ( chr = old_r . sID , s = s . start + <NUM_LIT:1> , e = s . end , pi = pb_id , j = isoform_index , strand = old_r . flag . strand ) ) <EOL> isoform_index = <NUM_LIT:2> <EOL> f_gff . write ( "<STR_LIT>" . format ( chr = r . sID , s = r . segments [ <NUM_LIT:0> ] . start + <NUM_LIT:1> , e = r . segments [ - <NUM_LIT:1> ] . end , pi = pb_id , j = isoform_index , strand = r . flag . strand ) ) <EOL> for s in r . segments : <EOL> f_gff . write ( "<STR_LIT>" . format ( chr = r . sID , s = s . start + <NUM_LIT:1> , e = s . end , pi = pb_id , j = isoform_index , strand = r . flag . strand ) ) <EOL> f_gff . close ( ) <EOL> for pb_id in rep_info : <EOL> best_id , best_seq , best_qual = rep_info [ pb_id ] <EOL> _id_ = "<STR_LIT>" . format ( pb_id , coords [ best_id ] , best_id ) <EOL> _seq_ = best_seq <EOL> if is_fq : <EOL> fout . writeRecord ( _id_ , _seq_ , best_qual ) <EOL> else : <EOL> fout . writeRecord ( _id_ , _seq_ ) <EOL> def sep_by_strand ( records ) : <EOL> output = { '<STR_LIT:+>' : [ ] , '<STR_LIT:->' : [ ] } <EOL> for r in records : <EOL> output [ r . flag . strand ] . append ( r ) <EOL> return output <EOL> def is_fusion_compatible ( r1 , r2 , max_fusion_point_dist , max_exon_end_dist , allow_extra_5_exons ) : <EOL> """<STR_LIT>""" <EOL> assert r1 . flag . strand == r2 . flag . strand <EOL> if r1 . qStart <= <NUM_LIT> * r1 . qLen : <EOL> if r2 . qStart > <NUM_LIT> * r2 . qLen : <EOL> return False <EOL> in_5_portion = True <EOL> else : <EOL> if r2 . qStart <= <NUM_LIT> * r2 . qLen : <EOL> return False <EOL> in_5_portion = False <EOL> plus_is_5end = ( r1 . flag . strand == '<STR_LIT:+>' ) <EOL> type = compare_junctions . compare_junctions ( r1 , r2 ) <EOL> if type == '<STR_LIT>' : <EOL> if len ( r1 . segments ) == <NUM_LIT:1> : <EOL> if len ( r2 . segments ) == <NUM_LIT:1> : <EOL> if in_5_portion and plus_is_5end : dist = abs ( r1 . sStart - r2 . sStart ) <EOL> else : dist = abs ( r1 . sEnd - r2 . sEnd ) <EOL> return dist <= max_fusion_point_dist <EOL> else : <EOL> raise Exception , "<STR_LIT>" + "<STR_LIT>" <EOL> else : <EOL> return True <EOL> elif type == '<STR_LIT>' or type == '<STR_LIT>' : <EOL> if allow_extra_5_exons : <EOL> if in_5_portion and plus_is_5end : <EOL> if abs ( r1 . segments [ - <NUM_LIT:1> ] . start - r2 . segments [ - <NUM_LIT:1> ] . start ) > max_exon_end_dist : return False <EOL> if abs ( r1 . segments [ - <NUM_LIT:1> ] . end - r2 . segments [ - <NUM_LIT:1> ] . end ) > max_fusion_point_dist : return False <EOL> return True <EOL> elif in_5_portion and ( not plus_is_5end ) : <EOL> if abs ( r1 . segments [ <NUM_LIT:0> ] . end - r2 . segments [ <NUM_LIT:0> ] . end ) > max_exon_end_dist : return False <EOL> if abs ( r1 . segments [ <NUM_LIT:0> ] . start - r2 . segments [ <NUM_LIT:0> ] . start ) > max_fusion_point_dist : return False <EOL> return True <EOL> else : <EOL> return False <EOL> else : <EOL> return False <EOL> else : <EOL> return False <EOL> def merge_fusion_exons ( records , max_fusion_point_dist , max_exon_end_dist , allow_extra_5_exons ) : <EOL> """<STR_LIT>""" <EOL> output = [ [ records [ <NUM_LIT:0> ] ] ] <EOL> for r1 in records [ <NUM_LIT:1> : ] : <EOL> merged = False <EOL> for i , r2s in enumerate ( output ) : <EOL> if all ( is_fusion_compatible ( r1 , r2 , max_fusion_point_dist , max_exon_end_dist , allow_extra_5_exons ) for r2 in r2s ) : <EOL> output [ i ] . append ( r1 ) <EOL> merged = True <EOL> break <EOL> if not merged : <EOL> output . append ( [ r1 ] ) <EOL> return output <EOL> def iter_gmap_sam_for_fusion ( gmap_sam_filename , fusion_candidates , transfrag_len_dict ) : <EOL> """<STR_LIT>""" <EOL> records = [ ] <EOL> iter = BioReaders . GMAPSAMReader ( gmap_sam_filename , True , query_len_dict = transfrag_len_dict ) <EOL> for r in iter : <EOL> if r . qID in fusion_candidates : <EOL> records = [ r ] <EOL> break <EOL> for r in iter : <EOL> if len ( records ) >= <NUM_LIT:1> and ( r . sID == records [ - <NUM_LIT:1> ] . sID and r . sStart < records [ - <NUM_LIT:1> ] . sStart ) : <EOL> print >> sys . stderr , "<STR_LIT>" <EOL> sys . exit ( - <NUM_LIT:1> ) <EOL> if len ( records ) >= <NUM_LIT:1> and ( r . sID != records [ <NUM_LIT:0> ] . sID or r . sStart > records [ - <NUM_LIT:1> ] . sEnd ) : <EOL> yield ( sep_by_strand ( records ) ) <EOL> records = [ ] <EOL> if r . qID in fusion_candidates : <EOL> records . append ( r ) <EOL> if len ( records ) > <NUM_LIT:0> : <EOL> yield ( sep_by_strand ( records ) ) <EOL> def find_fusion_candidates ( sam_filename , query_len_dict , min_locus_coverage = <NUM_LIT> , min_locus_coverage_bp = <NUM_LIT:1> , min_total_coverage = <NUM_LIT> , min_dist_between_loci = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> TmpRec = namedtuple ( '<STR_LIT>' , [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> def total_coverage ( tmprecs ) : <EOL> tree = ClusterTree ( <NUM_LIT:0> , <NUM_LIT:0> ) <EOL> for r in tmprecs : tree . insert ( r . qStart , r . qEnd , - <NUM_LIT:1> ) <EOL> return sum ( reg [ <NUM_LIT:1> ] - reg [ <NUM_LIT:0> ] for reg in tree . getregions ( ) ) <EOL> d = defaultdict ( lambda : [ ] ) <EOL> reader = BioReaders . GMAPSAMReader ( sam_filename , True , query_len_dict = query_len_dict ) <EOL> for r in reader : <EOL> if r . sID == '<STR_LIT:*>' : continue <EOL> if r . flag . strand == '<STR_LIT:+>' : <EOL> d [ r . qID ] . append ( TmpRec ( qCov = r . qCoverage , qLen = r . qLen , qStart = r . qStart , qEnd = r . qEnd , sStart = r . sStart , sEnd = r . sEnd , iden = r . identity ) ) <EOL> else : <EOL> d [ r . qID ] . append ( TmpRec ( qCov = r . qCoverage , qLen = r . qLen , qStart = r . qLen - r . qEnd , qEnd = r . qLen - r . qStart , sStart = r . sStart , sEnd = r . sEnd , iden = r . identity ) ) <EOL> fusion_candidates = [ ] <EOL> for k , data in d . iteritems ( ) : <EOL> if len ( data ) > <NUM_LIT:1> and all ( a . iden >= <NUM_LIT> for a in data ) and all ( a . qCov >= min_locus_coverage for a in data ) and all ( a . qCov * a . qLen >= min_locus_coverage_bp for a in data ) and total_coverage ( data ) * <NUM_LIT:1.> / data [ <NUM_LIT:0> ] . qLen >= min_total_coverage and all ( max ( a . sStart , b . sStart ) - min ( a . sEnd , b . sEnd ) >= min_dist_between_loci for a , b in itertools . combinations ( data , <NUM_LIT:2> ) ) : <EOL> fusion_candidates . append ( k ) <EOL> return fusion_candidates <EOL> def fusion_main ( fa_or_fq_filename , sam_filename , output_prefix , is_fq = False , allow_extra_5_exons = True , skip_5_exon_alt = True , prefix_dict_pickle_filename = None , min_locus_coverage = <NUM_LIT> , min_total_coverage = <NUM_LIT> , min_locus_coverage_bp = <NUM_LIT:1> , min_dist_between_loci = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> compressed_records_pointer_dict = defaultdict ( lambda : [ ] ) <EOL> merged_exons = [ ] <EOL> merged_i = <NUM_LIT:0> <EOL> check_ids_unique ( fa_or_fq_filename , is_fq = is_fq ) <EOL> bs = branch_simple2 . BranchSimple ( fa_or_fq_filename , is_fq = is_fq ) <EOL> fusion_candidates = find_fusion_candidates ( sam_filename , bs . transfrag_len_dict , min_locus_coverage , min_locus_coverage_bp , min_total_coverage , min_dist_between_loci ) <EOL> for recs in iter_gmap_sam_for_fusion ( sam_filename , fusion_candidates , bs . transfrag_len_dict ) : <EOL> for v in recs . itervalues ( ) : <EOL> if len ( v ) > <NUM_LIT:0> : <EOL> o = merge_fusion_exons ( v , max_fusion_point_dist = <NUM_LIT:100> , max_exon_end_dist = <NUM_LIT:0> , allow_extra_5_exons = allow_extra_5_exons ) <EOL> for group in o : <EOL> merged_exons . append ( group ) <EOL> for r in group : compressed_records_pointer_dict [ r . qID ] . append ( merged_i ) <EOL> merged_i += <NUM_LIT:1> <EOL> f_group = open ( '<STR_LIT>' , '<STR_LIT:w>' ) <EOL> gene_index = <NUM_LIT:1> <EOL> already_seen = set ( ) <EOL> for qid , indices in compressed_records_pointer_dict . iteritems ( ) : <EOL> combo = tuple ( indices ) <EOL> if combo in already_seen : <EOL> print "<STR_LIT>" , combo <EOL> continue <EOL> already_seen . add ( combo ) <EOL> for isoform_index , i in enumerate ( indices ) : <EOL> bs . cuff_index = gene_index <EOL> records = merged_exons [ i ] <EOL> f_group . write ( "<STR_LIT>" . format ( p = "<STR_LIT>" , i = gene_index , j = isoform_index , ids = "<STR_LIT:U+002C>" . join ( r . qID for r in records ) ) ) <EOL> gene_index += <NUM_LIT:1> <EOL> f_group . close ( ) <EOL> f_group = open ( output_prefix + '<STR_LIT>' , '<STR_LIT:w>' ) <EOL> group_info = { } <EOL> count = <NUM_LIT:0> <EOL> with open ( '<STR_LIT>' ) as f : <EOL> while True : <EOL> line = f . readline ( ) . strip ( ) <EOL> if len ( line ) == <NUM_LIT:0> : break <EOL> pbid1 , groups1 = line . strip ( ) . split ( '<STR_LIT:\t>' ) <EOL> pbid2 , groups2 = f . readline ( ) . strip ( ) . split ( '<STR_LIT:\t>' ) <EOL> assert pbid1 . split ( '<STR_LIT:.>' ) [ <NUM_LIT:1> ] == pbid2 . split ( '<STR_LIT:.>' ) [ <NUM_LIT:1> ] <EOL> group = set ( groups1 . split ( '<STR_LIT:U+002C>' ) ) . intersection ( groups2 . split ( '<STR_LIT:U+002C>' ) ) <EOL> f_group . write ( "<STR_LIT>" . format ( pbid1 [ : pbid1 . rfind ( '<STR_LIT:.>' ) ] , "<STR_LIT:U+002C>" . join ( group ) ) ) <EOL> group_info [ pbid1 [ : pbid1 . rfind ( '<STR_LIT:.>' ) ] ] = list ( group ) <EOL> count += <NUM_LIT:1> <EOL> f_group . close ( ) <EOL> gff_filename = output_prefix + '<STR_LIT>' <EOL> group_filename = output_prefix + '<STR_LIT>' <EOL> if is_fq : <EOL> output_filename = output_prefix + '<STR_LIT>' <EOL> else : <EOL> output_filename = output_prefix + '<STR_LIT>' <EOL> pick_rep ( fa_or_fq_filename , sam_filename , gff_filename , group_filename , output_filename , is_fq = is_fq , pick_least_err_instead = False ) <EOL> print >> sys . stderr , "<STR_LIT>" . format ( count ) <EOL> print >> sys . stderr , "<STR_LIT>" . format ( output_prefix , output_filename ) <EOL> if prefix_dict_pickle_filename is not None : <EOL> with open ( prefix_dict_pickle_filename ) as f : <EOL> d = load ( f ) <EOL> d1 = d [ '<STR_LIT>' ] <EOL> d1 . update ( d [ '<STR_LIT>' ] ) <EOL> tofu_wrap . get_abundance ( output_prefix , d1 , output_prefix ) <EOL> print >> sys . stderr , "<STR_LIT>" . format ( output_prefix ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> from argparse import ArgumentParser <EOL> parser = ArgumentParser ( ) <EOL> parser . add_argument ( "<STR_LIT>" , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , default = False , action = "<STR_LIT:store_true>" , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , required = True , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , required = True , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , action = "<STR_LIT>" , dest = "<STR_LIT>" , default = True , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , default = None , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT:-c>" , "<STR_LIT>" , type = float , default = <NUM_LIT> , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , type = int , default = <NUM_LIT:1> , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , type = float , default = <NUM_LIT> , help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , type = int , default = <NUM_LIT> , help = "<STR_LIT>" ) <EOL> args = parser . parse_args ( ) <EOL> fusion_main ( args . input , args . sam , args . prefix , <EOL> is_fq = args . fq , allow_extra_5_exons = args . allow_extra_5exon , <EOL> skip_5_exon_alt = False , prefix_dict_pickle_filename = args . prefix_dict_pickle_filename , <EOL> min_locus_coverage = args . min_locus_coverage , min_locus_coverage_bp = args . min_locus_coverage_bp , <EOL> min_total_coverage = args . min_total_coverage , <EOL> min_dist_between_loci = args . min_dist_between_loci ) </s>
