p.by.depth.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Name:       Estimate selection components
# adults :: ignore offspring
# v2:: Work on z-scale, treat male bulks with null var

import sys
from math import exp,log,sin,asin,sqrt




chrom = sys.argv[1] # chr2L
MinMax_MQ = [int(sys.argv[2]),60] # 20
MPoolDepth = [int(sys.argv[3]),int(sys.argv[4])] # [430,882]


invcf="BCF."+chrom+".vcf"

################ MAIN
MommyHet={}
for j in range(1,100):
    MommyHet[j]=[]

vcflines={}
Mdata={}
in1  =open("male.pools.txt", "r") # 
for line_idx, line in enumerate(in1):
    cols = line.replace('\n', '').split('\t') 
# 694    maleA1_S656    A
    try:
        vcflines[cols[2]].append(int(cols[0]))
        Mdata[cols[2]]=[] # r,a,m,z
    except KeyError:
        vcflines[cols[2]]=[int(cols[0])]

in1.close()

malepools={}
matfams={}
in1 =open("Sample.key.txt","r")
for line_idx, line in enumerate(in1):
    cols = line.replace('\n', '').split('\t') 
# vcf_j    sample    FamNo    type
# 530    6_c_S111    6    c
# 531    6_m_S305    6    m
    if line_idx>0:
        if int(cols[2])==0:
            malepools[int(cols[0])]=1
        else:
            try:
                uk=matfams[int(cols[2])]
            except KeyError:
                matfams[int(cols[2])]=[-9,-9]
            if cols[3]=="m":
                matfams[int(cols[2])][0]=int(cols[0])
            else:
                matfams[int(cols[2])][1]=int(cols[0])
in1.close()

Qstats=[[] for j in range(6)]
Qtats=[[] for j in range(6)]
Mpool=[]
pmp=0.0
pmf=0.0
inx=open(invcf, "r")
for line_idx, line in enumerate(inx):
# chr2L    10000016    .    C    A    999    .    DP=12589;VDB=0.714727;SGB=1128.3;RPB=0.879548;MQB=0.998021;MQSB=0.98709;BQB=0.514285;MQ0F=0.00111208;ICB=2.28618e-08;HOB=0.116428;AC=775;AN=1296;DP4=2249,421,4055,647;MQ=59    GT:PL:AD    1/1:255,51,0:0,17    1/1:235,18,0:0,6
    cols = line.replace('\n', '').split('\t')

    xyt = cols[7].split(";")
    mq = xyt[len(xyt)-1].split("=")
    if mq[0]=="MQ":
        MapQ = float(mq[1])
    else:
        print ("fail x",xyt)

    maledepth = 0
    Rcc=0
    for letter in vcflines:
        Mdata[letter]=[0,0,0,0]
        for j in vcflines[letter]:
            vv=cols[j].split(":") # 0/0:0,15,151:5,0
            if len(vv) == 3 and vv[0] != './.':
                Mdata[letter][0]+=int(vv[2].split(",")[0])
                Mdata[letter][1]+=int(vv[2].split(",")[1])
        maledepth+= Mdata[letter][0]+Mdata[letter][1]
        Rcc +=Mdata[letter][0]

    if maledepth>=MPoolDepth[0] and maledepth<=MPoolDepth[1] and MapQ>=MinMax_MQ[0] and MapQ<=MinMax_MQ[1]:

        Mpool.append(maledepth)
        pmp+= float(Rcc)/float(maledepth)
        p_pool=float(Rcc)/float(maledepth)

        RAfems=[0,0]

        Qs=[[0,0] for j in range(6)]
        Qt=[[0,0] for j in range(6)]
        momdepth = 0
        for z in matfams:

            j=matfams[z][0] # mom
            if j>=0:
                GT=cols[ j ].split(":")[0] # 1/1:101,14,0:1,12 but also really weird shit:: 1/1:0,9,0:26,25    0/0:0,5,5:3,1
                pl1=cols[ j ].split(":")[1]
                ad1=cols[ j ].split(":")[2]
                Rcc=int((cols[ j ].split(":")[2]).split(",")[0])
                Acc=int((cols[ j ].split(":")[2]).split(",")[1])

                if GT != "./." and len(pl1.split(","))==3 and (Rcc+Acc)>0:
                    gl = [int(pl1.split(",")[0]),int(pl1.split(",")[1]),int(pl1.split(",")[2])]
                    if min(gl)==0:    

                        cx=Rcc+Acc
                        if cx>99:
                            cx=99
                        RAfems[0]+=Rcc
                        RAfems[1]+=Acc        
                        if Rcc>0 and Acc==0:
                            gl=(1.0)
                            MommyHet[cx].append(0.0)
                        elif Rcc==0 and Acc>0:
                            gl=(0.0)
                            MommyHet[cx].append(0.0)
                        else:
                            gl=(0.5)
                            MommyHet[cx].append(1.0)
                        if (Rcc+Acc)<6:
                            Qs[Rcc+Acc-1][0]+=1
                            Qs[Rcc+Acc-1][1]+=gl
                        else:
                            Qs[5][0]+=1
                            Qs[5][1]+=gl

                        if GT=="0/0":
                            gl=(1.0)
                        elif GT=="1/1":
                            gl=(0.0)
                        else:
                            gl=(0.5)
                        if (Rcc+Acc)<6:
                            Qt[Rcc+Acc-1][0]+=1
                            Qt[Rcc+Acc-1][1]+=gl
                        else:
                            Qt[5][0]+=1
                            Qt[5][1]+=gl


        pmf+=float(RAfems[0])/float(RAfems[0]+RAfems[1])
        for j in range(6):
            nx =  Qs[j][0]
            if nx>0:
                px = float(Qs[j][1])/float(nx)
                for k in range(nx):
                    Qstats[j].append( (px-p_pool) )
                px = float(Qt[j][1])/float(nx)
                for k in range(nx):
                    Qtats[j].append( (px-p_pool) )

Mpool.sort()
snpx = len(Mpool)
pmp= pmp/float(snpx)
pmf= pmf/float(snpx)
for j in range(100):
    print ("pool D",float(j)/100, Mpool[ int(j*snpx/100) ])

print ("Pool Q ",snpx,pmp)
print ("Dev Fem-Male pool",pmf-pmp)
for j in range(6):
    print("depth",j+1,len(Qstats[j]),sum(Qstats[j])/float(len(Qstats[j])),sum(Qtats[j])/float(len(Qtats[j])))


for j in range(1,100):
    if len(MommyHet[j])>0:
        het = sum(MommyHet[j])/float(len(MommyHet[j]))
    print (j, len(MommyHet[j]), het)