<s> """<STR_LIT>""" <EOL> import sys <EOL> import logging <EOL> from itertools import combinations <EOL> from jcvi . formats . base import BaseFile , LineFile , must_open , read_block <EOL> from jcvi . formats . bed import Bed , fastaFromBed <EOL> from jcvi . utils . cbook import memoized <EOL> from jcvi . utils . counter import Counter <EOL> from jcvi . apps . base import OptionParser , ActionDispatcher , need_update <EOL> MSTheader = """<STR_LIT>""" <EOL> class BinMap ( BaseFile , dict ) : <EOL> def __init__ ( self , filename ) : <EOL> super ( BinMap , self ) . __init__ ( filename ) <EOL> fp = open ( filename ) <EOL> for header , seq in read_block ( fp , "<STR_LIT>" ) : <EOL> lg = header . split ( ) [ - <NUM_LIT:1> ] <EOL> self [ lg ] = [ ] <EOL> for s in seq : <EOL> if s . strip ( ) == '<STR_LIT>' or s [ <NUM_LIT:0> ] == '<STR_LIT:;>' : <EOL> continue <EOL> marker , pos = s . split ( ) <EOL> pos = int ( float ( pos ) * <NUM_LIT:1000> ) <EOL> self [ lg ] . append ( ( marker , pos ) ) <EOL> def print_to_bed ( self , filename = "<STR_LIT>" , switch = False ) : <EOL> fw = must_open ( filename , "<STR_LIT:w>" ) <EOL> for lg , markers in sorted ( self . items ( ) ) : <EOL> for marker , pos in markers : <EOL> if not switch : <EOL> line = ( lg , pos , pos + <NUM_LIT:1> , marker ) <EOL> else : <EOL> seqid , spos = marker . split ( "<STR_LIT:.>" ) <EOL> spos = int ( spos ) <EOL> marker = "<STR_LIT>" . format ( lg , pos / <NUM_LIT> ) <EOL> line = ( seqid , spos - <NUM_LIT:1> , spos , marker ) <EOL> print >> fw , "<STR_LIT:\t>" . join ( str ( x ) for x in line ) <EOL> fw . close ( ) <EOL> class MSTMapLine ( object ) : <EOL> def __init__ ( self , row , startidx = <NUM_LIT:3> ) : <EOL> args = row . split ( ) <EOL> self . id = args [ <NUM_LIT:0> ] <EOL> self . seqid , pos = self . id . split ( "<STR_LIT:.>" ) <EOL> self . pos = int ( pos ) <EOL> self . genotype = "<STR_LIT>" . join ( args [ startidx : ] ) <EOL> def __len__ ( self ) : <EOL> return len ( self . genotype ) <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" . format ( self . id , self . genotype ) <EOL> @ property <EOL> def bedline ( self ) : <EOL> return "<STR_LIT:\t>" . join ( str ( x ) for x in ( self . seqid , self . pos - <NUM_LIT:1> , self . pos , self . id ) ) <EOL> class MSTMap ( LineFile ) : <EOL> def __init__ ( self , filename ) : <EOL> super ( MSTMap , self ) . __init__ ( filename ) <EOL> fp = open ( filename ) <EOL> startidx = <NUM_LIT:1> <EOL> for row in fp : <EOL> if row . startswith ( "<STR_LIT>" ) : <EOL> if row . split ( ) [ <NUM_LIT:1> ] == "<STR_LIT>" : <EOL> startidx = <NUM_LIT:3> <EOL> self . header = row . split ( ) <EOL> break <EOL> for row in fp : <EOL> self . append ( MSTMapLine ( row , startidx = startidx ) ) <EOL> self . nmarkers = len ( self ) <EOL> self . nind = len ( self [ <NUM_LIT:0> ] . genotype ) <EOL> logging . debug ( "<STR_LIT>" . format ( self . nmarkers , self . nind ) ) <EOL> class MSTMatrix ( object ) : <EOL> def __init__ ( self , matrix , markerheader , population_type , missing_threshold ) : <EOL> self . matrix = matrix <EOL> self . markerheader = markerheader <EOL> self . population_type = population_type <EOL> self . missing_threshold = missing_threshold <EOL> self . ngenotypes = len ( matrix ) <EOL> self . nind = len ( markerheader ) - <NUM_LIT:1> <EOL> assert self . nind == len ( matrix [ <NUM_LIT:0> ] ) - <NUM_LIT:1> <EOL> logging . debug ( "<STR_LIT>" . format ( self . ngenotypes , self . nind ) ) <EOL> def write ( self , filename = "<STR_LIT>" , header = True ) : <EOL> fw = must_open ( filename , "<STR_LIT:w>" ) <EOL> if header : <EOL> print >> fw , MSTheader . format ( self . population_type , <EOL> self . missing_threshold , self . ngenotypes , self . nind ) <EOL> print >> fw , "<STR_LIT:\t>" . join ( self . markerheader ) <EOL> for m in self . matrix : <EOL> print >> fw , "<STR_LIT:\t>" . join ( m ) <EOL> def main ( ) : <EOL> actions = ( <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ) <EOL> p = ActionDispatcher ( actions ) <EOL> p . dispatch ( globals ( ) ) <EOL> @ memoized <EOL> def calc_ldscore ( a , b ) : <EOL> assert len ( a ) == len ( b ) , "<STR_LIT>" . format ( a , b ) <EOL> c = Counter ( zip ( a , b ) ) <EOL> c_aa = c [ ( '<STR_LIT:A>' , '<STR_LIT:A>' ) ] <EOL> c_ab = c [ ( '<STR_LIT:A>' , '<STR_LIT:B>' ) ] <EOL> c_ba = c [ ( '<STR_LIT:B>' , '<STR_LIT:A>' ) ] <EOL> c_bb = c [ ( '<STR_LIT:B>' , '<STR_LIT:B>' ) ] <EOL> n = c_aa + c_ab + c_ba + c_bb <EOL> if n == <NUM_LIT:0> : <EOL> return <NUM_LIT:0> <EOL> f = <NUM_LIT:1.> / n <EOL> x_aa = c_aa * f <EOL> x_ab = c_ab * f <EOL> x_ba = c_ba * f <EOL> x_bb = c_bb * f <EOL> p_a = x_aa + x_ab <EOL> p_b = x_ba + x_bb <EOL> q_a = x_aa + x_ba <EOL> q_b = x_ab + x_bb <EOL> D = x_aa - p_a * q_a <EOL> denominator = p_a * p_b * q_a * q_b <EOL> if denominator == <NUM_LIT:0> : <EOL> return <NUM_LIT:0> <EOL> r2 = D * D / denominator <EOL> return r2 <EOL> def ld ( args ) : <EOL> """<STR_LIT>""" <EOL> import numpy as np <EOL> from random import sample <EOL> from jcvi . algorithms . matrix import symmetrize <EOL> p = OptionParser ( ld . __doc__ ) <EOL> p . add_option ( "<STR_LIT>" , default = <NUM_LIT:1000> , type = "<STR_LIT:int>" , <EOL> help = "<STR_LIT>" ) <EOL> opts , args , iopts = p . set_image_options ( args , figsize = "<STR_LIT>" ) <EOL> if len ( args ) != <NUM_LIT:1> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mstmap , = args <EOL> subsample = opts . subsample <EOL> data = MSTMap ( mstmap ) <EOL> markerbedfile = mstmap + "<STR_LIT>" <EOL> ldmatrix = mstmap + "<STR_LIT>" <EOL> if subsample < data . nmarkers : <EOL> data = [ data [ x ] for x in sorted ( sample ( xrange ( len ( data ) ) , subsample ) ) ] <EOL> else : <EOL> logging . debug ( "<STR_LIT>" ) <EOL> nmarkers = len ( data ) <EOL> if need_update ( mstmap , ( ldmatrix , markerbedfile ) ) : <EOL> fw = open ( markerbedfile , "<STR_LIT:w>" ) <EOL> print >> fw , "<STR_LIT:\n>" . join ( x . bedline for x in data ) <EOL> logging . debug ( "<STR_LIT>" . format ( nmarkers , markerbedfile ) ) <EOL> fw . close ( ) <EOL> M = np . zeros ( ( nmarkers , nmarkers ) , dtype = float ) <EOL> for i , j in combinations ( range ( nmarkers ) , <NUM_LIT:2> ) : <EOL> a = data [ i ] <EOL> b = data [ j ] <EOL> M [ i , j ] = calc_ldscore ( a . genotype , b . genotype ) <EOL> M = symmetrize ( M ) <EOL> logging . debug ( "<STR_LIT>" . format ( ldmatrix ) ) <EOL> M . tofile ( ldmatrix ) <EOL> else : <EOL> nmarkers = len ( Bed ( markerbedfile ) ) <EOL> M = np . fromfile ( ldmatrix , dtype = "<STR_LIT:float>" ) . reshape ( nmarkers , nmarkers ) <EOL> logging . debug ( "<STR_LIT>" . format ( ldmatrix , nmarkers ) ) <EOL> from jcvi . graphics . base import plt , savefig , Rectangle , draw_cmap <EOL> plt . rcParams [ "<STR_LIT>" ] = <NUM_LIT:0> <EOL> fig = plt . figure ( <NUM_LIT:1> , ( iopts . w , iopts . h ) ) <EOL> root = fig . add_axes ( [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:1> ] ) <EOL> ax = fig . add_axes ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> ax . matshow ( M , cmap = iopts . cmap ) <EOL> bed = Bed ( markerbedfile ) <EOL> xsize = len ( bed ) <EOL> extent = ( <NUM_LIT:0> , nmarkers ) <EOL> chr_labels = [ ] <EOL> ignore_size = <NUM_LIT:20> <EOL> for ( seqid , beg , end ) in bed . get_breaks ( ) : <EOL> ignore = abs ( end - beg ) < ignore_size <EOL> pos = ( beg + end ) / <NUM_LIT:2> <EOL> chr_labels . append ( ( seqid , pos , ignore ) ) <EOL> if ignore : <EOL> continue <EOL> ax . plot ( ( end , end ) , extent , "<STR_LIT>" , lw = <NUM_LIT:1> ) <EOL> ax . plot ( extent , ( end , end ) , "<STR_LIT>" , lw = <NUM_LIT:1> ) <EOL> for label , pos , ignore in chr_labels : <EOL> pos = <NUM_LIT> + pos * <NUM_LIT> / xsize <EOL> if not ignore : <EOL> root . text ( pos , <NUM_LIT> , label , <EOL> ha = "<STR_LIT>" , va = "<STR_LIT>" , rotation = <NUM_LIT> , color = "<STR_LIT>" ) <EOL> root . text ( <NUM_LIT> , pos , label , <EOL> ha = "<STR_LIT:right>" , va = "<STR_LIT>" , color = "<STR_LIT>" ) <EOL> ax . set_xlim ( extent ) <EOL> ax . set_ylim ( extent ) <EOL> ax . set_axis_off ( ) <EOL> draw_cmap ( root , "<STR_LIT>" , <NUM_LIT:0> , <NUM_LIT:1> , cmap = iopts . cmap ) <EOL> root . add_patch ( Rectangle ( ( <NUM_LIT> , <NUM_LIT> ) , <NUM_LIT> , <NUM_LIT> , fill = False , ec = "<STR_LIT:k>" , lw = <NUM_LIT:2> ) ) <EOL> m = mstmap . split ( "<STR_LIT:.>" ) [ <NUM_LIT:0> ] <EOL> root . text ( <NUM_LIT> , <NUM_LIT> , "<STR_LIT>" . format ( m ) , ha = "<STR_LIT>" ) <EOL> root . set_xlim ( <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> root . set_ylim ( <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> root . set_axis_off ( ) <EOL> image_name = m + "<STR_LIT>" + "<STR_LIT:.>" + iopts . format <EOL> savefig ( image_name , dpi = iopts . dpi , iopts = iopts ) <EOL> def header ( args ) : <EOL> """<STR_LIT>""" <EOL> from jcvi . formats . base import DictFile <EOL> p = OptionParser ( header . __doc__ ) <EOL> p . add_option ( "<STR_LIT>" , default = "<STR_LIT>" , <EOL> help = "<STR_LIT>" ) <EOL> p . add_option ( "<STR_LIT>" , help = "<STR_LIT>" ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:2> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mstmap , conversion_table = args <EOL> data = MSTMap ( mstmap ) <EOL> hd = data . header <EOL> conversion = DictFile ( conversion_table ) <EOL> newhd = [ opts . prefix + conversion . get ( x , x ) for x in hd ] <EOL> print "<STR_LIT:\t>" . join ( hd ) <EOL> print "<STR_LIT>" <EOL> print "<STR_LIT:\t>" . join ( newhd ) <EOL> ids = opts . ids <EOL> if ids : <EOL> fw = open ( ids , "<STR_LIT:w>" ) <EOL> print >> fw , "<STR_LIT:\n>" . join ( newhd ) <EOL> fw . close ( ) <EOL> def rename ( args ) : <EOL> """<STR_LIT>""" <EOL> p = OptionParser ( rename . __doc__ ) <EOL> p . set_outfile ( ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:2> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mstmap , bedfile = args <EOL> markersbed = Bed ( bedfile ) <EOL> markers = markersbed . order <EOL> data = MSTMap ( mstmap ) <EOL> header = data . header <EOL> header = [ header [ <NUM_LIT:0> ] ] + [ "<STR_LIT>" , "<STR_LIT:start>" ] + header [ <NUM_LIT:1> : ] <EOL> renamed = [ ] <EOL> for b in data : <EOL> m , geno = b . id , b . genotype <EOL> om = m <EOL> if m not in markers : <EOL> m = m . rsplit ( "<STR_LIT:.>" , <NUM_LIT:1> ) [ <NUM_LIT:0> ] <EOL> if m not in markers : <EOL> continue <EOL> i , mb = markers [ m ] <EOL> renamed . append ( [ om , mb . seqid , mb . start , "<STR_LIT:\t>" . join ( list ( geno ) ) ] ) <EOL> renamed . sort ( key = lambda x : ( x [ <NUM_LIT:1> ] , x [ <NUM_LIT:2> ] ) ) <EOL> fw = must_open ( opts . outfile , "<STR_LIT:w>" ) <EOL> print >> fw , "<STR_LIT:\t>" . join ( header ) <EOL> for d in renamed : <EOL> print >> fw , "<STR_LIT:\t>" . join ( str ( x ) for x in d ) <EOL> def anchor ( args ) : <EOL> """<STR_LIT>""" <EOL> from jcvi . formats . blast import bed <EOL> p = OptionParser ( anchor . __doc__ ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:2> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mapbed , blastfile = args <EOL> bedfile = bed ( [ blastfile ] ) <EOL> markersbed = Bed ( bedfile ) <EOL> markers = markersbed . order <EOL> mapbed = Bed ( mapbed , sorted = False ) <EOL> for b in mapbed : <EOL> m = b . accn <EOL> if m not in markers : <EOL> continue <EOL> i , mb = markers [ m ] <EOL> new_accn = "<STR_LIT>" . format ( mb . seqid , mb . start , mb . end ) <EOL> b . accn = new_accn <EOL> print b <EOL> def bed ( args ) : <EOL> """<STR_LIT>""" <EOL> p = OptionParser ( bed . __doc__ ) <EOL> p . add_option ( "<STR_LIT>" , default = False , action = "<STR_LIT:store_true>" , <EOL> help = "<STR_LIT>" ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:1> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mapout , = args <EOL> pf = mapout . split ( "<STR_LIT:.>" ) [ <NUM_LIT:0> ] <EOL> mapbed = pf + "<STR_LIT>" <EOL> bm = BinMap ( mapout ) <EOL> bm . print_to_bed ( mapbed , switch = opts . switch ) <EOL> return mapbed <EOL> def fasta ( args ) : <EOL> """<STR_LIT>""" <EOL> from jcvi . formats . sizes import Sizes <EOL> p = OptionParser ( fasta . __doc__ ) <EOL> p . add_option ( "<STR_LIT>" , default = <NUM_LIT:1000> , type = "<STR_LIT:int>" , <EOL> help = "<STR_LIT>" ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:2> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mapout , sfasta = args <EOL> Flank = opts . extend <EOL> pf = mapout . split ( "<STR_LIT:.>" ) [ <NUM_LIT:0> ] <EOL> mapbed = pf + "<STR_LIT>" <EOL> bm = BinMap ( mapout ) <EOL> bm . print_to_bed ( mapbed ) <EOL> bed = Bed ( mapbed , sorted = False ) <EOL> markersbed = pf + "<STR_LIT>" <EOL> fw = open ( markersbed , "<STR_LIT:w>" ) <EOL> sizes = Sizes ( sfasta ) . mapping <EOL> for b in bed : <EOL> accn = b . accn <EOL> scf , pos = accn . split ( "<STR_LIT:.>" ) <EOL> pos = int ( pos ) <EOL> start = max ( <NUM_LIT:0> , pos - Flank ) <EOL> end = min ( pos + Flank , sizes [ scf ] ) <EOL> print >> fw , "<STR_LIT:\t>" . join ( str ( x ) for x in ( scf , start , end , accn ) ) <EOL> fw . close ( ) <EOL> fastaFromBed ( markersbed , sfasta , name = True ) <EOL> def hamming_distance ( a , b , ignore = None ) : <EOL> dist = <NUM_LIT:0> <EOL> for x , y in zip ( a , b ) : <EOL> if ignore and ignore in ( x , y ) : <EOL> continue <EOL> if x != y : <EOL> dist += <NUM_LIT:1> <EOL> return dist <EOL> OK , BREAK , END = range ( <NUM_LIT:3> ) <EOL> def check_markers ( a , b , maxdiff ) : <EOL> if a . seqid != b . seqid : <EOL> return END , None <EOL> diff = hamming_distance ( a . genotype , b . genotype , ignore = "<STR_LIT:->" ) <EOL> max_allowed = len ( a ) * maxdiff <EOL> if diff <= max_allowed : <EOL> return OK , None <EOL> return BREAK , ( a . seqid , a . pos , b . pos ) <EOL> def breakpoint ( args ) : <EOL> """<STR_LIT>""" <EOL> from jcvi . utils . iter import pairwise <EOL> p = OptionParser ( breakpoint . __doc__ ) <EOL> p . add_option ( "<STR_LIT>" , default = <NUM_LIT> , type = "<STR_LIT:float>" , <EOL> help = "<STR_LIT>" ) <EOL> opts , args = p . parse_args ( args ) <EOL> if len ( args ) != <NUM_LIT:1> : <EOL> sys . exit ( not p . print_help ( ) ) <EOL> mstmap , = args <EOL> diff = opts . diff <EOL> data = MSTMap ( mstmap ) <EOL> good = [ ] <EOL> nsingletons = <NUM_LIT:0> <EOL> for i in xrange ( <NUM_LIT:1> , len ( data ) - <NUM_LIT:1> ) : <EOL> a = data [ i ] <EOL> left_label , left_rr = check_markers ( data [ i - <NUM_LIT:1> ] , a , diff ) <EOL> right_label , right_rr = check_markers ( a , data [ i + <NUM_LIT:1> ] , diff ) <EOL> if left_label == BREAK and right_label == BREAK : <EOL> nsingletons += <NUM_LIT:1> <EOL> continue <EOL> good . append ( a ) <EOL> logging . debug ( "<STR_LIT>" . format ( nsingletons ) ) <EOL> for a , b in pairwise ( good ) : <EOL> label , rr = check_markers ( a , b , diff ) <EOL> if label == BREAK : <EOL> print "<STR_LIT:\t>" . join ( str ( x ) for x in rr ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> main ( ) </s>
