<s> import sys <EOL> import os <EOL> import re <EOL> import pysam <EOL> from eta import ETA <EOL> import ngsutils . support <EOL> def bam_open ( fname , mode = '<STR_LIT:r>' , * args , ** kwargs ) : <EOL> if fname . lower ( ) [ - <NUM_LIT:4> : ] == '<STR_LIT>' : <EOL> return pysam . Samfile ( fname , '<STR_LIT>' % mode , * args , ** kwargs ) <EOL> return pysam . Samfile ( fname , '<STR_LIT:%s>' % mode , * args , ** kwargs ) <EOL> def bam_pileup_iter ( bam , mask = <NUM_LIT> , quiet = False , callback = None ) : <EOL> if not quiet and bam . filename : <EOL> eta = ETA ( os . stat ( bam . filename ) . st_size ) <EOL> else : <EOL> eta = None <EOL> for pileup in bam . pileup ( mask = mask ) : <EOL> pos = bam . tell ( ) <EOL> bgz_offset = pos >> <NUM_LIT:16> <EOL> if not quiet : <EOL> if callback : <EOL> eta . print_status ( bgz_offset , extra = callback ( pileup ) ) <EOL> else : <EOL> eta . print_status ( bgz_offset , extra = '<STR_LIT>' % ( bam . getrname ( pileup . tid ) , pileup . pos ) ) <EOL> yield pileup <EOL> if eta : <EOL> eta . done ( ) <EOL> def bam_iter ( bam , quiet = False , show_ref_pos = False , ref = None , start = None , end = None , callback = None ) : <EOL> '''<STR_LIT>''' <EOL> if os . path . exists ( '<STR_LIT>' % bam . filename ) : <EOL> show_ref_pos = True <EOL> eta = None <EOL> if not ref : <EOL> if not quiet and bam . filename : <EOL> eta = ETA ( os . stat ( bam . filename ) . st_size ) <EOL> for read in bam : <EOL> pos = bam . tell ( ) <EOL> bgz_offset = pos >> <NUM_LIT:16> <EOL> if not quiet and eta : <EOL> if callback : <EOL> eta . print_status ( bgz_offset , extra = callback ( read ) ) <EOL> elif ( show_ref_pos ) : <EOL> if read . tid > - <NUM_LIT:1> : <EOL> eta . print_status ( bgz_offset , extra = '<STR_LIT>' % ( bam . getrname ( read . tid ) , read . pos , read . qname ) ) <EOL> else : <EOL> eta . print_status ( bgz_offset , extra = '<STR_LIT>' % ( read . qname ) ) <EOL> else : <EOL> eta . print_status ( bgz_offset , extra = '<STR_LIT:%s>' % read . qname ) <EOL> yield read <EOL> else : <EOL> working_chrom = None <EOL> if ref in bam . references : <EOL> working_chrom = ref <EOL> elif ref [ <NUM_LIT:0> : <NUM_LIT:3> ] == '<STR_LIT>' : <EOL> if ref [ <NUM_LIT:3> : ] in bam . references : <EOL> working_chrom = ref [ <NUM_LIT:3> : ] <EOL> if not working_chrom : <EOL> raise ValueError ( '<STR_LIT>' % ref ) <EOL> tid = bam . gettid ( working_chrom ) <EOL> if not start : <EOL> start = <NUM_LIT:0> <EOL> if not end : <EOL> end = bam . lengths [ tid ] <EOL> if not quiet and bam . filename : <EOL> eta = ETA ( end - start ) <EOL> for read in bam . fetch ( working_chrom , start , end ) : <EOL> if not quiet and eta : <EOL> if callback : <EOL> eta . print_status ( read . pos - start , extra = callback ( read ) ) <EOL> else : <EOL> eta . print_status ( read . pos - start , extra = '<STR_LIT>' % ( bam . getrname ( read . tid ) , read . pos , read . qname ) ) <EOL> yield read <EOL> if eta : <EOL> eta . done ( ) <EOL> def bam_batch_reads ( bam , quiet = False ) : <EOL> '''<STR_LIT>''' <EOL> reads = [ ] <EOL> last = None <EOL> for read in bam_iter ( bam , quiet = quiet ) : <EOL> if last and read . qname != last : <EOL> yield reads <EOL> reads = [ ] <EOL> last = read . qname <EOL> reads . append ( read ) <EOL> if reads : <EOL> yield reads <EOL> bam_cigar = [ '<STR_LIT:M>' , '<STR_LIT:I>' , '<STR_LIT:D>' , '<STR_LIT:N>' , '<STR_LIT:S>' , '<STR_LIT:H>' , '<STR_LIT:P>' , '<STR_LIT:=>' , '<STR_LIT:X>' ] <EOL> bam_cigar_op = { <EOL> '<STR_LIT:M>' : <NUM_LIT:0> , <EOL> '<STR_LIT:I>' : <NUM_LIT:1> , <EOL> '<STR_LIT:D>' : <NUM_LIT:2> , <EOL> '<STR_LIT:N>' : <NUM_LIT:3> , <EOL> '<STR_LIT:S>' : <NUM_LIT:4> , <EOL> '<STR_LIT:H>' : <NUM_LIT:5> , <EOL> '<STR_LIT:P>' : <NUM_LIT:6> , <EOL> '<STR_LIT:=>' : <NUM_LIT:7> , <EOL> '<STR_LIT:X>' : <NUM_LIT:8> , <EOL> } <EOL> def cigar_fromstr ( s ) : <EOL> '''<STR_LIT>''' <EOL> ret = [ ] <EOL> spl = re . split ( '<STR_LIT>' , s ) [ <NUM_LIT:1> : ] <EOL> for op , size in zip ( spl [ <NUM_LIT:1> : : <NUM_LIT:2> ] , spl [ : : <NUM_LIT:2> ] ) : <EOL> ret . append ( ( bam_cigar_op [ op ] , int ( size ) ) ) <EOL> return ret <EOL> def cigar_tostr ( cigar ) : <EOL> '''<STR_LIT>''' <EOL> s = '<STR_LIT>' <EOL> for op , size in cigar : <EOL> s += '<STR_LIT>' % ( size , bam_cigar [ op ] ) <EOL> return s <EOL> def cigar_read_len ( cigar ) : <EOL> '''<STR_LIT>''' <EOL> read_pos = <NUM_LIT:0> <EOL> for op , length in cigar : <EOL> if op == <NUM_LIT:0> : <EOL> read_pos += length <EOL> elif op == <NUM_LIT:1> : <EOL> read_pos += length <EOL> elif op == <NUM_LIT:2> : <EOL> pass <EOL> elif op == <NUM_LIT:3> : <EOL> pass <EOL> elif op == <NUM_LIT:4> : <EOL> read_pos += length <EOL> elif op == <NUM_LIT:5> : <EOL> pass <EOL> elif op == <NUM_LIT:7> : <EOL> read_pos += length <EOL> elif op == <NUM_LIT:8> : <EOL> read_pos += length <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % op ) <EOL> return read_pos <EOL> def read_calc_mismatches ( read ) : <EOL> inserts = <NUM_LIT:0> <EOL> deletions = <NUM_LIT:0> <EOL> indels = <NUM_LIT:0> <EOL> edits = int ( read . opt ( '<STR_LIT>' ) ) <EOL> for op , length in read . cigar : <EOL> if op == <NUM_LIT:1> : <EOL> inserts += length <EOL> indels += <NUM_LIT:1> <EOL> elif op == <NUM_LIT:2> : <EOL> deletions += length <EOL> indels += <NUM_LIT:1> <EOL> return edits - inserts - deletions + indels <EOL> def _extract_md_matches ( md , maxlength ) : <EOL> md_pos = <NUM_LIT:0> <EOL> while md and md_pos < maxlength : <EOL> tmp = '<STR_LIT:0>' <EOL> while md and md [ <NUM_LIT:0> ] in '<STR_LIT>' : <EOL> tmp += md [ <NUM_LIT:0> ] <EOL> md = md [ <NUM_LIT:1> : ] <EOL> pos = int ( tmp ) <EOL> if pos > maxlength : <EOL> return ( maxlength , '<STR_LIT>' % ( pos - maxlength , md ) ) <EOL> return ( pos , md ) <EOL> def read_calc_variations ( read ) : <EOL> '<STR_LIT>' <EOL> for tup in _read_calc_variations ( read . pos , read . cigar , read . opt ( '<STR_LIT>' ) , read . seq ) : <EOL> yield tup <EOL> def _read_calc_variations ( start_pos , cigar , md , seq ) : <EOL> '''<STR_LIT>''' <EOL> ref_pos = start_pos <EOL> read_pos = <NUM_LIT:0> <EOL> for op , length in cigar : <EOL> if md and md [ <NUM_LIT:0> ] == '<STR_LIT:0>' : <EOL> md = md [ <NUM_LIT:1> : ] <EOL> if op == <NUM_LIT:0> : <EOL> md_pos = <NUM_LIT:0> <EOL> last = None <EOL> while md and md_pos < length : <EOL> if last == ( op , length , md ) : <EOL> sys . stderr . write ( '<STR_LIT>' % ( ref_pos , op , length ) ) <EOL> sys . exit ( <NUM_LIT:1> ) <EOL> last = ( op , length , md ) <EOL> chunk_size , md = _extract_md_matches ( md , length - md_pos ) <EOL> md_pos += chunk_size <EOL> while md_pos < length and md and md [ <NUM_LIT:0> ] not in '<STR_LIT>' : <EOL> yield ( op , ref_pos + md_pos , seq [ read_pos + md_pos ] ) <EOL> md = md [ <NUM_LIT:1> : ] <EOL> md_pos += <NUM_LIT:1> <EOL> ref_pos += length <EOL> read_pos += length <EOL> elif op == <NUM_LIT:1> : <EOL> yield ( op , ref_pos , seq [ read_pos : read_pos + length ] ) <EOL> read_pos += length <EOL> elif op == <NUM_LIT:2> : <EOL> if md [ <NUM_LIT:0> ] == '<STR_LIT>' : <EOL> md = md [ <NUM_LIT:1> : ] <EOL> yield ( op , ref_pos , md [ : length ] ) <EOL> md = md [ length : ] <EOL> ref_pos += length <EOL> elif op == <NUM_LIT:3> : <EOL> ref_pos += length <EOL> def read_calc_mismatches_gen ( ref , read , chrom ) : <EOL> start = read . pos <EOL> ref_pos = <NUM_LIT:0> <EOL> read_pos = <NUM_LIT:0> <EOL> for op , length in read . cigar : <EOL> if op == <NUM_LIT:1> : <EOL> yield ref_pos , op , None <EOL> read_pos += length <EOL> elif op == <NUM_LIT:2> : <EOL> yield ref_pos , op , None <EOL> ref_pos += length <EOL> elif op == <NUM_LIT:3> : <EOL> ref_pos += length <EOL> elif op == <NUM_LIT:0> : <EOL> refseq = ref . fetch ( chrom , start + ref_pos , start + ref_pos + length ) <EOL> if not refseq : <EOL> raise ValueError ( "<STR_LIT>" % ( chrom , ref . filename ) ) <EOL> cur_pos = start + ref_pos <EOL> for refbase , readbase in zip ( refseq . upper ( ) , read . seq [ read_pos : read_pos + length ] . upper ( ) ) : <EOL> if refbase != readbase : <EOL> yield op , cur_pos , readbase <EOL> cur_pos += <NUM_LIT:1> <EOL> ref_pos += length <EOL> read_pos += length <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % op ) <EOL> def read_calc_mismatches_ref ( ref , read , chrom ) : <EOL> edits = <NUM_LIT:0> <EOL> for op , pos , base in read_calc_mismatches_gen ( ref , read , chrom ) : <EOL> edits += <NUM_LIT:1> <EOL> return edits <EOL> __region_cache = { } <EOL> def region_pos_to_genomic_pos ( name , start , cigar ) : <EOL> '''<STR_LIT>''' <EOL> if name in __region_cache : <EOL> chrom , fragments = __region_cache [ name ] <EOL> else : <EOL> c1 = name . split ( '<STR_LIT::>' ) <EOL> chrom = c1 [ <NUM_LIT:0> ] <EOL> fragments = [ ] <EOL> for fragment in c1 [ <NUM_LIT:1> ] . split ( '<STR_LIT:U+002C>' ) : <EOL> s , e = fragment . split ( '<STR_LIT:->' ) <EOL> fragments . append ( ( int ( s ) , int ( e ) ) ) <EOL> __region_cache [ name ] = ( chrom , fragments ) <EOL> chr_cigar = [ ] <EOL> chr_start = fragments [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] <EOL> read_start = int ( start ) <EOL> frag_idx = <NUM_LIT:0> <EOL> frag_start = <NUM_LIT:0> <EOL> frag_end = <NUM_LIT:0> <EOL> for i , ( s , e ) in enumerate ( fragments ) : <EOL> if chr_start + read_start < e : <EOL> chr_start += read_start <EOL> frag_idx = i <EOL> frag_start = s <EOL> frag_end = e <EOL> break <EOL> else : <EOL> chr_start += ( e - s ) <EOL> read_start -= ( e - s ) <EOL> cur_pos = chr_start <EOL> for op , length in cigar : <EOL> if op in [ <NUM_LIT:1> , <NUM_LIT:4> , <NUM_LIT:5> ] : <EOL> chr_cigar . append ( ( op , length ) ) <EOL> elif op in [ <NUM_LIT:0> , <NUM_LIT:2> ] : <EOL> if cur_pos + length <= frag_end : <EOL> cur_pos += length <EOL> chr_cigar . append ( ( op , length ) ) <EOL> else : <EOL> while cur_pos + length > frag_end : <EOL> if frag_end - cur_pos > <NUM_LIT:0> : <EOL> chr_cigar . append ( ( op , frag_end - cur_pos ) ) <EOL> length -= ( frag_end - cur_pos ) <EOL> cur_pos = frag_end <EOL> frag_idx += <NUM_LIT:1> <EOL> if len ( fragments ) <= frag_idx : <EOL> print '<STR_LIT>' , name , fragments <EOL> return ( chrom , <NUM_LIT:0> , chr_cigar ) <EOL> frag_start , frag_end = fragments [ frag_idx ] <EOL> chr_cigar . append ( ( <NUM_LIT:3> , frag_start - cur_pos ) ) <EOL> cur_pos = frag_start <EOL> cur_pos = cur_pos + length <EOL> chr_cigar . append ( ( op , length ) ) <EOL> else : <EOL> print "<STR_LIT>" % bam_cigar [ op ] <EOL> sys . exit ( <NUM_LIT:1> ) <EOL> return ( chrom , chr_start , chr_cigar ) <EOL> def is_junction_valid ( cigar , min_overlap = <NUM_LIT:4> ) : <EOL> '''<STR_LIT>''' <EOL> first = True <EOL> pre_gap = True <EOL> pre_gap_count = <NUM_LIT:0> <EOL> post_gap_count = <NUM_LIT:0> <EOL> has_gap = False <EOL> for op , length in cigar : <EOL> if first and op == <NUM_LIT:3> : <EOL> return ( False , '<STR_LIT>' % cigar_tostr ( cigar ) ) <EOL> first = False <EOL> if op == <NUM_LIT:3> : <EOL> pre_gap = False <EOL> post_gap_count = <NUM_LIT:0> <EOL> has_gap = True <EOL> elif pre_gap : <EOL> pre_gap_count += length <EOL> else : <EOL> post_gap_count += length <EOL> if not has_gap : <EOL> return ( False , "<STR_LIT>" ) <EOL> elif pre_gap_count < min_overlap : <EOL> return ( False , "<STR_LIT>" % cigar_tostr ( cigar ) ) <EOL> elif post_gap_count < min_overlap : <EOL> return ( False , "<STR_LIT>" % cigar_tostr ( cigar ) ) <EOL> return True , '<STR_LIT>' <EOL> def read_alignment_fragments_gen ( read ) : <EOL> '''<STR_LIT>''' <EOL> for start , end in _read_alignment_fragments_gen ( read . pos , read . cigar ) : <EOL> yield ( start , end ) <EOL> def _read_alignment_fragments_gen ( pos , cigar ) : <EOL> '''<STR_LIT>''' <EOL> ref_pos = pos <EOL> for op , length in cigar : <EOL> if op == <NUM_LIT:0> : <EOL> yield ( ref_pos , ref_pos + length ) <EOL> ref_pos += length <EOL> elif op == <NUM_LIT:1> : <EOL> pass <EOL> elif op == <NUM_LIT:2> : <EOL> ref_pos += length <EOL> elif op == <NUM_LIT:3> : <EOL> ref_pos += length <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % op ) <EOL> def read_cigar_at_pos ( cigar , qpos , is_del ) : <EOL> '''<STR_LIT>''' <EOL> pos = <NUM_LIT:0> <EOL> returnnext = False <EOL> for op , length in cigar : <EOL> if returnnext : <EOL> return op <EOL> if op == <NUM_LIT:0> : <EOL> pos += length <EOL> elif op == <NUM_LIT:1> : <EOL> pos += length <EOL> elif op == <NUM_LIT:2> : <EOL> pass <EOL> elif op == <NUM_LIT:3> : <EOL> pass <EOL> elif op == <NUM_LIT:4> : <EOL> pos += length <EOL> elif op == <NUM_LIT:5> : <EOL> pass <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % op ) <EOL> if pos > qpos : <EOL> return op <EOL> if is_del and pos == qpos : <EOL> returnnext = True <EOL> return None <EOL> def cleancigar ( cigar ) : <EOL> '''<STR_LIT>''' <EOL> newcigar = [ ] <EOL> changed = False <EOL> zero = False <EOL> for op , size in cigar : <EOL> if size > <NUM_LIT:0> : <EOL> if zero and newcigar : <EOL> if newcigar [ - <NUM_LIT:1> ] [ <NUM_LIT:0> ] == op : <EOL> newsize = newcigar [ - <NUM_LIT:1> ] [ <NUM_LIT:1> ] + size <EOL> newcigar [ - <NUM_LIT:1> ] = ( op , newsize ) <EOL> zero = <NUM_LIT:0> <EOL> else : <EOL> newcigar . append ( ( op , size ) ) <EOL> else : <EOL> changed = True <EOL> zero = True <EOL> if changed : <EOL> return newcigar <EOL> return None <EOL> def read_cleancigar ( read ) : <EOL> '''<STR_LIT>''' <EOL> if read . is_unmapped : <EOL> return False <EOL> newcigar = [ ] <EOL> newcigar = cleancigar ( read . cigar ) <EOL> if newcigar : <EOL> read . cigar = newcigar <EOL> return True <EOL> return False <EOL> def read_to_unmapped ( read , ref = None ) : <EOL> '''<STR_LIT>''' <EOL> newread = pysam . AlignedRead ( ) <EOL> if ref : <EOL> tags = [ ( '<STR_LIT>' , '<STR_LIT>' % ( ref , read . pos , cigar_tostr ( read . cigar ) ) ) ] <EOL> newread . is_unmapped = True <EOL> newread . mapq = <NUM_LIT:0> <EOL> newread . tlen = <NUM_LIT:0> <EOL> newread . pos = - <NUM_LIT:1> <EOL> newread . pnext = - <NUM_LIT:1> <EOL> newread . rnext = - <NUM_LIT:1> <EOL> newread . tid = - <NUM_LIT:1> <EOL> newread . qname = read . qname <EOL> if read . is_paired : <EOL> newread . is_paired = True <EOL> if not read . is_unmapped and read . is_reverse : <EOL> newread . seq = ngsutils . support . revcomp ( read . seq ) <EOL> newread . qual = read . qual [ : : - <NUM_LIT:1> ] <EOL> else : <EOL> newread . seq = read . seq <EOL> newread . qual = read . qual <EOL> newread . tags = tags <EOL> return newread <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> import doctest <EOL> doctest . testmod ( ) </s>
