Pysam

python scripts/inspect_hts.py sample.bam
python scripts/inspect_hts.py cohort.vcf.gz
python scripts/inspect_hts.py reference.fa

python scripts/alignment_qc.py sample.bam --output sample.qc.json
python scripts/alignment_qc.py sample.cram \
  --reference reference.fa \
  --max-records 100000

import pysam

with pysam.AlignmentFile("sample.bam", "rb", require_index=True) as bam:
    for stats in bam.get_index_statistics():
        print(stats.contig, stats.mapped, stats.unmapped, stats.total)

import pysam


def base_depth(
    bam: pysam.AlignmentFile,
    contig: str,
    start: int,
    stop: int,
    *,
    min_base_quality: int = 20,
) -> list[int]:
    if not 0 <= start < stop <= bam.get_reference_length(contig):
        raise ValueError("coverage interval must be nonempty and within the contig")
    a, c, g, t = bam.count_coverage(
        contig,
        start,
        stop,
        quality_threshold=min_base_quality,
        read_callback="all",
    )
    return [sum(counts) for counts in zip(a, c, g, t)]


with pysam.AlignmentFile("sample.bam", "rb") as bam:
    depth = base_depth(bam, "chr1", 1_000, 2_000)

def usable_read(read: pysam.AlignedSegment) -> bool:
    return (
        not read.is_unmapped
        and not read.is_secondary
        and not read.is_supplementary
        and not read.is_qcfail
        and not read.is_duplicate
        and read.mapping_quality >= 20
    )


a, c, g, t = bam.count_coverage(
    "chr1",
    1_000,
    2_000,
    quality_threshold=20,
    read_callback=usable_read,
)

def below_threshold_intervals(
    depths: list[int],
    start: int,
    threshold: int,
):
    interval_start = None

    for offset, value in enumerate(depths):
        position = start + offset
        if value < threshold and interval_start is None:
            interval_start = position
        elif value >= threshold and interval_start is not None:
            yield interval_start, position
            interval_start = None

    if interval_start is not None:
        yield interval_start, start + len(depths)

def aligned_depth_at(
    bam: pysam.AlignmentFile,
    fasta: pysam.FastaFile,
    contig: str,
    position: int,
) -> int:
    iterator = bam.pileup(
        contig,
        position,
        position + 1,
        truncate=True,
        stepper="samtools",
        fastafile=fasta,
        min_mapping_quality=20,
        min_base_quality=20,
        max_depth=100_000,
        ignore_overlaps=True,
        ignore_orphans=True,
    )
    for column in iterator:
        if column.reference_pos == position:
            return sum(
                not item.is_del and not item.is_refskip
                and item.query_position is not None
                for item in column.pileups
            )
    return 0

from collections import Counter


def snp_base_counts(
    bam: pysam.AlignmentFile,
    fasta: pysam.FastaFile,
    record: pysam.VariantRecord,
) -> Counter[str]:
    if (
        record.ref.upper() not in {"A", "C", "G", "T"}
        or not record.alts
        or any(alt.upper() not in {"A", "C", "G", "T"} for alt in record.alts)
        or any(alt.upper() == record.ref.upper() for alt in record.alts)
    ):
        raise ValueError("snp_base_counts only supports simple SNP records")

    # Validate REF without inferring a variant from a mismatched assembly.
    observed_ref = fasta.fetch(record.contig, record.start, record.start + 1)
    if observed_ref.upper() != record.ref.upper():
        raise ValueError("VCF REF does not match the selected FASTA")
    counts: Counter[str] = Counter()
    iterator = bam.pileup(
        record.contig,
        record.start,
        record.start + 1,
        truncate=True,
        stepper="samtools",
        fastafile=fasta,
        min_mapping_quality=20,
        min_base_quality=20,
        max_depth=100_000,
        ignore_overlaps=True,
        ignore_orphans=True,
    )

    for column in iterator:
        for pileup_read in column.pileups:
            if pileup_read.is_del or pileup_read.is_refskip:
                continue
            query_position = pileup_read.query_position
            if query_position is None:
                continue
            base = pileup_read.alignment.query_sequence[query_position]
            counts[base.upper()] += 1
    return counts

