ABBABABA-4AF.py

import sys
from collections import defaultdict as d
from optparse import OptionParser, OptionGroup

#Author: Martin Kapun

#########################################################   HELP   #########################################################################
usage="python3 %prog --AF file.af --SNPs 500 --Order 4,5,6,7 --Output file "
parser = OptionParser(usage=usage)
group=OptionGroup(parser,'''
          /\
         /  \
        /\   \
       /  \   \
      /\   \   \
     /  \   \   \
    H1  H2  H3  H4

    A   B   B   A
    B   A   B   A
''')

#########################################################   CODE   #########################################################################

parser.add_option("--AF", dest="AF", help="Input file; First three columns: Chrom, Pos, MajorAllele/MinorAllele")
parser.add_option("--Output", dest="OUT", help="Output prefix")
parser.add_option("--Order", dest="ORDER", help="column positions of populations H1,H2,H3,H4 in input file")
parser.add_option("--SNPs", dest="SNP", help="number of SNPS",default="NA")
parser.add_option("--window", dest="window", help="window size in bp",default="NA")

(options, args) = parser.parse_args()
parser.add_option_group(group)

def D1(x,XL,YL,fXL,fYL,N):
    '''calculate nominator X and demoninator Y a la Soraggi et al. 2018, assumes X=BABA-ABBA '''
    A,B,C,D=x
    X=(A-B)*(C-D)
    Y=(A+B-2*A*B)*(C+D-2*C*D)
    XL.append(X)
    fXL[N].append(X)
    YL.append(Y)
    fYL[N].append(Y)

def D2(x,XL,YL,fXL,fYL,N):
    '''calculate nominator X and demoninator Y a la Durand et al 2011, assumes X=ABBA-BABA'''
    A,B,C,D=x
    X=(1-A)*B*C*(1-D)-A*(1-B)*C*(1-D)
    Y=(1-A)*B*C*(1-D)+A*(1-B)*C*(1-D)
    XL.append(X)
    fXL[N].append(X)
    YL.append(Y)
    fYL[N].append(Y)


def blockJackEven(fx,fy,D):
    ''' see Busing et al. 1999 '''

    ## at first retain only windows that don't result in window-wise D="NA"
    fxnew=[]
    fynew=[]
    for i in range(len(D)):
        if D[i]== "NA":
            continue
        fxnew.append(fx[i])
        fynew.append(fy[i])

    ## then calculate overall expected D
    Dest=sum([sum(x) for x in fxnew])/float(sum([sum(y) for y in fynew]))

    # calculate D when removing one group at a time:
    Dmean=[]
    for i in range(len(fxnew)):
        fxprime=fxnew[:i] + fxnew[i+1 :]
        fyprime=fynew[:i] + fynew[i+1 :]
        Dmean.append(sum([sum(x) for x in fxprime])/float(sum([sum(y) for y in fyprime])))

    ## g is the number of independent groups
    g=float(len(fxnew))

    ## caluclate the mean of all group-wise estimates of D
    meanJack=sum(Dmean)/float(len(Dmean))

    ## calculate the jackknife estimate and variance of D and z-score
    DjackEst=g*Dest-(g-1)*meanJack
    DjackVar=((g-1)/g)*sum([(Di-Dest)**2 for Di in Dmean])
    Z=Dest/(DjackVar**0.5)

    ## return
    return Dest,DjackEst,DjackVar,Z

def test0(x,y,z):
    ''' calculate the ratio of x and y if bin contains enough SNPs and if y is not 0, then append to list z '''

    if sum(y)==0:
        xy="NA"
    elif options.SNP!="NA" and len(x)!=SNPs:
        xy="NA"
    else:
        xy=str(sum(x)/float(sum(y)))
        if float(xy)>1 or float(xy)<-1:
            xy="NA"
    z.append(xy)
    return xy

def load_data(x):
  ''' import data either from a gzipped or or uncrompessed file or from STDIN'''

  import gzip
  if x=="-":
      y=sys.stdin
  elif x.endswith(".gz"):
      y=gzip.open(x,"rt", encoding="latin-1")
  else:
      y=open(x,"r", encoding="latin-1")
  return y

if options.SNP != "NA":
    SNPs=int(options.SNP)
else:
    WS=int(options.window)

ORDER=[int(x) for x in options.ORDER.split(",")]

Chr=""

# Store the position of every SNP
SNPwindows=[]

## to calculate a global value
fSNPX1,fSNPX2=d(list),d(list)
fSNPY1,fSNPY2=d(list),d(list)

