Python并发实践_03_并发实战之一

16S数据质控流程，一次下机lane包括很多的项目，每个项目有独立的合同号，一个项目可能包含16S或者ITS两种，通过一个完整的pipeline，将上游拆分好的数据全部整理成可以直接分析的数据。原本这个工作是通过并行的sge实现，是运行层面的并行，现在在程序层面实现并行处理，可以脱离sge系统工作。

  1 import os
  2 import sys
  3 import re
  4 import time
  5 import collections
  6 from multiprocessing import Process,JoinableQueue,Queue,cpu_count
  7 from threading import Thread
  8 from settings import primer,pandaseq_soft
  9 from programs import *
 10 
 11 Result = collections.namedtuple("Result","compact sample_name HQ_fq")
 12 
 13 def parse_sam_barcode_file(sam_barcode_file):
 14     for line in open(sam_barcode_file):
 15         yield line.strip().split('\t')
 16 
 17 def proc(compact,sample_name,work_path,lib_method,data_type):
 18     split_path = '%s/Split'%work_path
 19     QC_path = '%s/QC'%work_path
 20     compact_path = '%s/%s'%(QC_path,compact)
 21     if not os.path.exists(compact_path):
 22         os.makedirs(compact_path)
 23     sample_path = '%s/%s'%(compact_path,sample_name)
 24     if not os.path.exists(sample_path):
 25         os.makedirs(sample_path)
 26     original_path = '%s/%s/%s'%(split_path,compact,sample_name)
 27     (read1,read2) = os.popen('ls %s/*'%original_path).read().strip().split('\n')
 28     pandaseq_fq = '%s/pandaseq.fq'%sample_path
 29     pandaseq_log = '%s/pandaseq.log'%sample_path
 30     pandaseq(pandaseq_soft,read1,read2,pandaseq_fq,primer[lib_method][data_type]['forward'],primer[lib_method][data_type]['reverse'],pandaseq_log)
 31     high_quality_fq = '%s/high_quality.fq'%sample_path
 32     high_quality_log = '%s/high_quality.stat'%sample_path
 33     QC(pandaseq_fq,high_quality_fq,high_quality_log,data_type)
 34     return Result(compact,sample_name,high_quality_fq)
 35 
 36 def worker(work_path,jobs,results):
 37     while True:
 38         try:
 39             compact,sample_name,lib_method,data_type = jobs.get()
 40             try:
 41                 result = proc(compact,sample_name,work_path,lib_method,data_type)
 42                 sys.stderr.write( 'Process %s is finished doing with compact:%s sample_name:%s\n'%(os.getpid(),compact,sample_name) )
 43                 results.put(result)
 44             except:
 45                 sys.stderr.write('Process %s is FIALED !!! %s/%s may be some problem!\n'%(os.getpid(),compact,sample_name))
 46                 jobs.put((compact,sample_name,lib_method,data_type))
 47                 sys.stderr.write('The job is repushed into the queue,with compact:%s sample_name:%s\n'%(compact,sample_name))
 48         finally:
 49             jobs.task_done()
 50 
 51 def add_jobs(work_path,sam_barcode_file_list,jobs):
 52     job_num = 0
 53     data_type_hash = {}
 54     for todo,sam_barcode_file in enumerate(sam_barcode_file_list):
 55         sam_barcode_file = sam_barcode_file.strip()
 56         if not os.path.isfile(sam_barcode_file):
 57             continue
 58         lib_method = get_lib_method(sam_barcode_file)
 59         if lib_method is None:
 60             continue
 61         print 'sam_barcode_file loading: %s             ......  ok\n'%sam_barcode_file
 62         for compact,sample_name,barcode_info,data_type in parse_sam_barcode_file(sam_barcode_file):
 63         print 'sam_barcode_file loading: %s             ......  ok\n'%sam_barcode_file
 64         for compact,sample_name,barcode_info,data_type in parse_sam_barcode_file(sam_barcode_file):
 65             if not data_type_hash.has_key(compact):
 66                 data_type_hash[compact] = {}
 67             if not data_type_hash[compact].has_key(data_type):
 68                 data_type_hash[compact][data_type] = []
 69             data_type_hash[compact][data_type].append(sample_name)
 70             jobs.put((compact,sample_name,lib_method,data_type))
 71             job_num += 1
 72             sys.stderr.write('The job is pushed into the queue,with compact:%s sample_name:%s\n'%(compact,sample_name))
 73     sys.stderr.write('\n### All %s jobs have been pushed into the queue ###\n'%job_num)
 74     return data_type_hash
 75 
 76 def create_processes(concurrency,jobs,work_path,results):
 77     print '\nBegin create jobs with %s Process...\n'%concurrency
 78     for _ in range(concurrency):
 79         process = Process(target=worker,args=(work_path,jobs,results))
 80         process.daemon = True
 81         process.start()
 82 
 83 def main(work_path,sam_barcode_file_list):
 84     global concurrency
 85     split_path = '%s/Split'%work_path
 86     QC_path = '%s/QC'%work_path
 87     jobs = JoinableQueue()
 88     results = Queue()
 89 
 90     canceled = False
 91     data_type_hash = add_jobs(split_path,sam_barcode_file_list,jobs)
 92     create_processes(concurrency,jobs,work_path,results)
 93     try:
 94         jobs.join()
 95     except KeyboardInterrupt:
 96         sys.stderr.write('cancelling ...\n')
 97         canceled = True
 98     finally:
 99         job_num = 0
