Coverage for src / optwps / read_processing.py: 97%
145 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-25 22:34 +0200
1import random
2from dataclasses import dataclass
4import joblib
5import numpy as np
7from .utils import is_soft_clipped, ref_aln_length
10@dataclass
11class ReadInfo:
12 is_duplicate: bool
13 is_qcfail: bool
14 is_unmapped: bool
15 cigartuples: tuple
16 reference_name: str
17 reference_start: int
18 reference_end: int
19 is_paired: bool
20 mate_is_unmapped: bool
21 rnext: int
22 tid: int
23 is_read1: bool
24 is_read2: bool
25 pnext: int
26 qlen: int
27 template_length: int
28 pos: int
29 query_sequence: str
32def read_info(read):
33 try:
34 reference_name = read.reference_name
35 except ValueError:
36 reference_name = None
37 return ReadInfo(
38 read.is_duplicate,
39 read.is_qcfail,
40 read.is_unmapped,
41 tuple(read.cigartuples or ()),
42 reference_name,
43 read.reference_start,
44 read.reference_end,
45 read.is_paired,
46 read.mate_is_unmapped,
47 read.rnext,
48 read.tid,
49 read.is_read1,
50 read.is_read2,
51 read.pnext,
52 read.qlen,
53 read.template_length,
54 read.pos,
55 read.query_sequence or "",
56 )
59def read_info_batches(reads, buffer_size):
60 batch = []
61 for read in reads:
62 batch.append(read_info(read))
63 if len(batch) >= buffer_size:
64 yield batch
65 batch = []
66 if batch:
67 yield batch
70def iter_pysam_reads(reads):
71 while True:
72 try:
73 yield next(reads)
74 except StopIteration:
75 return
76 except ValueError as error:
77 if "Firing event 10 with no exception set" not in str(error):
78 raise
79 return
82def valid_read_info(
83 read,
84 min_insert_size=None,
85 max_insert_size=None,
86 mappability_file=None,
87 min_mappability_threshold=0.9,
88 upstream_limit=None,
89 downsample_ratio=None,
90):
91 if read.is_duplicate or read.is_qcfail or read.is_unmapped:
92 return False
93 if is_soft_clipped(read.cigartuples):
94 return False
95 if mappability_file is not None:
96 if read.reference_name is None or read.reference_end is None:
97 return False
98 try:
99 mappability = mappability_file.values(
100 read.reference_name, read.reference_start, read.reference_end
101 )
102 if not mappability or any(value != value for value in mappability):
103 return False
104 if min_mappability_threshold is not None and (
105 sum(mappability) / len(mappability) < min_mappability_threshold
106 ):
107 return False
108 except RuntimeError:
109 return False
110 if read.is_paired:
111 if read.mate_is_unmapped or read.rnext != read.tid:
112 return False
113 if not (
114 read.is_read1
115 or (
116 read.is_read2
117 and upstream_limit is not None
118 and read.pnext + read.qlen < upstream_limit
119 )
120 ):
121 return False
122 lseq = abs(read.template_length)
123 if lseq == 0:
124 return False
125 else:
126 lseq = ref_aln_length(read.cigartuples)
127 if downsample_ratio is not None and random.random() >= downsample_ratio:
128 return False
129 if min_insert_size is not None and lseq < min_insert_size:
130 return False
131 if max_insert_size is not None and lseq > max_insert_size:
132 return False
133 return True
136def fragment_interval(read):
137 if read.is_paired:
138 lseq = abs(read.template_length)
139 rstart = min(read.pos, read.pnext)
140 else:
141 lseq = ref_aln_length(read.cigartuples)
142 rstart = read.pos
143 return rstart, rstart + lseq - 1
146def fragment_features(read):
147 fragment_length = (
148 abs(read.template_length)
149 if read.is_paired
150 else ref_aln_length(read.cigartuples)
151 )
152 sequence = read.query_sequence
153 if not sequence:
154 return [fragment_length, 0.0]
155 gc_count = sum(1 for base in sequence if base in "GCgc")
156 return [fragment_length, gc_count / len(sequence)]
159def weight_from_features(features, bin_edges, weights):
160 if weights is None:
161 return 1.0
162 bin_indices = []
163 for i, feature in enumerate(features):
164 bin_index = np.searchsorted(bin_edges[i], feature, side="right") - 1
165 bin_indices.append(np.clip(bin_index, 0, weights.shape[i] - 1))
166 return weights[tuple(bin_indices)]
169def _open_mappability(path):
170 if path is None:
171 return None
172 import pyBigWig
174 return pyBigWig.open(path)
177def valid_fragment_features_batch(
178 reads,
179 min_insert_size=None,
180 max_insert_size=None,
181 mappability_path=None,
182 min_mappability_threshold=0.9,
183 downsample_ratio=None,
184):
185 mappability_file = _open_mappability(mappability_path)
186 try:
187 return [
188 fragment_features(read)
189 for read in reads
190 if valid_read_info(
191 read,
192 min_insert_size=min_insert_size,
193 max_insert_size=max_insert_size,
194 mappability_file=mappability_file,
195 min_mappability_threshold=min_mappability_threshold,
196 downsample_ratio=downsample_ratio,
197 )
198 ]
199 finally:
200 if mappability_file is not None:
201 mappability_file.close()
204def valid_fragment_intervals_batch(
205 reads,
206 min_insert_size=None,
207 max_insert_size=None,
208 mappability_path=None,
209 min_mappability_threshold=0.9,
210 upstream_limit=None,
211 downsample_ratio=None,
212 bin_edges=None,
213 weight_values=None,
214 use_weights=False,
215):
216 mappability_file = _open_mappability(mappability_path)
217 starts, ends, weights = [], [], []
218 try:
219 for read in reads:
220 if not valid_read_info(
221 read,
222 min_insert_size=min_insert_size,
223 max_insert_size=max_insert_size,
224 mappability_file=mappability_file,
225 min_mappability_threshold=min_mappability_threshold,
226 upstream_limit=upstream_limit,
227 downsample_ratio=downsample_ratio,
228 ):
229 continue
230 start, end = fragment_interval(read)
231 starts.append(start)
232 ends.append(end)
233 if use_weights:
234 features = fragment_features(read)
235 weights.append(weight_from_features(features, bin_edges, weight_values))
236 finally:
237 if mappability_file is not None:
238 mappability_file.close()
239 return starts, ends, weights
242def process_read_batches(batches, njobs, batch_func, **kwargs):
243 if njobs == 1:
244 for batch in batches:
245 yield batch_func(batch, **kwargs)
246 else:
247 yield from joblib.Parallel(n_jobs=njobs)(
248 joblib.delayed(batch_func)(batch, **kwargs) for batch in batches
249 )
252def collect_fragment_features(
253 batches,
254 min_insert_size=None,
255 max_insert_size=None,
256 mappability_path=None,
257 min_mappability_threshold=0.9,
258 downsample_ratio=None,
259 njobs=1,
260 read_buffer_size=10000,
261):
262 return [
263 feature
264 for batch in process_read_batches(
265 batches,
266 njobs,
267 valid_fragment_features_batch,
268 min_insert_size=min_insert_size,
269 max_insert_size=max_insert_size,
270 mappability_path=mappability_path,
271 min_mappability_threshold=min_mappability_threshold,
272 downsample_ratio=downsample_ratio,
273 )
274 for feature in batch
275 ]
278def collect_fragment_intervals(
279 batches,
280 min_insert_size=None,
281 max_insert_size=None,
282 mappability_path=None,
283 min_mappability_threshold=0.9,
284 upstream_limit=None,
285 downsample_ratio=None,
286 bin_edges=None,
287 weight_values=None,
288 use_weights=False,
289 njobs=1,
290 read_buffer_size=10000,
291):
292 starts, ends, weights = [], [], []
293 for read_starts, read_ends, read_weights in process_read_batches(
294 batches,
295 njobs,
296 valid_fragment_intervals_batch,
297 min_insert_size=min_insert_size,
298 max_insert_size=max_insert_size,
299 mappability_path=mappability_path,
300 min_mappability_threshold=min_mappability_threshold,
301 upstream_limit=upstream_limit,
302 downsample_ratio=downsample_ratio,
303 bin_edges=bin_edges,
304 weight_values=weight_values,
305 use_weights=use_weights,
306 ):
307 starts.extend(read_starts)
308 ends.extend(read_ends)
309 weights.extend(read_weights)
310 return starts, ends, weights