## to calculate window-wise values
SNPX1,SNPX2=[],[]
SNPY1,SNPY2=[],[]

## make list to store window-wise values of D
D1L,D2L=[],[]

## open Outputfile for window-wise estimates
outWindow=open(options.OUT,"w")

## Counter of SNP groups
if options.SNP!="NA":
    N=0
else:
    N=""

outWindow.write("Chromosome\tStartPos\tAveragePos\tEndPos\tWindowSize\tNoOfSNPs\tSoraggi\tDurand\n")

for l in load_data(options.AF):
    a=l.rstrip().split()

    if options.window!="NA" and N=="":
        N=round(float(a[1])/WS,0)

    ## only keep allele frequencies of populations used for ABBA BABA and retain order
    AFs=[a[x] for x in ORDER]
    if Chr=="":
        Chr=a[0]

    ## continue if one AF missing
    if "NA" in AFs or "na" in AFs:
        continue

    AFs=[float(x) for x in AFs]

    ## continue if site is not informative according to Sorragi et al., i.e. if H1==H2 or H3==H4
    if AFs[0]==AFs[1] or AFs[2]==AFs[3]:
        continue

    # condition AFs for the ancestral allele which is more common in the outgroup population
    if AFs[-1]<0.5:
        AFs=[1-x for x in AFs]


    # test if condition for new window:

    if options.SNP!="NA":
        if len(SNPX1)==SNPs:
            cond=True
        else:
            cond=False
    else:
        if round(float(a[1])/WS,0)!=N:
            cond=True
        else:
            cond=False

    # if the end of a window is reached or a new contig starts, calculate the window-specific D, print it to the outputfile and move on
    if Chr!=a[0] or cond:

        #print(Chr,a[0],len(SNPX1),SNPwindows)

        # calculate D
        D1s=test0(SNPX1,SNPY1,D1L)
        D2s=test0(SNPX2,SNPY2,D2L)

        if len(SNPX1)!=0:
            START=min(SNPwindows)
            END=max(SNPwindows)
            MIDDLE=sum(SNPwindows)/float(len(SNPwindows))
            LENGTH=END-START

            # only print the values of D if the number of SNPs is equivalent to the number of SNPs defined as the binsize
            outWindow.write(Chr+"\t"+"\t".join([str(x) for x in [START,MIDDLE,END,LENGTH]])+"\t"+str(len(SNPX1))+"\t"+D1s+"\t"+D2s+"\n")

        # reset the lists
        SNPwindows=[]
        SNPX1,SNPX2=[],[]
        SNPY1,SNPY2=[],[]
        Chr=a[0]

        SNPwindows.append(int(a[1]))

        # store the numerator and denominator of D
        D1([float(x) for x in AFs],SNPX1,SNPY1,fSNPX1,fSNPY1,N)
        D2([float(x) for x in AFs],SNPX2,SNPY2,fSNPX2,fSNPY2,N)

        ## increase counter by one to start a new group
        if options.SNP!="NA":
            N+=1
        else:
            N=round(float(a[1])/float(options.window),0)

        continue

    SNPwindows.append(int(a[1]))
    Chr=a[0]

    D1([float(x) for x in AFs],SNPX1,SNPY1,fSNPX1,fSNPY1,N)
    D2([float(x) for x in AFs],SNPX2,SNPY2,fSNPX2,fSNPY2,N)

D1s=test0(SNPX1,SNPY1,D1L)
D2s=test0(SNPX2,SNPY2,D2L)

if len(SNPX1)!=0:

    START=min(SNPwindows)
    END=max(SNPwindows)
    MIDDLE=sum(SNPwindows)/float(len(SNPwindows))
    LENGTH=END-START

    outWindow.write(Chr+"\t"+"\t".join([str(x) for x in [START,MIDDLE,END,LENGTH]])+"\t"+str(len(SNPX1))+"\t"+D1s+"\t"+D2s+"\n")


## open Outputfile for genome-wide estimates
outGW=open(options.OUT+"_GenomeWide.D","w")

## calculate Jacknife stat
outGW.write("Method\tDest\tDjackEst\tDjackVar\tZ\n")

Z1=blockJackEven(fSNPX1,fSNPY1,D1L)
outGW.write("SoraggiEtAl\t"+"\t".join([str(x) for x in Z1])+"\n")

Z2=blockJackEven(fSNPX2,fSNPY2,D2L)
outGW.write("DurandEtAl\t"+"\t".join([str(x) for x in Z2])+"\n")