100         finished_hash = {}
101         while not results.empty():
102             result = results.get_nowait()
103             job_num += 1
104             if not finished_hash.has_key(result.compact):
105                 finished_hash[result.compact] = []
106             finished_hash[result.compact].append(result.sample_name)
107         sys.stderr.write('all %s work finished!\n\n'%job_num)
108         log_out = open('%s/work.log'%QC_path,'w')
109         for compact,sample_list in finished_hash.iteritems():
110             for sample_name in sample_list:
111                 log_out.write('%s\t%s has been finished\n'%(compact,sample_name))
112         log_out.close()
113     if canceled:
114         return False
115 
116     for compact in os.listdir(QC_path):
117         compact_dir = '%s/%s'%(QC_path,compact)
118         if not os.path.isdir(compact_dir):
119             continue
120         sys.stderr.write('Begin stat compact: %s\n'%compact)
121         reads_stat(compact_dir)
122     sys.stderr.write('All campact stat finished!\n\n')
123 
124     reads_stat_all(QC_path,split_path)
125 
126     merge_threads = set()
127     for compact,subitem in data_type_hash.iteritems():
128         compact_dir = '%s/%s'%(QC_path,compact)
129         for data_type,sample_list in subitem.iteritems():
130             merged_file = '%s/%s/%s.together.fna'%(QC_path,compact,data_type)
131             t = Thread(target=sampleMerge,args=(sample_list,data_type,compact_dir,merged_file))
132             merge_threads.add(t)
133             t.start()
134             while True:
135                 if threading.activeCount() < concurrency:
136                     break
137     for t in threading.enumerate():
138         if t in merge_threads:
139             t.join()
140 
141     sys.stderr.write('\n All pipeline is done ! \n')
142 
143 if __name__ == '__main__':
144     sys.argv.pop(0)
145     if len(sys.argv) < 1:
146         sys.stderr.write('Usage: python run_pipeline.py work_path [process_num] \n process_num default is cpu_count\n')
147         sys.exit()
148     work_path = sys.argv.pop(0)
149     work_path = os.path.abspath(work_path)
150     sys.stderr.write('Workdir is %s,pipeline begin\n'%work_path)
151     sam_barcode_file_list = os.popen('ls %s/Split/sam_barcode.*'%work_path).read().strip().split('\n')
152     if len(sys.argv) != 0:
153         concurrency = int(sys.argv.pop(0))
154     else:
155         concurrency = cpu_count()
156 
157     main(work_path,sam_barcode_file_list)

