create012AllParents.py

#!/usr/bin/env python

import sys, os, gzip
import numpy as np
from collections import Counter,defaultdict
from subprocess import Popen, PIPE
from time import time
from collections import Counter


if len(sys.argv) < 5:
    print("Usage: create012AllParents.py <HDB> <out-012-file> phased=1 or 0 proband=1 or 0 ")
    exit(1)

hdb = sys.argv[1]
out = sys.argv[2]
phased = int(sys.argv[3])
proband = int(sys.argv[4])
#faiFn = sys.argv[5]
#th = float(sys.argv[6])
#hdb = "/mnt/wigclust8/data/safe/autism/HT/SSC510/SSC510-22"
#out = "a.vcf"
#phased = False
print(hdb, out, phased, file=sys.stderr)

badGens = set(list('oucd'))
ambiguous = {
    'AC':'M', 
    'AG':'R', 
    'AT':'W', 
    'CG':'S', 
    'CT':'Y', 
    'GT':'K'
}
ramb = {v:k for k,v in list(ambiguous.items())}
for t,c in list(ambiguous.items()):
    ambiguous[t[1] + t[0]] = c

class P:
    def __init__(p,cs):
        p.ch,p.pos,p.refA,p.A,p.C,p.G,p.T = cs
        p.pos,p.A,p.C,p.G,p.T = list(map(int,[p.pos,p.A,p.C,p.G,p.T]))
        aad = {a:n for a,n in zip("ACGT",[p.A,p.C,p.G,p.T])} 
        ad = {a:n for a,n in list(aad.items()) if a!=p.refA and n>0}
        p.AI = {p.refA:0}
        for ai,(a,n) in enumerate(sorted(list(ad.items()),key=lambda x: (-x[1],x[0]))):
            p.AI[a] = ai+1

        p.unphasedD = {}
        p.phasedD = {}
        for a,ai in list(p.AI.items()):
            p.unphasedD[a + a] = "%d/%d" % (ai,ai)
            p.phasedD[a + a] = "%d|%d" % (ai,ai)
            for b,bi in list(p.AI.items()):
                if a==b: continue
                fi,si = sorted([ai,bi])
                
                p.unphasedD[a + b] = "%d/%d" % (fi,si)
                p.phasedD[a + b] = "%d|%d" % (ai,bi)
                p.unphasedD[ambiguous[a + b]*2] = "%d/%d" % (fi,si)
                p.phasedD[ambiguous[a + b]*2] = "%d/%d" % (fi,si)

        # unphasedD {"AA": "0/0", "AC": 0/1, "CA": 0/1, "RR": "0/1"}
        # phasedD   {"AA": "0|0", "AC": 0|1, "CA": 1|0, "RR": "0/1"}
        
       

    def vcfGen(p,AT,ANT, phased,fmId):
        if AT in badGens:
            return '-1'
        D = p.phasedD if phased else p.unphasedD
        #return D[AT+ANT]
        try:
            return sum(map(int,[D[AT+ANT][0], D[AT+ANT][2]]))
        except KeyError:
            print(p.ch, p.pos, p.AI, p.A, p.C, p.G, p.T, AT, ANT, fmId, file=sys.stderr)
            
        # if phased==True:  0|0  1|0 0|1, ...
        # else 0/0, 0/1, 0/2
        # * if AT or ANT are small letters: ./. or .|. 
        # * if AT is a two-nucleotide code, assert that AT == ANT, and 
        # output unphased regardless of phased parameter

POS = []
posFn = hdb + '-pos.txt'
with open(posFn, 'r') as f:
    f.readline()

    for l in f:
        POS.append(P(l.strip('\n\r').split('\t')))


FM_CH = {} 
FM_HPTH = defaultdict(dict) 

hpthFn = hdb + '-hpth.txt'

if os.path.isfile(hpthFn):
    f = open(hpthFn)
elif  os.path.isfile(hpthFn + '.gz'):
    f = gzip.open(hpthFn + '.gz')
