forked from BarasLab/miRge
-
Notifications
You must be signed in to change notification settings - Fork 0
/
trim_file.py
151 lines (129 loc) · 5 KB
/
trim_file.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
from distutils.version import StrictVersion
import argparse
import sys
from multiprocessing import Process, Queue
from cutadapt.adapters import Adapter, gather_adapters
from cutadapt.scripts.cutadapt import AdapterCutter
from cutadapt.modifiers import QualityTrimmer, UnconditionalCutter
from cutadapt.seqio import FastqReader
import cutadapt
if StrictVersion(cutadapt.__version__) < StrictVersion('1.8.1'):
raise ImportError('miRge requires cutadapt >= version 1.8.1 to operate.')
class Worker(Process):
def __init__(self, queue=None, results=None, adapter=None, phred64=False):
super(Worker, self).__init__()
self.queue=queue
self.results = results
self.phred = 64 if phred64 else 33
self.modifiers = [QualityTrimmer(0, 10, self.phred)]
self.adapters = []
self.error_rate = 0.12
self.min_length = 16
if adapter.startswith('+'):
self.modifiers.append(UnconditionalCutter(int(adapter)))
elif adapter == 'none':
self.adapter = None
else:
for name,seq,where in gather_adapters(adapter.split(','), [], []):
self.adapters.append(Adapter(seq, where, self.error_rate, name=name))
adapter_cutter = AdapterCutter(self.adapters)
self.modifiers.append(adapter_cutter)
def run(self):
# we can't use the sentinel iter(self.queue.get, None) because of some issue with Sequence classes
results = self.results
get_func = self.queue.get
reads = get_func()
modifiers = self.modifiers
min_length = self.min_length
while reads is not None:
result_batch = []
for read in reads:
for modifier in modifiers:
read = modifier(read)
if len(read.sequence) >= min_length:
try:
second_header = read.name2
except AttributeError:
second_header = read.name if getattr(read, 'twoheaders', False) else ''
if second_header:
read_out = '@' + read.name + '\n' + read.sequence + '\n+' + second_header + '\n' + read.qualities + '\n'
else:
read_out = '@' + read.name + '\n' + read.sequence + '\n+\n' + read.qualities + '\n'
result_batch.append(read_out)
results.put(result_batch)
reads = get_func()
class Writer(Process):
def __init__(self, queue=None, trimmed=None, outfile=None):
super(Writer, self).__init__()
self.queue = queue
self.trimmed = trimmed
self.outfile = outfile
def run(self):
get_func = self.queue.get
reads = get_func()
with self.outfile as outfile:
kept = 0
while reads is not None:
for read in reads:
outfile.write(read)
kept += 1
reads = get_func()
self.trimmed.put(kept)
parser = argparse.ArgumentParser()
parser.add_argument('--cutadapt', type=str)
parser.add_argument('--adapter', type=str)
parser.add_argument('--infile', type=argparse.FileType('rb'))
parser.add_argument('--outfile', type=argparse.FileType('wb'))
parser.add_argument('--threads', type=int, default=1)
parser.add_argument('--phred64', action='store_true')
def main():
args = parser.parse_args()
dest = args.outfile
phred = args.phred64 or 33
threads = args.threads
logfile = args.infile.name
adapter = args.adapter
read_queue = Queue()
result_queue = Queue()
trimmed_queue = Queue()
workers = []
def start_workers():
for i in xrange(threads):
worker = Worker(queue=read_queue, results=result_queue, phred64=phred==64, adapter=adapter)
workers.append(worker)
worker.start()
writer = Writer(queue=result_queue, trimmed=trimmed_queue, outfile=dest)
writer.start()
batch = []
for index, read in enumerate(FastqReader(args.infile.name)):
batch.append(read)
if index < 1000 and phred == 33:
if any([i for i in read.qualities if ord(i) > 74]):
phred = 64
if index % 10000 == 0:
if not workers:
start_workers()
read_queue.put(batch)
batch = []
if not workers:
start_workers()
read_queue.put(batch)
processed = index+1
# poison pill to stop workers
for i in xrange(threads):
read_queue.put(None)
for i in workers:
i.join()
# poison pill for writers
result_queue.put(None)
# wait for writing to finish
writer.join()
trimmed_queue.put(None)
kept_reads = sum([i for i in iter(trimmed_queue.get, None)])
with open('{0}.log'.format(logfile), 'wb') as o:
o.write('Starting reads: {0}\n'.format(processed))
o.write('Processed reads: {0}\n'.format(kept_reads))
sys.stdout.write('{0}\n'.format(phred))
return phred
if __name__ == "__main__":
sys.exit(main())