下面是一些辅助程序：

  1 from __future__ import division
  2 from threading import Thread,Lock
  3 from multiprocessing import cpu_count
  4 import threading
  5 import sys
  6 import os
  7 import re
  8 import types
  9 from Bio import SeqIO
 10 from Bio.SeqRecord import SeqRecord
 11 def fq_reads_num(fq_file):
 12     wc_out = os.popen('wc -l %s'%fq_file).read().strip()
 13     result = int(re.search('^(\d+)',wc_out).group(1)) / 4
 14     return int(result)
 15 
 16 def Q_ave(self):
 17     Q_sum = 0
 18     for qlist in  self.letter_annotations.itervalues():
 19         for q in qlist:
 20             Q_sum += q
 21         Q_ave = Q_sum / len(self)
 22         return Q_ave
 23 
 24 def QC(file,out_file,out_stat_file,data_type):
 25     SeqRecord.Q_ave = Q_ave
 26     out_stat = open(out_stat_file,'w')
 27     out = open(out_file,'w')
 28 
 29     count = 0
 30     high_count = 0
 31     for record in SeqIO.parse(open(file),'fastq'):
 32         count += 1
 33         if record.Q_ave() < 20:
 34             continue
 35         if len(record) < 220 or len(record) > 500:
 36             continue
 37         out.write(record.format('fastq'))
 38         high_count += 1
 39     high_ratio = high_count / count
 40     out_stat.write('%s\t%s\t%s\t%s\n'%(data_type,count,high_count,high_ratio))
 41 
 42 class MyList(list):
 43     def __str__(self):
 44         out_str = ''
 45         for item in self:
 46             out_str += item
 47             out_str += '\t'
 48         return out_str.strip()
 49 
 50 def parse_stat(stat_file):
 51     tabs = os.popen('cat %s'%stat_file).read().strip().split('\t')
 52     yield tabs
 53 
 54 def parse_stat_files(compact_path):
 55     for f in os.popen('ls %s/*/*.stat'%compact_path):
 56         stat_file = f.strip()
 57         sample_name =  re.search('%s\/(\S+)\/high_quality\.stat'%compact_path,stat_file).group(1)
 58         yield stat_file,sample_name
 59 
 60 def reads_stat(compact_path):
 61     out = open('%s/reads_stat.xls'%compact_path,'w')
 62     sample_reads = {}
 63     out = open('%s/reads_stat.xls'%compact_path,'w')
 64     sample_reads = {}
 65     for stat_file,sample_name in parse_stat_files(compact_path):
 66         for tabs in parse_stat(stat_file):
 67             sample_reads[sample_name] = tabs
 68 
 69     out.write('sample_name\tsample_type\traw_reads\tHQ_reads\tHQ_ratio\n')
 70     for sample,tabs in sample_reads.iteritems():
 71         tabs = MyList(tabs)
 72         out.write('%s\t%s\n'%(sample,str(tabs)))
 73     out.close()
 74 
 75 def raw_stat_thread(fq_file,lock,compact,sample_name,tabs,out):
 76     global total_reads
 77 #    sys.stderr.write('thread %s stat with %s %s\n'%(threading.currentThread().ident,compact,sample_name))
 78     raw_reads = fq_reads_num(fq_file)
 79     lock.acquire()
 80     total_reads += raw_reads
 81     data_type = tabs.pop(0)
 82     ratio = int(tabs[1]) / raw_reads * 100
 83     tabs = str(MyList(tabs))
 84     out.write('%s\t%s\t%s\t%s\t%s\t%2.2f%%\n'%(compact,sample_name,data_type,raw_reads,tabs,ratio))
 85     lock.release()
 86 #    sys.stderr.write('thread %s finished doing with %s %s\n'%(threading.currentThread().ident,compact,sample_name))
 87 
 88 total_reads = 0
 89 def reads_stat_all(work_path,original_path):
 90     global total_reads
 91     sys.stderr.write('\nmerge stat is begin ... \n')
 92     out = open('%s/reads_stat.xls'%work_path,'w')
 93     compact_hash = {}
 94     for f in os.listdir(work_path):
 95         compact = f.strip()
 96         compact_path = '%s/%s'%(work_path,compact)
 97         if not os.path.isdir(compact_path):
 98             continue
 99         if not compact_hash.has_key(compact):