else:
    print("no file ", hpthFn + '[.gz]')
    exit(1)

f.readline()
for li,l in enumerate(f):
    if li % 100 == 0:
        print(li, "...")
    fmId,perId,thrT,thr = l.strip('\n\r').split('\t')
    #if fmId in FM_CH and perId != FM_CH[fmId]:
    #if fmId in FM_CH and ((proband and 's1' in perId) or (not proband and 'p1' in perId)):
    #    continue
    #FM_CH[fmId] = perId # my addition
    FM_HPTH[fmId][thrT] = thr
    #if li == 50000:
    #    break

f.close()

NH ,= set(len(x) for x in list(FM_HPTH.values()))
assert NH == 4

fmIdOrder = sorted(FM_HPTH.keys())


### TODO: get that somehow!!!!
"""
with open(faiFn, 'r') as f:
    chrLen = {'chr'+str(x):0 for x in range(1,23)}
    for l in f:
        cs = l.strip('\n\r').split('\t')
        if cs[0] in chrLen:
            chrLen[cs[0]] = int(cs[1])
"""

OF = open(out,"w")
# OF = sys.stdout 

#OF.write('##fileformat=VCFv4.2\n')
#OF.write('##fileDate=20171012\n')
#OF.write('##source=createVcfAllParents.py\n')
#for ch in sorted({p.ch for p in POS}):
#    OF.write('##contig=<ID=%s,length=%d>\n' % (ch,chrLen[ch]))
#OF.write('##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\n')


OF.write("\t".join('#CHROM POS ID REF ALT QUAL FILTER INFO FORMAT'.split(' ') + 
            ["%s.%s" % (fmId,rl) for fmId in fmIdOrder for rl in ["mo", "fa"]]) + "\n")

filterStats = defaultdict(int)
for pi,p in enumerate(POS):
    if (pi % 1000) == 0:
        print("writing out position", pi, "...", file=sys.stderr)

    """
    # no laternative alleles
    if len(p.AI) == 1:
        filterStats['no_alt'] += 1
        print >>sys.stderr, 'no_alt', p.ch, p.pos
        continue
    
    # more than one alternative 
    if len(p.AI) > 2:
        filterStats['more_alt'] += 1
        print >>sys.stderr, 'more_alt', p.ch, p.pos
        continue
    
    # refA not seen
    cntsA = dict(zip('ACGT',[p.A,p.C,p.G,p.T]));
    if cntsA[p.refA] == 0:
        filterStats['no_ref'] += 1
        print >>sys.stderr, 'no_ref', p.ch, p.pos        
        continue
    
    # less than 90% genotyped
    NG = sum(cntsA.values())
    if NG < 0.9 * len(FM_HPTH) * 2 * 2:
        filterStats['less_then_95%'] += 1
        print >>sys.stderr, 'less_then_90%', p.ch, p.pos, NG, len(FM_HPTH)
        continue
    """
    #if 100.*(NG - cntsA[p.refA])/NG < th:
    #    continue
    
    cs = [p.ch,
          str(p.pos), 
          "%s:%d" % (p.ch,p.pos), 
          p.refA, 
          ",".join([a for a,ai in sorted(list(p.AI.items()),key=lambda x: -x[1]) if ai>0]), 
          ".", 
          ".", 
          ".", 
          "GT"]
    for fmId in fmIdOrder:
        fmD = FM_HPTH[fmId]
        cs += [p.vcfGen(fmD['MT'][pi],fmD['MNT'][pi],phased,fmId), 
               p.vcfGen(fmD['FT'][pi],fmD['FNT'][pi],phased,fmId)]
    OF.write("\t".join(map(str, cs)) + "\n") 
    #if pi > 10000:
    #     break
OF.close()
with open(hdb+'-filterStats.txt', 'w') as f:
    f.write('\t'.join(map(str,sorted(filterStats)))+'\n')
    f.write('\t'.join(map(str,[filterStats[k] for k in sorted(filterStats)]))+'\n')