Find file
Fetching contributors…
Cannot retrieve contributors at this time
executable file 355 lines (279 sloc) 11.6 KB
#!/usr/bin/env python
'''
given a SORTED file containing uniqued lines (see preprocess_radtag_lane.py) with cluster and node label prepended
computes multiple alignments across all cluster sequences and outputs SAM formatted alignments taking the most prevalent longest sequence as reference.
'''
import os, sys, re
import musclemap
from collections import defaultdict
from config import RTDROOT
def next_cluster_lines(fh):
this_cl = None
cl_lines = []
for l in fh:
if l.split()[0] != this_cl:
if this_cl is not None:
return cl_lines
else:
this_cl = l.split()[0]
cl_lines.append(l)
else:
cl_lines.append(l)
def samline_from_alnpair(rname,raln,qname,qaln,qqual):
if set(qqual) == set(['#']):
return None
leader,qseq = re.search('^(-*)(.*?)$',qaln).groups()
pos = len(leader)+1
cigar = []
nm = 0
md = []
qi = 0
for r,q in zip(raln[len(leader):].upper(),qseq.rstrip('-').upper()):
if q != '-':
qq = qqual[qi]
qi += 1
else:
qq = None
if qq == '#' or q == 'N' or r =='N':
cigar.append('S')
elif r in ['A','C','G','T'] and q in ['A','C','G','T']:
cigar.append('M')
if r != q:
nm += 1
md.append(r)
else:
md.append(1)
elif r == '-' and q == '-':
cigar.append('P')
elif r == '-':
cigar.append('I')
nm += 1
elif q == '-':
cigar.append('D')
nm += 1
md.append('^'+r)
#print ''.join(cigar)
if 'S' in ''.join(cigar).strip('S'):
return None
#figure out cigar
ccnt = 1
cli = []
cstate = None
for c in cigar:
if cstate == c:
ccnt += 1
else:
if cstate is not None:
cli.append('%d%s' % (ccnt,cstate))
cstate = c
ccnt = 1
cli.append('%d%s' % (ccnt,cstate))
cstr = ''.join(cli)
#figure out md
mdli = []
mddel = []
mdcnt = 0
for c in md+['A']:
if isinstance(c,int):
mdcnt += c
if len(mddel) > 0:
mdli.append('^'+(''.join(mddel)))
mddel = []
else:
if mdcnt:
mdli.append(str(mdcnt))
mdcnt = 0
if c.startswith('^'):
mddel.append(c[1:])
else:
if len(mddel) > 0:
mdli.append('^'+(''.join(mddel)))
mddel = []
mdli.append(c)
mdstr = ''.join(mdli[:-1])
if mdstr == '':
mdstr = '0'
return '\t'.join([qname,'0',rname,str(pos),'30',cstr,'*','0','0',qaln.replace('-',''), qqual, 'NM:i:%s\tMD:Z:%s' % (nm,mdstr)])
def ref_seq_from_clust(clname,cl_aln):
ref_seq = cl_aln[0][1].replace('-','')
fa_str = '>%s\n%s\n' % (clname,ref_seq)
return fa_str
def indiv_in_clust(cl_lines,rep_cut = 0):
if isinstance(cl_lines[0],str):
cl_lines = [l.strip().split() for l in cl_lines]
ind_cts = defaultdict(int)
for l in cl_lines:
for ind,ct in zip( l[5].split(','), [int(i) for i in l[6].split(',')] ):
if ct >= rep_cut:
ind_cts[ind] += ct
return ind_cts
def aln_from_clust(clname,cl_lines,keep_seqs=None,seq_len=0,break_on_error=True):
if isinstance(cl_lines[0],str):
cl_lines = [l.strip().split() for l in cl_lines]
if keep_seqs is not None and len(cl_lines) > keep_seqs:
orig_ind_ct = indiv_in_clust(cl_lines)
orig_ind = len(indiv_in_clust(cl_lines))
orig_len = len(cl_lines)
cl_lines.sort(key = lambda l: (len(l[5].split(',')),sum([int(i) for i in l[6].split(',')]), len(l[2])),reverse=True)
cl_lines = cl_lines[:keep_seqs]
now_ind = len(indiv_in_clust(cl_lines))
now_len = len(cl_lines)
drop_indiv = set(orig_ind_ct.keys()) - set(indiv_in_clust(cl_lines).keys())
#summarize!
print >> sys.stderr, '\tcluster %s abbreviated: orig %s lines, %s indiv now %s lines, %s indiv (dropped: %s)' % \
(clname, orig_len, orig_ind, now_len, now_ind,[(ind,orig_ind_ct[ind]) for ind in drop_indiv])
cl_seqs = [l[2] for l in cl_lines]
cl_nodes = [l[1] for l in cl_lines]
#20110919 qscore translation functionality moved to get_uniqued_lines_by_cluster.py
cl_quals = [l[4] for l in cl_lines]
if seq_len != 0: #truncate sequences
cl_seqs = [s[:seq_len] for s in cl_seqs]
cl_quals = [s[:seq_len] for s in cl_quals]
lastnode = None
cl_node_ids = []
for node in cl_nodes:
if node != lastnode:
ct = 0
lastnode = node
else:
ct += 1
cl_node_ids.append('%s.%03d' % (node,ct))
try:
cl_aln = sorted( zip( cl_node_ids, \
musclemap.muscle(cl_seqs,1), \
cl_quals, \
[zip( l[5].split(','), [int(i) for i in l[6].split(',')] ) for l in cl_lines] ) , \
key=lambda x: (len(x[1].replace('-','').replace('N','')),len(x[3]),len(x[2].replace('#',''))),reverse=True)
except:
print >> sys.stderr, 'alignment failed for cluster %s (%s lines)' % (clname,len(cl_lines))
if break_on_error:
raise
else:
print >> sys.stderr, '--skip_errors requested; proceeding'
return None
return cl_aln
def write_sam_from_aln(clname,cl_aln,rg_dict,samheader_fh,sambody_fh,ref_fh):
raln = cl_aln[0][1]
#sbfh = open(samfile+'.body','w')
#rofh = open(ref_fasta_file,'w')
rseq = ref_seq_from_clust(clname,cl_aln)
ref_fh.write(rseq)
#headers (@SQ lines)
headline = '@SQ\tSN:%s\tLN:%s\n' % (clname,len(cl_aln[0][2]))
samheader_fh.write(headline)
#body
for qname,qaln,qqual,inds_cts in cl_aln:
samline = samline_from_alnpair(clname,raln,qname,qaln,qqual)
if samline is None: continue
samfields = samline.split()
rg_lane = qname.split('.')[1]
#try:
# if any([len(el) != 2 for el in inds_cts]):
# print inds_cts
#except:
# print cl_aln
for ind,ct in inds_cts:
rg = '%s_%s' % (ind,rg_lane)
rg_dict[rg] = ind
for i in range(ct):
this_samline = '\t'.join([samfields[0]+'.%s.%04d' % (ind,i)] + samfields[1:])
sambody_fh.write('%s\tRG:Z:%s\n' % (this_samline,rg))
def calc_cluster_dirt(cl_lines):
cl_ind_ct = defaultdict(list)
for l in cl_lines:
f = l.split()
for ind,ct in zip(f[5].split(','),f[6].split(',')):
cl_ind_ct[(ind,f[1].split('.')[1])].append(int(ct))
totct = sum([sum(v) for v in cl_ind_ct.values()])
dirtct = sum([sum(sorted(v,reverse=True)[2:]) for v in cl_ind_ct.values()])
ctdirt = dirtct/float(totct)
return ctdirt
if __name__ == '__main__':
import argparse
ds = ' [%(default)s]'
#create command line parser
parser = argparse.ArgumentParser(description='generates SAM/BAM by multiple alignment within graph clusters')
parser.add_argument('-d','--clust_dirt_max',default=0.10,type=float,help='cluster "dirt" threshold for processing (see documentation)'+ds)
parser.add_argument('-i','--min_indiv',default=20,type=int,help='minimum number of individuals with at least one sequence in a cluster to include cluster'+ds)
parser.add_argument('-k','--keep_seqs',default=100,type=int,help='only retain this many sequences for processing'+ds)
parser.add_argument('-l','--seq_len',default=0,type=int,help='arbitrarily truncate sequences in SAM/BAM output at this length if not 0'+ds)
parser.add_argument('-cs','--calc_only',action='store_true',help='calculate cluster statistics at supplied thresholds; do not generate alignments'+ds)
parser.add_argument('-s','--skip_errors',action='store_true',help=''+ds)
parser.add_argument('cluniq',help='sorted .cluniq file containing cluster-associated unique sequences')
parser.add_argument('fbase',help='basename for output files')
opts = parser.parse_args()
cluniq = opts.cluniq
fbase = opts.fbase
clust_dirt_max = opts.clust_dirt_max
min_indiv = opts.min_indiv
keep_seqs = opts.keep_seqs
seq_len = opts.seq_len
fdir = os.path.dirname(fbase)
try:
os.makedirs(fdir)
except:
pass
if opts.skip_errors:
break_on_error = False
print >> sys.stderr, 'skip_errors invoked; problem clusters will be skipped entirely'
else:
break_on_error = True
print >> sys.stderr, 'skip_errors not set; problem clusters will halt analysis'
fh = open(cluniq)
if not opts.calc_only:
samheader_fh = open(fbase+'.sam.header','w')
sambody_fh = open(fbase+'.sam.body','w')
ref_fh = open(fbase+'.fa','w')
clstats_fh = open(fbase+'.clstats','w')
rg_dict = {}
this_cl = None
cl_lines = []
cl_on = 0
for l in fh:
if l.split()[0] != this_cl:
if this_cl is not None:
cl_dirt = calc_cluster_dirt(cl_lines)
cl_indiv = len(indiv_in_clust(cl_lines))
clstats_fh.write('%s\t%s\t%s\t%s\n' % (this_cl,len(cl_lines),cl_indiv,cl_dirt))
if cl_on % 100 == 0: print >> sys.stderr, '%s\tcluster: %s\tunique seqs: %s\tindiv: %s\tdirt: %s' % (cl_on,this_cl,len(cl_lines),cl_indiv,cl_dirt)
if not opts.calc_only and cl_dirt < clust_dirt_max and cl_indiv >= min_indiv:
cl_aln = aln_from_clust(this_cl,cl_lines,keep_seqs,seq_len,break_on_error)
write_sam_from_aln(this_cl,cl_aln,rg_dict,samheader_fh,sambody_fh,ref_fh)
cl_on += 1
this_cl = l.split()[0]
cl_lines = []
cl_lines.append(l)
clstats_fh.write('%s\t%s\t%s\t%s\n' % (this_cl,len(cl_lines),cl_indiv,cl_dirt))
if not opts.calc_only:
cl_aln = aln_from_clust(this_cl,cl_lines,keep_seqs)
if cl_aln is not None and calc_cluster_dirt(cl_lines) < clust_dirt_max and len(indiv_in_clust(cl_lines)) >= min_indiv:
write_sam_from_aln(this_cl,cl_aln,rg_dict,samheader_fh,sambody_fh,ref_fh)
clstats_fh.close()
os.system(os.path.join(RTDROOT,'plot_error.py %s > %s' % (fbase+'.clstats',fbase+'.clstats.cdest' )))
#finish headers (@RG lines)
if not opts.calc_only:
if len(rg_dict) == 0:
print >> sys.stderr, 'readgroup dict is empty; no individuals included in final dataset. Check number of individuals and cluster dirt cutoffs and re-run'
print >> sys.stderr, 'close output files ...',
samheader_fh.close()
sambody_fh.close()
ref_fh.close()
print >> sys.stderr, 'done.\nremove output files ...',
os.unlink(samheader_fh.name)
os.unlink(sambody_fh.name)
os.unlink(ref_fh.name)
print >> sys.stderr, 'done'
sys.exit(1)
for rg in rg_dict:
headline = '@RG\tID:%s\tPL:Illumina\tLB:%s\tSM:%s\n' % (rg,rg_dict[rg],rg_dict[rg])
samheader_fh.write(headline)
samheader_fh.close()
sambody_fh.close()
ref_fh.close()
print >> sys.stderr, 'index reference'
os.system('samtools faidx %s.fa' % (fbase))
print >> sys.stderr, 'add headers and sort'
os.system('cat %s.sam.header %s.sam.body | samtools view -bS - | samtools sort - %s' % (fbase,fbase,fbase))
print >> sys.stderr, 'index bam'
os.system('samtools index %s.bam' % (fbase))
print >> sys.stderr, 'done'