100             compact_hash[compact] = {}
101         for stat_file,sample_name in parse_stat_files(compact_path):
102             for tabs in parse_stat(stat_file):
103                 compact_hash[compact][sample_name] = tabs
104     out.write('compact\tsample_name\tdata_type\traw_reads\tpandaseq_reads\tHQ_reads\tratio\n')
105     lock = Lock()
106     active_threads = set()
107     for compact,sample in compact_hash.iteritems():
108         sys.stderr.write('doing with %s stat\n'%compact)
109         for sample_name,tabs in sample.iteritems():
110             original_fq = os.popen('ls %s/%s/%s/*'%(original_path,compact,sample_name)).read().strip().split('\n').pop(0)
111             t = Thread(target=raw_stat_thread,args=(original_fq,lock,compact,sample_name,tabs,out))
112             active_threads.add(t)
113             t.start()
114             while True:
115                 if threading.activeCount() < cpu_count():
116                     break
117     out.flush()
118     for t in threading.enumerate():
119         if t in active_threads:
120             sys.stderr.write('thread %s  is still alive, wait ...\n'%t.ident)
121             t.join()
122     sys.stderr.write('Unaligned stating ...\n')
123     out.write('\n###\n')
124     unalign_fq = os.popen('ls %s/Unalign/*'%original_path).read().strip().split('\n').pop(0)
125     unalign_reads = fq_reads_num(unalign_fq)
126     total_reads += unalign_reads
127     ratio = unalign_reads / total_reads * 100
128     out.write('Unalign\t%s\t%2.2f%%\n'%(unalign_reads,ratio))
129     out.close()
130     sys.stderr.write('merge stat is all finished!\n\n')
131 
132 def pandaseq(pandaseq_soft,read1,read2,fa_out,f_primer,r_primer,log_out):
133     cmd = '%s -F -f %s -r %s -w %s -p %s -q %s -g %s -l 220 -L 500'%(pandaseq_soft,read1,read2,fa_out,f_primer,r_primer,log_out)
134     os.system(cmd)
135 
136 def sampleMerge(sample_list,data_type,file_path,outfile):
137     outhandle = open(outfile,'w')
138 #    sys.stderr.write('Begin merge into %s\n'%file_path)
139     reads_num = {}
140     f_template = '%s/%s/high_quality.fq'
141     for sample in sample_list:
142         f = f_template%(file_path,sample)
143         sample = re.sub('[-_]','.',sample)
144         sample = '%s%s'%(data_type,sample)
145         if not reads_num.has_key(sample):
146             reads_num[sample] = 0
147         for record in SeqIO.parse(open(f),'fastq'):
148             reads_num[sample] += 1
149             outhandle.write('>%s_%s\n%s\n'%(sample,reads_num[sample],str(record.seq)))
150     outhandle.close()
151     sys.stderr.write('merge file: %s is finished\n'%outfile)
152 
153 def get_lib_method(file):
154     file = os.path.basename(file)
155     if re.match('^sam_barcode.l$',file):
156         lib_method = 'Self'
157     elif re.match('^sam_barcode.s\d+$',file):
158         lib_method = 'HXT'
159     else:
160         lib_method = None
161     return lib_method

settings.py中包含不同建库方式的引物序列。

这个程序也算是前几天的学习成果展示了