root/galaxy-central/eggs/bx_python-0.5.0_dev_f74aec067563-py2.6-macosx-10.6-universal-ucs2.egg/EGG-INFO/scripts/mMK_bitset.py @ 3

#!/usr/bin/python2.6

import sys

import bx.align.maf
import bx.bitset
from bx.bitset_builders import *
from itertools import *
from optparse import OptionParser
from rpy import *

def main():

        # Parse the command line
        parser = OptionParser(usage = "usage: %prog [options] maf_file snp_file neutral_file window_size step_size") 
        
        parser.add_option("-o", "--outfile", help = "Specify file for output")
        
        parser.add_option("-s", "--species", type = "string", default = "panTro2")
        
        parser.add_option("-b", "--build", type = "string", default = "hg18")
        
        (options, args) = parser.parse_args()

        if len(args) != 5:
                parser.error("Incorrect number of arguments")
        else:
                maf_filename = args[0]
                snp_filename = args[1]
                neutral_filename = args[2]
                window_size = int(args[3])
                step_size = int(args[4])
                
        if options.outfile != None:
                out_file = open(options.outfile, 'w')
                
        #Generate snp and neutral bitsets
        AR_snp_bitsets = binned_bitsets_from_file(open(snp_filename))
        neutral_bitsets = binned_bitsets_from_file(open(neutral_filename))

        # Generate divergence bitset from maf file
        AR_div_bitsets = dict()
        chr_lens = dict()
        reader = bx.align.maf.Reader( open (maf_filename) )
        
        for block in reader:
                comp1 = block.get_component_by_src_start( options.build )
                comp2 = block.get_component_by_src_start( options.species )
                
                if comp1 is None or comp2 is None:
                        continue
                        
                # Chromosome, start, and stop of reference species alignment
                chr = comp1.src.split( '.' )[1]
                start = comp1.start
                end = comp1.end
                
                # Get or create bitset for this chromosome
                if chr in AR_div_bitsets:
                        bitset = AR_div_bitsets[chr]
                else:
                        bitset = AR_div_bitsets[chr] = bx.bitset.BinnedBitSet()
                        chr_lens[chr] = comp1.get_src_size()
                        
                # Iterate over text and set diverged bit
                pos = start
                for ch1, ch2 in izip( comp1.text.upper(), comp2.text.upper() ):
                        if ch1 == '-': continue
                        if ch2 == '-':
                                pos += 1
                                continue
                        
                        if ch1 != ch2 and not AR_snp_bitsets[chr][pos]:
                                bitset.set( pos )
                        pos += 1
                        
        # Debugging Code
#       for chr in AR_div_bitsets:
#               for pos in range(0, AR_div_bitsets[chr].size):
#                       if AR_div_bitsets[pos]:
#                               print >> sys.stderr, chr, pos, pos+1
        
        # Copy div and snp bitsets
        nonAR_snp_bitsets = dict()
        for chr in AR_snp_bitsets:
                nonAR_snp_bitsets[chr] = bx.bitset.BinnedBitSet()
                nonAR_snp_bitsets[chr].ior(AR_snp_bitsets[chr])
                
        nonAR_div_bitsets = dict()
        for chr in AR_div_bitsets:
                nonAR_div_bitsets[chr] = bx.bitset.BinnedBitSet()
                nonAR_div_bitsets[chr].ior(AR_div_bitsets[chr])
        
        # Generates AR snps by intersecting with neutral intervals
        for chr in AR_snp_bitsets:
                AR_snp_bitsets[chr].iand(neutral_bitsets[chr])
        
        # Generates AR divs by intersecting with neutral intervals      
        for chr in AR_div_bitsets:
                AR_div_bitsets[chr].iand(neutral_bitsets[chr])

        # Inverts the neutral intervals so now represents nonAR
        for chr in neutral_bitsets:
                neutral_bitsets[chr].invert()
        
        # Generates nonAR snps by intersecting with masked neutral intervals    
        for chr in nonAR_snp_bitsets:
                nonAR_snp_bitsets[chr].iand(neutral_bitsets[chr])
        # Generates nonAR divs by intersecting with masked neutral intervals            
        for chr in nonAR_div_bitsets:
                nonAR_div_bitsets[chr].iand(neutral_bitsets[chr])

        for chr in AR_div_bitsets:
                for window in range(0, chr_lens[chr] - window_size, step_size):
#                       neutral_size = neutral_bitsets[chr].count_range(window, window_size)
#                       if neutral_size < 9200: continue
                        AR_snp = AR_snp_bitsets[chr].count_range(window, window_size)
                        AR_div = AR_div_bitsets[chr].count_range(window, window_size)
                        nonAR_snp = nonAR_snp_bitsets[chr].count_range(window, window_size)
                        nonAR_div = nonAR_div_bitsets[chr].count_range(window, window_size)
                        
                        if nonAR_snp >= 6 and nonAR_div >= 6 and AR_snp >= 6 and AR_div >= 6:
                                MK_pval = MK_chi_pvalue(nonAR_snp, nonAR_div, AR_snp, AR_div)
                        else:
                                MK_pval = MK_fisher_pvalue(nonAR_snp, nonAR_div, AR_snp, AR_div)
                                
                        if options.outfile != None:
                                out_file.write("%s\t%d\t%d\t%d\t%d\t%d\t%d\t%1.15f\n" % (chr, window, window+window_size, nonAR_snp, nonAR_div, AR_snp, AR_div, MK_pval))
                        else:
                                print "%s\t%d\t%d\t%d\t%d\t%d\t%d\t%1.15f" % (chr, window, window+window_size, nonAR_snp, nonAR_div, AR_snp, AR_div, MK_pval)
        
        if options.outfile != None:
                out_file.close()
                        
def MK_fisher_pvalue(win_snp, win_div, AR_snp, AR_div):

        if win_snp == 0 and win_div == 0 and AR_snp == 0 and AR_div == 0:
                return 1.0
        
        fisher_result = r.fisher_test(r.matrix(r.c([win_snp, win_div, AR_snp, AR_div]), nr = 2))
        
        return fisher_result['p.value']

def MK_chi_pvalue(win_snp, win_div, AR_snp, AR_div):
        
        chi_result = r.chisq_test(r.matrix(r.c([win_snp, win_div, AR_snp, AR_div]), nr = 2))
        
        return chi_result['p.value']

main()