import pysam


def annotate_depth(
    input_vcf: str,
    input_bam: str,
    reference_fasta: str,
    output_vcf: str,
) -> None:
    from pathlib import Path

    if Path(output_vcf).exists():
        raise FileExistsError(output_vcf)
    with pysam.FastaFile(reference_fasta) as fasta, pysam.AlignmentFile(
        input_bam,
        "rb",
        reference_filename=reference_fasta,
    ) as bam, pysam.VariantFile(input_vcf) as source:
        header = source.header.copy()
        if "BAM_BASE_DP" in header.info:
            raise ValueError("BAM_BASE_DP already exists in input header")
        header.info.add(
            "BAM_BASE_DP",
            number=1,
            type="Integer",
            description=(
                "Aligned base depth at POS; MAPQ>=20, baseQ>=20, "
                "samtools pileup filters, paired overlaps collapsed"
            ),
        )

        with pysam.VariantFile(output_vcf, "w", header=header) as output:
            for source_record in source:
                record = source_record.copy()
                record.translate(header)
                record.info["BAM_BASE_DP"] = aligned_depth_at(
                    bam,
                    fasta,
                    record.contig,
                    record.start,
                )
                output.write(record)

def reference_matches(
    record: pysam.VariantRecord,
    fasta: pysam.FastaFile,
) -> bool:
    observed = fasta.fetch(
        record.contig,
        record.start,
        record.start + len(record.ref),
    )
    return observed.upper() == record.ref.upper()


with pysam.FastaFile("reference.fa") as fasta, pysam.VariantFile(
    "variants.vcf.gz"
) as variants:
    mismatches = [
        (record.contig, record.pos, record.ref)
        for record in variants
        if not reference_matches(record, fasta)
    ]

python scripts/alignment_qc.py input.bam --max-records 10000
python scripts/filter_alignments.py input.bam filtered.bam \
  --min-mapq 20 --exclude-unmapped --exclude-secondary \
  --exclude-qcfail --exclude-duplicates --exclude-supplementary --index

with pysam.AlignmentFile("input.bam", "rb") as source, pysam.AlignmentFile(
    "filtered.bam", "wb", template=source
) as output:
    for read in source.fetch(until_eof=True):
        if usable_read(read):
            output.write(read)

with pysam.TabixFile(
    "genes.bed.gz", parser=pysam.asBed()
) as genes, pysam.FastaFile("reference.fa") as fasta, open(
    "genes.fa", "x", encoding="utf-8"
) as output:
    for gene in genes.fetch():
        sequence = fasta.fetch(gene.contig, gene.start, gene.end)
        if gene.strand == "-":
            sequence = pysam.reverse_complement(sequence)
        output.write(f">{gene.name}\n{sequence}\n")

with pysam.AlignmentFile("rna.bam", "rb") as bam:
    primary_reads = (
        read
        for read in bam.fetch("chr1")
        if not read.is_unmapped
        and not read.is_qcfail
        and not read.is_secondary
        and not read.is_supplementary
        and not read.is_duplicate
        and read.mapping_quality >= 20
    )
    junction_counts = bam.find_introns(primary_reads)

import pysam.samtools
import pysam.bcftools

pysam.samtools.sort(
    "-@", "4",
    "-o", "sorted.bam",
    "input.bam",
    catch_stdout=False,
)
pysam.samtools.index(
    "-@", "4",
    "sorted.bam",
    catch_stdout=False,
)

pysam.bcftools.norm(
    "-f", "reference.fa",
    "-m", "-any",
    "--multi-overlaps", ".",  # other ALT alleles become missing after splitting
    "-Oz",
    "-o", "normalized.vcf.gz",
    "input.vcf.gz",
    catch_stdout=False,
)
pysam.bcftools.index(
    "--csi",
    "normalized.vcf.gz",
    catch_stdout=False,
)

pysam.samtools.quickcheck("-v", "filtered.bam")

with pysam.VariantFile("annotated.vcf.gz") as variants:
    assert "BAM_BASE_DP" in variants.header.info