Skills › Scientific Agent Skills › pysam › filter_alignments.py
filter_alignments.py
387 lines · skills/pysam/scripts/filter_alignments.py at 5f2bc03
1#!/usr/bin/env python3
2"""Filter a local SAM/BAM/CRAM into a new alignment file.
4The script preserves record order and the input header; it does not sort.
5Whole-file iteration includes unplaced unmapped records and needs no index.
6--region is a 1-based inclusive samtools region and requires an input index.
7CRAM input or output requires --reference. Existing outputs are never replaced.
8"""
10from __future__ import annotations
12import argparse
13import json
14import sys
15from collections import Counter
16from pathlib import Path
17from typing import Any, Iterator, Optional
19import pysam
20import pysam.samtools
23def positive_int(value: str) -> int:
24 parsed = int(value)
25 if parsed <= 0:
26 raise argparse.ArgumentTypeError("must be greater than zero")
27 return parsed
30def nonnegative_int(value: str) -> int:
31 parsed = int(value)
32 if parsed < 0:
33 raise argparse.ArgumentTypeError("must be zero or greater")
34 return parsed
37def build_parser() -> argparse.ArgumentParser:
38 parser = argparse.ArgumentParser(
39 description=__doc__,
40 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
41 )
42 parser.add_argument("input", type=Path, help="local SAM/BAM/CRAM input")
43 parser.add_argument("output", type=Path, help="new SAM/BAM/CRAM output")
44 parser.add_argument(
45 "--reference",
46 type=Path,
47 help="local indexed reference FASTA; required for CRAM",
48 )
49 parser.add_argument(
50 "--input-index",
51 type=Path,
52 help="explicit local BAI/CSI/CRAI input index",
53 )
54 parser.add_argument(
55 "--region",
56 help="1-based inclusive samtools region, for example chr1:1-1000000",
57 )
58 parser.add_argument(
59 "--min-mapq",
60 type=nonnegative_int,
61 default=0,
62 help="minimum mapping quality; MAPQ 255 passes this numeric threshold",
63 )
64 parser.add_argument(
65 "--exclude-unmapped",
66 action="store_true",
67 help="exclude unmapped records",
68 )
69 parser.add_argument(
70 "--exclude-secondary",
71 action="store_true",
72 help="exclude secondary alignments",
73 )
74 parser.add_argument(
75 "--exclude-supplementary",
76 action="store_true",
77 help="exclude supplementary alignments",
78 )
79 parser.add_argument(
80 "--exclude-duplicates",
81 action="store_true",
82 help="exclude records marked duplicate",
83 )
84 parser.add_argument(
85 "--exclude-qcfail",
86 action="store_true",
87 help="exclude records marked QC fail",
88 )
89 parser.add_argument(
90 "--proper-pairs-only",
91 action="store_true",
92 help="keep only records marked as proper paired alignments",
93 )
94 parser.add_argument(
95 "--threads",
96 type=positive_int,
97 default=1,
98 help="HTSlib compression/decompression threads",
99 )
100 parser.add_argument(
101 "--index",
102 action="store_true",
103 help="index output; requires coordinate sort order and BAM/CRAM",
104 )
105 parser.add_argument(
106 "--csi",
107 action="store_true",
108 help="use CSI for BAM output (requires --index); supports large contigs",
109 )
110 parser.add_argument(
111 "--summary",
112 type=Path,
113 help="write JSON summary to a new file instead of stdout",
114 )
115 return parser
118def require_local_file(path: Path, label: str) -> Path:
119 resolved = path.expanduser()
120 if not resolved.exists():
121 raise FileNotFoundError(f"{label} does not exist: {resolved}")
122 if not resolved.is_file():
123 raise ValueError(f"{label} is not a regular file: {resolved}")
124 return resolved
127def alignment_mode(path: Path, *, writing: bool) -> str:
128 name = path.name.lower()
129 if name.endswith(".sam"):
130 return "w" if writing else "r"
131 if name.endswith(".bam"):
132 return "wb" if writing else "rb"
133 if name.endswith(".cram"):
134 return "wc" if writing else "rc"
135 raise ValueError("alignment suffix must be .sam, .bam, or .cram")
138def selected_reads(
139 alignments: pysam.AlignmentFile,
140 region: Optional[str],
141) -> Iterator[pysam.AlignedSegment]:
142 if region is None:
143 return alignments.fetch(until_eof=True)
144 return alignments.fetch(region=region)
147def exclusion_reasons(
148 read: pysam.AlignedSegment,
149 args: argparse.Namespace,
150) -> list[str]:
151 reasons = []
152 if args.exclude_unmapped and read.is_unmapped:
153 reasons.append("unmapped")
154 if args.exclude_secondary and read.is_secondary:
155 reasons.append("secondary")
156 if args.exclude_supplementary and read.is_supplementary:
157 reasons.append("supplementary")
158 if args.exclude_duplicates and read.is_duplicate:
159 reasons.append("duplicate")
160 if args.exclude_qcfail and read.is_qcfail:
161 reasons.append("qc_fail")
162 if args.proper_pairs_only and not read.is_proper_pair:
163 reasons.append("not_proper_pair")
164 if read.mapping_quality < args.min_mapq:
165 reasons.append("mapq_below_threshold")
166 return reasons
169def validate_destinations(
170 input_path: Path,
171 output_path: Path,
172 summary_path: Optional[Path],
173) -> None:
174 if not output_path.parent.exists():
175 raise FileNotFoundError(
176 f"output parent does not exist: {output_path.parent}"
177 )
178 if output_path.exists() or output_path.is_symlink():
179 raise FileExistsError(f"output already exists: {output_path}")
180 if input_path.resolve() == output_path.resolve():
181 raise ValueError("output must not be the input file")
183 # Stale sidecars can be silently opened for a newly written alignment;
184 # samtools index can also overwrite them. Refuse both naming conventions.
185 suffixes = (".crai",) if output_path.suffix.lower() == ".cram" else (".bai", ".csi")
186 sidecars = {
187 candidate
188 for suffix in suffixes
189 for candidate in (Path(str(output_path) + suffix), output_path.with_suffix(suffix))
190 }
191 for candidate in sidecars:
192 if candidate.exists() or candidate.is_symlink():
193 raise FileExistsError(f"output index already exists: {candidate}")
195 if summary_path is not None:
196 if not summary_path.parent.exists():
197 raise FileNotFoundError(
198 f"summary parent does not exist: {summary_path.parent}"
199 )
200 if summary_path.exists() or summary_path.is_symlink():
201 raise FileExistsError(
202 f"summary output already exists: {summary_path}"
203 )
204 if summary_path.resolve() in {
205 input_path.resolve(),
206 output_path.resolve(),
207 *(candidate.resolve() for candidate in sidecars),
208 }:
209 raise ValueError("summary path must differ from input, output, and indexes")
212def filter_file(args: argparse.Namespace) -> dict[str, Any]:
213 input_path = require_local_file(args.input, "input")
214 output_path = args.output.expanduser()
215 summary_path = (
216 args.summary.expanduser() if args.summary is not None else None
217 )
218 validate_destinations(input_path, output_path, summary_path)
220 reference = (
221 require_local_file(args.reference, "reference")
222 if args.reference is not None
223 else None
224 )
225 input_index = (
226 require_local_file(args.input_index, "input index")
227 if args.input_index is not None
228 else None
229 )
230 input_mode = alignment_mode(input_path, writing=False)
231 output_mode = alignment_mode(output_path, writing=True)
233 if (input_mode == "rc" or output_mode == "wc") and reference is None:
234 raise ValueError("CRAM input or output requires --reference")
235 if args.index and output_mode == "w":
236 raise ValueError("SAM output cannot be indexed")
237 if args.csi and (not args.index or output_mode != "wb"):
238 raise ValueError("--csi requires --index and BAM output")
240 input_kwargs: dict[str, Any] = {
241 "threads": args.threads,
242 "require_index": args.region is not None,
243 }
244 if reference is not None:
245 input_kwargs["reference_filename"] = str(reference)
246 if input_index is not None:
247 input_kwargs["index_filename"] = str(input_index)
249 output_kwargs: dict[str, Any] = {"threads": args.threads}
250 if reference is not None:
251 output_kwargs["reference_filename"] = str(reference)
253 reason_counts: Counter[str] = Counter()
254 selected = 0
255 written = 0
256 excluded = 0
257 sort_order: Optional[str] = None
258 output_created = False
260 try:
261 with pysam.AlignmentFile(
262 str(input_path),
263 input_mode,
264 **input_kwargs,
265 ) as source:
266 sort_order = source.header.to_dict().get("HD", {}).get("SO")
267 if args.index and sort_order != "coordinate":
268 raise ValueError(
269 "--index requires input header sort order 'coordinate'; "
270 "this script preserves order but does not sort"
271 )
273 with output_path.open("xb") as output_handle:
274 output_created = True
275 with pysam.AlignmentFile(
276 output_handle,
277 output_mode,
278 template=source,
279 **output_kwargs,
280 ) as destination:
281 for read in selected_reads(source, args.region):
282 selected += 1
283 reasons = exclusion_reasons(read, args)
284 if reasons:
285 excluded += 1
286 reason_counts.update(reasons)
287 continue
288 destination.write(read)
289 written += 1
290 except Exception:
291 if output_created and output_path.exists():
292 output_path.unlink()
293 raise
295 if args.index:
296 pysam.samtools.index(
297 *(["-c"] if args.csi else []),
298 "-@",
299 str(args.threads),
300 str(output_path),
301 catch_stdout=False,
302 )
304 return {
305 "schema_version": "1.0",
306 "pysam_version": pysam.__version__,
307 "samtools_version": getattr(pysam, "__samtools_version__", None),
308 "input_file_name": input_path.name,
309 "output_file_name": output_path.name,
310 "input_mode": input_mode,
311 "output_mode": output_mode,
312 "sort_order_preserved": sort_order,
313 "selection": {
314 "region": args.region,
315 "region_coordinate_system": (
316 "1-based inclusive samtools string"
317 if args.region is not None
318 else None
319 ),
320 },
321 "filters": {
322 "min_mapq": args.min_mapq,
323 "exclude_unmapped": args.exclude_unmapped,
324 "exclude_secondary": args.exclude_secondary,
325 "exclude_supplementary": args.exclude_supplementary,
326 "exclude_duplicates": args.exclude_duplicates,
327 "exclude_qcfail": args.exclude_qcfail,
328 "proper_pairs_only": args.proper_pairs_only,
329 },
330 "counts": {
331 "selected_records": selected,
332 "written_records": written,
333 "excluded_records": excluded,
334 "exclusion_reasons": dict(sorted(reason_counts.items())),
335 },
336 "indexed_output": args.index,
337 "index_format": ("CSI" if args.csi else "BAI" if output_mode == "wb" else "CRAI") if args.index else None,
338 "semantics": (
339 "Exclusion reason counts can exceed excluded_records because one "
340 "record can fail multiple filters. Counts are alignment records. "
341 "Filtering is per record: mates outside the selection are not "
342 "retrieved, and mate flags/coordinates/TLEN are not recomputed."
343 ),
344 }
347def write_summary(
348 report: dict[str, Any],
349 output: Optional[Path],
350) -> None:
351 text = json.dumps(
352 report,
353 indent=2,
354 ensure_ascii=False,
355 allow_nan=False,
356 )
357 if output is None:
358 sys.stdout.write(text + "\n")
359 return
360 with output.expanduser().open("x", encoding="utf-8") as handle:
361 handle.write(text + "\n")
364def main(argv: Optional[list[str]] = None) -> int:
365 parser = build_parser()
366 args = parser.parse_args(argv)
368 try:
369 report = filter_file(args)
370 write_summary(report, args.summary)
371 except (
372 FileExistsError,
373 FileNotFoundError,
374 IndexError,
375 KeyError,
376 OSError,
377 TypeError,
378 ValueError,
379 pysam.SamtoolsError,
380 ) as error:
381 print(f"filter_alignments: {error}", file=sys.stderr)
382 return 2
383 return 0
386if __name__ == "__main__":
387 raise SystemExit(main())