Skills › Scientific Agent Skills › pysam › alignment_qc.py
alignment_qc.py
337 lines · skills/pysam/scripts/alignment_qc.py at 5f2bc03
1#!/usr/bin/env python3
2"""Stream a local SAM/BAM/CRAM file and write aggregate QC counts as JSON.
4Counts are alignment-record counts, not unique templates. A whole-file scan
5uses fetch(until_eof=True) and needs no index. --region is a 1-based inclusive
6samtools region string and requires an index. CRAM requires --reference.
7"""
9from __future__ import annotations
11import argparse
12import json
13import sys
14from pathlib import Path
15from typing import Any, Iterator, Optional
17import pysam
20def positive_int(value: str) -> int:
21 parsed = int(value)
22 if parsed <= 0:
23 raise argparse.ArgumentTypeError("must be greater than zero")
24 return parsed
27def nonnegative_int(value: str) -> int:
28 parsed = int(value)
29 if parsed < 0:
30 raise argparse.ArgumentTypeError("must be zero or greater")
31 return parsed
34def build_parser() -> argparse.ArgumentParser:
35 parser = argparse.ArgumentParser(
36 description=__doc__,
37 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
38 )
39 parser.add_argument("input", type=Path, help="local SAM/BAM/CRAM file")
40 parser.add_argument(
41 "--reference",
42 type=Path,
43 help="local indexed reference FASTA; required for CRAM",
44 )
45 parser.add_argument(
46 "--index",
47 type=Path,
48 help="explicit local BAI/CSI/CRAI index",
49 )
50 parser.add_argument(
51 "--region",
52 help="1-based inclusive samtools region, for example chr1:1-1000000",
53 )
54 parser.add_argument(
55 "--threads",
56 type=positive_int,
57 default=1,
58 help="HTSlib decompression threads",
59 )
60 parser.add_argument(
61 "--max-records",
62 type=nonnegative_int,
63 default=0,
64 help="stop after this many records; 0 scans the complete selection",
65 )
66 parser.add_argument(
67 "--output",
68 type=Path,
69 help="write JSON to a new file instead of stdout; refuses overwrite",
70 )
71 return parser
74def require_local_file(path: Path, label: str) -> Path:
75 resolved = path.expanduser()
76 if not resolved.exists():
77 raise FileNotFoundError(f"{label} does not exist: {resolved}")
78 if not resolved.is_file():
79 raise ValueError(f"{label} is not a regular file: {resolved}")
80 return resolved
83def alignment_mode(path: Path) -> str:
84 name = path.name.lower()
85 if name.endswith(".sam"):
86 return "r"
87 if name.endswith(".cram"):
88 return "rc"
89 if name.endswith(".bam"):
90 return "rb"
91 raise ValueError("input suffix must be .sam, .bam, or .cram")
94def ratio(numerator: int, denominator: int) -> Optional[float]:
95 return numerator / denominator if denominator else None
98def average(total: int, count: int) -> Optional[float]:
99 return total / count if count else None
102def selected_reads(
103 alignments: pysam.AlignmentFile,
104 region: Optional[str],
105) -> Iterator[pysam.AlignedSegment]:
106 if region is None:
107 return alignments.fetch(until_eof=True)
108 return alignments.fetch(region=region)
111def empty_counts() -> dict[str, int]:
112 return {
113 "total_records": 0,
114 "primary_records": 0,
115 "mapped_records": 0,
116 "unmapped_records": 0,
117 "secondary_records": 0,
118 "supplementary_records": 0,
119 "duplicate_records": 0,
120 "qc_fail_records": 0,
121 "paired_records": 0,
122 "proper_pair_records": 0,
123 "read1_records": 0,
124 "read2_records": 0,
125 "reverse_records": 0,
126 "mate_unmapped_records": 0,
127 "cigar_missing_records": 0,
128 "sequence_missing_records": 0,
129 "mapq_255_records": 0,
130 }
133def inspect_alignments(args: argparse.Namespace) -> dict[str, Any]:
134 path = require_local_file(args.input, "input")
135 reference = (
136 require_local_file(args.reference, "reference")
137 if args.reference is not None
138 else None
139 )
140 index = (
141 require_local_file(args.index, "index")
142 if args.index is not None
143 else None
144 )
145 mode = alignment_mode(path)
147 if mode == "rc" and reference is None:
148 raise ValueError("CRAM QC requires --reference")
150 kwargs: dict[str, Any] = {
151 "threads": args.threads,
152 "require_index": args.region is not None,
153 }
154 if reference is not None:
155 kwargs["reference_filename"] = str(reference)
156 if index is not None:
157 kwargs["index_filename"] = str(index)
159 counts = empty_counts()
160 query_length_sum = 0
161 query_length_count = 0
162 aligned_query_bases = 0
163 mapq_sum = 0
164 mapq_count = 0
165 record_limit_reached = False
167 with pysam.AlignmentFile(str(path), mode, **kwargs) as alignments:
168 header = alignments.header.to_dict()
169 for read in selected_reads(alignments, args.region):
170 if args.max_records and counts["total_records"] >= args.max_records:
171 record_limit_reached = True
172 break
174 counts["total_records"] += 1
176 is_primary = not read.is_secondary and not read.is_supplementary
177 if is_primary:
178 counts["primary_records"] += 1
179 if read.is_unmapped:
180 counts["unmapped_records"] += 1
181 else:
182 counts["mapped_records"] += 1
183 if read.mapping_quality == 255:
184 counts["mapq_255_records"] += 1
185 else:
186 mapq_sum += int(read.mapping_quality)
187 mapq_count += 1
188 if read.is_secondary:
189 counts["secondary_records"] += 1
190 if read.is_supplementary:
191 counts["supplementary_records"] += 1
192 if read.is_duplicate:
193 counts["duplicate_records"] += 1
194 if read.is_qcfail:
195 counts["qc_fail_records"] += 1
196 if read.is_paired:
197 counts["paired_records"] += 1
198 if read.is_proper_pair:
199 counts["proper_pair_records"] += 1
200 if read.is_read1:
201 counts["read1_records"] += 1
202 if read.is_read2:
203 counts["read2_records"] += 1
204 if read.is_reverse:
205 counts["reverse_records"] += 1
206 if read.mate_is_unmapped:
207 counts["mate_unmapped_records"] += 1
208 if read.cigartuples is None:
209 counts["cigar_missing_records"] += 1
210 if read.query_sequence is None:
211 counts["sequence_missing_records"] += 1
213 # query_length is 0, not None, when SEQ is absent. Do not make
214 # missing sequence look like a measured zero-length read.
215 if read.query_sequence is not None:
216 query_length_sum += int(read.query_length)
217 query_length_count += 1
218 # pysam reports the full query length for unmapped/CIGAR-less
219 # records; that is not evidence of aligned bases.
220 if (
221 not read.is_unmapped
222 and read.cigartuples is not None
223 and read.query_sequence is not None
224 ):
225 aligned_query_bases += int(read.query_alignment_length)
227 try:
228 has_index = bool(alignments.has_index())
229 except (AttributeError, OSError, ValueError):
230 has_index = False
232 report = {
233 "schema_version": "1.0",
234 "pysam_version": pysam.__version__,
235 "samtools_version": getattr(
236 pysam, "__samtools_version__", None
237 ),
238 "file_name": path.name,
239 "format": (
240 "CRAM"
241 if alignments.is_cram
242 else "BAM"
243 if alignments.is_bam
244 else "SAM"
245 ),
246 "sort_order": header.get("HD", {}).get("SO"),
247 "has_index": has_index,
248 "selection": {
249 "region": args.region,
250 "region_coordinate_system": (
251 "1-based inclusive samtools string"
252 if args.region is not None
253 else None
254 ),
255 "whole_file_order": args.region is None,
256 "max_records": args.max_records,
257 "record_limit_reached": record_limit_reached,
258 },
259 "semantics": (
260 "Counts are alignment records, not unique query names, "
261 "templates, or fragments. MAPQ 255 is reported as unavailable "
262 "and excluded from mean_mapping_quality. Mean query length "
263 "excludes missing SEQ. Aligned query bases count non-soft-clipped "
264 "query bases (including insertions) only in mapped records "
265 "with CIGAR and SEQ; they are not reference coverage."
266 ),
267 "counts": counts,
268 "derived": {
269 "mapping_rate_all_records": ratio(
270 counts["mapped_records"],
271 counts["total_records"],
272 ),
273 "duplicate_rate_all_records": ratio(
274 counts["duplicate_records"],
275 counts["total_records"],
276 ),
277 "proper_pair_rate_paired_records": ratio(
278 counts["proper_pair_records"],
279 counts["paired_records"],
280 ),
281 "mean_query_length": average(
282 query_length_sum,
283 query_length_count,
284 ),
285 "mean_mapping_quality": average(mapq_sum, mapq_count),
286 "aligned_query_bases": aligned_query_bases,
287 },
288 }
289 return report
292def write_report(report: dict[str, Any], output: Optional[Path]) -> None:
293 text = json.dumps(
294 report,
295 indent=2,
296 ensure_ascii=False,
297 allow_nan=False,
298 )
299 if output is None:
300 sys.stdout.write(text + "\n")
301 return
303 destination = output.expanduser()
304 with destination.open("x", encoding="utf-8") as handle:
305 handle.write(text + "\n")
308def main(argv: Optional[list[str]] = None) -> int:
309 parser = build_parser()
310 args = parser.parse_args(argv)
312 if args.output is not None:
313 try:
314 if args.output.expanduser().resolve() == args.input.expanduser().resolve():
315 parser.error("--output must not overwrite the input file")
316 except OSError:
317 pass
319 try:
320 report = inspect_alignments(args)
321 write_report(report, args.output)
322 except (
323 FileExistsError,
324 FileNotFoundError,
325 IndexError,
326 KeyError,
327 OSError,
328 TypeError,
329 ValueError,
330 ) as error:
331 print(f"alignment_qc: {error}", file=sys.stderr)
332 return 2
333 return 0
336if __name__ == "__main__":
337 raise SystemExit(main())