Skills › Scientific Agent Skills › pysam › variant_summary.py
variant_summary.py
377 lines · skills/pysam/scripts/variant_summary.py at 5f2bc03
1#!/usr/bin/env python3
2"""Stream a local VCF/BCF and emit variant/genotype summary counts as JSON.
4Normal iteration scans in file order without an index. --region is a 1-based
5inclusive samtools region string and requires TBI/CSI. Sample identifiers are
6omitted unless --include-sample-names is supplied.
7"""
9from __future__ import annotations
11import argparse
12import json
13import sys
14from collections import Counter
15from pathlib import Path
16from typing import Any, Iterator, Optional
18import pysam
21def positive_int(value: str) -> int:
22 parsed = int(value)
23 if parsed <= 0:
24 raise argparse.ArgumentTypeError("must be greater than zero")
25 return parsed
28def nonnegative_int(value: str) -> int:
29 parsed = int(value)
30 if parsed < 0:
31 raise argparse.ArgumentTypeError("must be zero or greater")
32 return parsed
35def build_parser() -> argparse.ArgumentParser:
36 parser = argparse.ArgumentParser(
37 description=__doc__,
38 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
39 )
40 parser.add_argument("input", type=Path, help="local VCF/VCF.gz/BCF file")
41 parser.add_argument(
42 "--index",
43 type=Path,
44 help="explicit local TBI/CSI index",
45 )
46 parser.add_argument(
47 "--region",
48 help="1-based inclusive samtools region, for example chr1:1-1000000",
49 )
50 parser.add_argument(
51 "--sample",
52 action="append",
53 default=[],
54 help="sample to decode; repeat for multiple samples (default: all)",
55 )
56 parser.add_argument(
57 "--include-sample-names",
58 action="store_true",
59 help="include decoded sample identifiers in the JSON report",
60 )
61 parser.add_argument(
62 "--threads",
63 type=positive_int,
64 default=1,
65 help="HTSlib decompression threads",
66 )
67 parser.add_argument(
68 "--max-records",
69 type=nonnegative_int,
70 default=0,
71 help="stop after this many records; 0 scans the complete selection",
72 )
73 parser.add_argument(
74 "--output",
75 type=Path,
76 help="write JSON to a new file instead of stdout; refuses overwrite",
77 )
78 return parser
81def require_local_file(path: Path, label: str) -> Path:
82 resolved = path.expanduser()
83 if not resolved.exists():
84 raise FileNotFoundError(f"{label} does not exist: {resolved}")
85 if not resolved.is_file():
86 raise ValueError(f"{label} is not a regular file: {resolved}")
87 return resolved
90def variant_records(
91 variants: pysam.VariantFile,
92 region: Optional[str],
93) -> Iterator[pysam.VariantRecord]:
94 if region is None:
95 return iter(variants)
96 return variants.fetch(region=region)
99def is_symbolic_or_breakend(allele: str) -> bool:
100 return (
101 allele == "*"
102 or allele.startswith("<")
103 or "[" in allele
104 or "]" in allele
105 or (allele != "." and (allele.startswith(".") or allele.endswith(".")))
106 )
109def classify_record(record: pysam.VariantRecord) -> str:
110 alts = record.alts or ()
111 if not alts:
112 return "no_alt"
113 if any(is_symbolic_or_breakend(alt) for alt in alts):
114 return "symbolic_or_breakend"
115 if any(
116 not allele or any(base not in "ACGTN" for base in allele.upper())
117 for allele in (record.ref, *alts)
118 ) or all(alt.upper() == record.ref.upper() for alt in alts):
119 return "other"
120 if len(record.ref) == 1 and all(len(alt) == 1 for alt in alts):
121 return "snv"
122 if all(len(alt) == len(record.ref) for alt in alts):
123 return "mnv"
124 if any(len(alt) != len(record.ref) for alt in alts):
125 return "indel_or_mixed"
126 return "other"
129def filter_status(record: pysam.VariantRecord) -> str:
130 labels = tuple(record.filter.keys())
131 if "PASS" in labels:
132 return "pass"
133 if not labels:
134 return "unfiltered"
135 return "failed"
138def update_substitution_counts(
139 record: pysam.VariantRecord,
140 counts: Counter[str],
141) -> None:
142 alts = record.alts or ()
143 if len(alts) != 1:
144 return
145 ref = record.ref.upper()
146 alt = alts[0].upper()
147 if (
148 ref not in "ACGT" or alt not in "ACGT"
149 or len(ref) != 1 or len(alt) != 1 or ref == alt
150 ):
151 return
153 pair = frozenset((ref, alt))
154 if pair in (frozenset(("A", "G")), frozenset(("C", "T"))):
155 counts["transitions"] += 1
156 else:
157 counts["transversions"] += 1
160def update_genotype_counts(
161 record: pysam.VariantRecord,
162 counts: Counter[str],
163 ploidy_counts: Counter[int],
164) -> None:
165 for call in record.samples.values():
166 genotype = call.get("GT")
167 counts["genotypes_seen"] += 1
168 if genotype:
169 ploidy_counts[len(genotype)] += 1
171 if genotype is None or all(allele is None for allele in genotype):
172 counts["missing_genotypes"] += 1
173 continue
175 if any(allele is None for allele in genotype):
176 counts["partially_missing_genotypes"] += 1
177 continue
179 called = tuple(int(allele) for allele in genotype)
180 counts["called_genotypes"] += 1
181 counts["called_alleles"] += len(called)
182 counts["alternate_alleles"] += sum(
183 allele > 0 for allele in called
184 )
185 if call.phased:
186 counts["phased_called_genotypes"] += 1
188 unique = set(called)
189 if unique == {0}:
190 counts["reference_only_genotypes"] += 1
191 elif len(called) == 1:
192 counts["haploid_alternate_genotypes"] += 1
193 elif len(unique) == 1:
194 counts["same_alternate_allele_genotypes"] += 1
195 else:
196 counts["mixed_allele_genotypes"] += 1
199def summarize_variants(args: argparse.Namespace) -> dict[str, Any]:
200 path = require_local_file(args.input, "input")
201 index = (
202 require_local_file(args.index, "index")
203 if args.index is not None
204 else None
205 )
207 kwargs: dict[str, Any] = {"threads": args.threads}
208 if index is not None:
209 kwargs["index_filename"] = str(index)
211 record_counts: Counter[str] = Counter()
212 filter_label_counts: Counter[str] = Counter()
213 substitution_counts: Counter[str] = Counter()
214 genotype_counts: Counter[str] = Counter()
215 ploidy_counts: Counter[int] = Counter()
216 quality_sum = 0.0
217 quality_count = 0
218 record_limit_reached = False
220 with pysam.VariantFile(str(path), **kwargs) as variants:
221 available_samples = list(variants.header.samples)
222 if args.sample:
223 missing = sorted(set(args.sample) - set(available_samples))
224 if missing:
225 raise ValueError(
226 "samples not present in header: " + ", ".join(missing)
227 )
228 variants.subset_samples(args.sample)
230 decoded_samples = list(variants.header.samples)
232 for record in variant_records(variants, args.region):
233 if args.max_records and record_counts["total"] >= args.max_records:
234 record_limit_reached = True
235 break
237 record_counts["total"] += 1
238 record_counts[classify_record(record)] += 1
239 record_counts[filter_status(record)] += 1
241 alts = record.alts or ()
242 if len(alts) == 1:
243 record_counts["biallelic"] += 1
244 elif len(alts) > 1:
245 record_counts["multiallelic"] += 1
247 if record.qual is None:
248 record_counts["missing_quality"] += 1
249 else:
250 quality_sum += float(record.qual)
251 quality_count += 1
253 labels = tuple(record.filter.keys())
254 if not labels:
255 filter_label_counts["."] += 1
256 else:
257 filter_label_counts.update(labels)
259 update_substitution_counts(record, substitution_counts)
260 update_genotype_counts(
261 record,
262 genotype_counts,
263 ploidy_counts,
264 )
266 try:
267 has_index = variants.index is not None
268 except (AttributeError, OSError, ValueError):
269 has_index = False
271 report: dict[str, Any] = {
272 "schema_version": "1.0",
273 "pysam_version": pysam.__version__,
274 "samtools_version": getattr(
275 pysam, "__samtools_version__", None
276 ),
277 "file_name": path.name,
278 "format": (
279 "BCF"
280 if bool(getattr(variants, "is_bcf", False))
281 else "VCF"
282 ),
283 "has_index": has_index,
284 "selection": {
285 "region": args.region,
286 "region_coordinate_system": (
287 "1-based inclusive samtools string"
288 if args.region is not None
289 else None
290 ),
291 "max_records": args.max_records,
292 "record_limit_reached": record_limit_reached,
293 "available_sample_count": len(available_samples),
294 "decoded_sample_count": len(decoded_samples),
295 },
296 "record_counts": dict(sorted(record_counts.items())),
297 "filter_label_counts": dict(
298 sorted(filter_label_counts.items())
299 ),
300 "substitution_counts": dict(
301 sorted(substitution_counts.items())
302 ),
303 "quality": {
304 "records_with_quality": quality_count,
305 "mean_quality": (
306 quality_sum / quality_count
307 if quality_count
308 else None
309 ),
310 },
311 "genotypes": dict(sorted(genotype_counts.items())),
312 "observed_genotype_ploidy_counts": {
313 str(ploidy): count
314 for ploidy, count in sorted(ploidy_counts.items())
315 },
316 "semantics": (
317 "PASS, unfiltered '.', and failed FILTER states are distinct. "
318 "Record classes are mutually exclusive; indel_or_mixed can "
319 "include multiallelic records with mixed allele lengths. "
320 "SNV/MNV classes describe literal allele lengths, including N. "
321 "Substitution counts use distinct biallelic A/C/G/T alleles. "
322 "Called allele totals exclude partially missing genotypes. "
323 "Observed ploidy counts include missing GT slots, which do not "
324 "establish biological ploidy."
325 ),
326 }
327 if args.include_sample_names:
328 report["selection"]["decoded_samples"] = decoded_samples
329 return report
332def write_report(report: dict[str, Any], output: Optional[Path]) -> None:
333 text = json.dumps(
334 report,
335 indent=2,
336 ensure_ascii=False,
337 allow_nan=False,
338 )
339 if output is None:
340 sys.stdout.write(text + "\n")
341 return
343 destination = output.expanduser()
344 with destination.open("x", encoding="utf-8") as handle:
345 handle.write(text + "\n")
348def main(argv: Optional[list[str]] = None) -> int:
349 parser = build_parser()
350 args = parser.parse_args(argv)
352 if args.output is not None:
353 try:
354 if args.output.expanduser().resolve() == args.input.expanduser().resolve():
355 parser.error("--output must not overwrite the input file")
356 except OSError:
357 pass
359 try:
360 report = summarize_variants(args)
361 write_report(report, args.output)
362 except (
363 FileExistsError,
364 FileNotFoundError,
365 IndexError,
366 KeyError,
367 OSError,
368 TypeError,
369 ValueError,
370 ) as error:
371 print(f"variant_summary: {error}", file=sys.stderr)
372 return 2
373 return 0
376if __name__ == "__main__":
377 raise SystemExit(main())