ffpe_damage/ffpe_damage_v2.py

#!/usr/bin/env python3
"""FFPE deamination profiling from an aligned BAM.

Methodology
-----------
For every read position i the normalized frequency is:

    C>T_i = #(C->T at position i) / #(C observations at position i)
    G>A_i = #(G->A at position i) / #(G observations at position i)

where the denominator counts ALL high-quality observations of that
reference base at position i, including matches.

Separate profiles are produced for R1 / R2 and for forward / reverse
alignment strand, using both distance-from-5'-end and distance-from-3'-end.

All 12 substitution types are counted globally so that C>T / G>A can be
judged against the other mismatch classes.

This tool reports measurable frequencies only; it does NOT assign a
low/moderate/high classification.  Classification requires comparison
against validated control samples (see ffpe_compare.py).

Outputs (in --outdir):
  sample_summary.csv
  substitution_summary.csv
  substitution_all12.png
  normalized_damage_by_read_position.csv          (5' profile)
  normalized_damage_by_read_position_3prime.csv   (3' profile)
  CtoT_GtoA_normalized_profile.csv                (5' profile)
  CtoT_GtoA_normalized_profile_3prime.csv         (3' profile)
  strand_damage_profile.csv                       (5' profile)
  strand_damage_profile_3prime.csv                (3' profile)
  end_enrichment.csv
  candidate_CtoT.bed / candidate_GtoA.bed         (IGV tracks)
  candidate_variants.tsv                          (per-position context)
  CtoT_R1.png CtoT_R2.png GtoA_R1.png GtoA_R2.png (5' plots)
  CtoT_R1_3prime.png ... GtoA_R2_3prime.png       (3' plots)
"""

import argparse
import os
from collections import defaultdict

import pysam
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.stats import fisher_exact

SUBSTITUTIONS = [f"{r}>{a}" for r in "ACGT" for a in "ACGT" if r != a]
END_SIZES = [1, 3, 5, 10, 20]


def parse_args():
    ap = argparse.ArgumentParser(description="FFPE damage analysis")
    ap.add_argument("--bam", required=True)
    ap.add_argument("--reference", required=True)
    ap.add_argument("--outdir", default="FFPE_QC")
    ap.add_argument("--mapq", type=int, default=20)
    ap.add_argument("--baseq", type=int, default=20)
    ap.add_argument("--min-depth", type=int, default=20)
    ap.add_argument("--min-alt-count", type=int, default=5)
    ap.add_argument("--include-duplicates", action="store_true")
    return ap.parse_args()


def process_bam(bam, fasta, args):
    total_reads = 0
    used_reads = 0
    usable_bases = 0

    sub_counts = defaultdict(int)
    ref_base_counts = defaultdict(int)

    profile = defaultdict(lambda: defaultdict(lambda: defaultdict(
        lambda: defaultdict(int))))
    profile3 = defaultdict(lambda: defaultdict(lambda: defaultdict(
        lambda: defaultdict(int))))
    strand_profile = defaultdict(lambda: defaultdict(lambda: defaultdict(
        lambda: defaultdict(int))))
    strand_profile3 = defaultdict(lambda: defaultdict(lambda: defaultdict(
        lambda: defaultdict(int))))

    genomic = defaultdict(lambda: {
        "ref": None,
        "depth": 0,
        "depth_plus": 0,
        "depth_minus": 0,
        "alts": defaultdict(lambda: {
            "count": 0,
            "bq_sum": 0,
            "pos5_sum": 0,
            "pos3_sum": 0,
            "plus": 0,
            "minus": 0,
            "r1": 0,
            "r2": 0,
        }),
    })

    for read in bam.fetch(until_eof=True):
        total_reads += 1
        if read.is_unmapped or read.is_secondary or read.is_supplementary:
            continue
        if read.mapping_quality < args.mapq:
            continue
        if read.is_duplicate and not args.include_duplicates:
            continue
        seq = read.query_sequence
        quals = read.query_qualities
        if seq is None or quals is None:
            continue
        used_reads += 1

        if read.is_read1:
            read_label = "R1"
        elif read.is_read2:
            read_label = "R2"
        else:
            read_label = "single"
        strand = "-" if read.is_reverse else "+"
        read_len = read.query_length
        chrom = bam.get_reference_name(read.reference_id)

        for query_pos, ref_pos in read.get_aligned_pairs(matches_only=True):
            if quals[query_pos] < args.baseq:
                continue
            ref_base = fasta.fetch(chrom, ref_pos, ref_pos + 1).upper()
            if ref_base not in "ACGT":
                continue
            alt_base = seq[query_pos].upper()
            if alt_base not in "ACGT":
                continue
            usable_bases += 1

            pos5 = query_pos + 1
            pos3 = read_len - query_pos

            ref_base_counts[ref_base] += 1

            g = genomic[(chrom, ref_pos)]
            g["ref"] = ref_base
            g["depth"] += 1
            if strand == "+":
                g["depth_plus"] += 1
            else:
                g["depth_minus"] += 1

            profile[read_label][pos5][ref_base]["total"] += 1
            profile3[read_label][pos3][ref_base]["total"] += 1
            strand_profile[strand][pos5][ref_base]["total"] += 1
            strand_profile3[strand][pos3][ref_base]["total"] += 1

            if ref_base == alt_base:
                continue

            sub = f"{ref_base}>{alt_base}"
            sub_counts[sub] += 1

            profile[read_label][pos5][ref_base][sub] += 1
            profile3[read_label][pos3][ref_base][sub] += 1
            strand_profile[strand][pos5][ref_base][sub] += 1
            strand_profile3[strand][pos3][ref_base][sub] += 1

            alt = g["alts"][alt_base]
            alt["count"] += 1
            alt["bq_sum"] += quals[query_pos]
            alt["pos5_sum"] += pos5
            alt["pos3_sum"] += pos3
            if strand == "+":
                alt["plus"] += 1
            else:
                alt["minus"] += 1
            if read_label == "R1":
                alt["r1"] += 1
            elif read_label == "R2":
                alt["r2"] += 1

    return {
        "total_reads": total_reads,
        "used_reads": used_reads,
        "usable_bases": usable_bases,
        "sub_counts": sub_counts,
        "ref_base_counts": ref_base_counts,
        "profile": profile,
        "profile3": profile3,
        "strand_profile": strand_profile,
        "strand_profile3": strand_profile3,
        "genomic": genomic,
    }


def write_substitution_summary(d, args):
    rows = []
    for sub in SUBSTITUTIONS:
        ref = sub[0]
        count = d["sub_counts"][sub]
        denom = d["ref_base_counts"][ref]
        freq = count / denom if denom else 0.0
        rows.append({
            "substitution": sub,
            "count": count,
            "reference_base_observations": denom,
            "frequency": freq,
            "frequency_percent": freq * 100,
        })
    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(args.outdir, "substitution_summary.csv"),
              index=False)
    return df


def write_read_position_profiles(d, args):
    def rows_for(store, pos_key, side):
        out = []
        for label in store:
            for pos in sorted(store[label]):
                for ref in "ACGT":
                    total = store[label][pos][ref]["total"]
                    if total == 0:
                        continue
                    for alt in "ACGT":
                        if alt == ref:
                            continue
                        sub = f"{ref}>{alt}"
                        count = store[label][pos][ref][sub]
                        out.append({
                            "read": label,
                            pos_key: pos,
                            "end_side": side,
                            "reference_base": ref,
                            "substitution": sub,
                            "count": count,
                            "denominator": total,
                            "frequency": count / total,
                            "frequency_percent": count / total * 100,
                        })
        return out

    df5 = pd.DataFrame(rows_for(d["profile"], "position_5prime", "5prime"))
    df3 = pd.DataFrame(rows_for(d["profile3"], "position_3prime", "3prime"))
    df5.to_csv(os.path.join(args.outdir,
                            "normalized_damage_by_read_position.csv"),
               index=False)
    df3.to_csv(os.path.join(args.outdir,
                            "normalized_damage_by_read_position_3prime.csv"),
               index=False)
    return df5, df3


def cg_rows(store, pos_key):
    out = []
    for label in store:
        for pos in sorted(store[label]):
            c_total = store[label][pos]["C"]["total"]
            ct = store[label][pos]["C"]["C>T"]
            g_total = store[label][pos]["G"]["total"]
            ga = store[label][pos]["G"]["G>A"]
            out.append({
                "read": label,
                pos_key: pos,
                "C_total": c_total,
                "C_to_T": ct,
                "C_to_T_frequency_percent": ct / c_total * 100 if c_total else 0,
                "G_total": g_total,
                "G_to_A": ga,
                "G_to_A_frequency_percent": ga / g_total * 100 if g_total else 0,
            })
    return out


def write_cg_profiles(d, args):
    df5 = pd.DataFrame(cg_rows(d["profile"], "position_5prime"))
    df3 = pd.DataFrame(cg_rows(d["profile3"], "position_3prime"))
    df5.to_csv(os.path.join(args.outdir,
                            "CtoT_GtoA_normalized_profile.csv"), index=False)
    df3.to_csv(os.path.join(args.outdir,
                            "CtoT_GtoA_normalized_profile_3prime.csv"),
               index=False)
    return df5, df3


def write_strand_profiles(d, args):
    df5 = pd.DataFrame(cg_rows(d["strand_profile"], "position_5prime"))
    df3 = pd.DataFrame(cg_rows(d["strand_profile3"], "position_3prime"))
    df5.to_csv(os.path.join(args.outdir, "strand_damage_profile.csv"),
               index=False)
    df3.to_csv(os.path.join(args.outdir,
                            "strand_damage_profile_3prime.csv"), index=False)
    return df5, df3


def write_end_enrichment(d, args):
    rows = []
    for store, side in ((d["profile"], "5prime"), (d["profile3"], "3prime")):
        for label in store:
            for size in END_SIZES:
                ct_num = ct_den = ga_num = ga_den = 0
                for pos in range(1, size + 1):
                    ct_den += store[label][pos]["C"]["total"]
                    ct_num += store[label][pos]["C"]["C>T"]
                    ga_den += store[label][pos]["G"]["total"]
                    ga_num += store[label][pos]["G"]["G>A"]
                rows.append({
                    "read": label,
                    "end": side,
                    "window_size": size,
                    "C_to_T_frequency_percent":
                        ct_num / ct_den * 100 if ct_den else 0,
                    "G_to_A_frequency_percent":
                        ga_num / ga_den * 100 if ga_den else 0,
                })
    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(args.outdir, "end_enrichment.csv"), index=False)
    return df


def write_beds_and_variants(d, args):
    ct_lines = []
    ga_lines = []
    var_rows = []

    for (chrom, pos), data in d["genomic"].items():
        depth = data["depth"]
        if depth < args.min_depth:
            continue
        ref = data["ref"]
        for alt_base, alt in data["alts"].items():
            count = alt["count"]
            if count < args.min_alt_count:
                continue
            vaf = count / depth
            name = f"{ref}>{alt_base};depth={depth};alt={count};VAF={vaf:.3f}"
            start, end = pos, pos + 1
            line = f"{chrom}\t{start}\t{end}\t{name}\n"
            if ref == "C" and alt_base == "T":
                ct_lines.append(line)
            if ref == "G" and alt_base == "A":
                ga_lines.append(line)

            plus = alt["plus"]
            minus = alt["minus"]
            depth_plus = data["depth_plus"]
            depth_minus = data["depth_minus"]
            ref_plus = depth_plus - plus
            ref_minus = depth_minus - minus
            table = [[plus, ref_plus], [minus, ref_minus]]
            fisher_p = 1.0
            try:
                if plus + ref_plus > 0 and minus + ref_minus > 0:
                    _, fisher_p = fisher_exact(table)
            except Exception:
                fisher_p = 1.0

            alt_count = alt["count"]
            mean_pos5 = alt["pos5_sum"] / alt_count
            mean_pos3 = alt["pos3_sum"] / alt_count

            var_rows.append({
                "chrom": chrom,
                "position": pos,
                "ref": ref,
                "alt": alt_base,
                "depth": depth,
                "alt_count": count,
                "VAF": vaf,
                "mean_base_quality": alt["bq_sum"] / alt_count,
                "mean_read_position_5prime": mean_pos5,
                "mean_read_position_3prime": mean_pos3,
                "alt_fraction_plus": plus / alt_count,
                "strand_bias_pvalue_fisher": fisher_p,
                "alt_fraction_R1": alt["r1"] / alt_count,
                "alt_fraction_R2": alt["r2"] / alt_count,
            })

    with open(os.path.join(args.outdir, "candidate_CtoT.bed"), "w") as f:
        f.writelines(ct_lines)
    with open(os.path.join(args.outdir, "candidate_GtoA.bed"), "w") as f:
        f.writelines(ga_lines)
    if var_rows:
        pd.DataFrame(var_rows).to_csv(
            os.path.join(args.outdir, "candidate_variants.tsv"),
            index=False, sep="\t")


def write_summary(d, args):
    ct = d["sub_counts"]["C>T"]
    ga = d["sub_counts"]["G>A"]
    ct_den = d["ref_base_counts"]["C"]
    ga_den = d["ref_base_counts"]["G"]
    row = {
        "total_reads": d["total_reads"],
        "used_reads": d["used_reads"],
        "usable_bases": d["usable_bases"],
        "C_to_T_count": ct,
        "C_observations": ct_den,
        "C_to_T_frequency_percent": ct / ct_den * 100 if ct_den else 0,
        "G_to_A_count": ga,
        "G_observations": ga_den,
        "G_to_A_frequency_percent": ga / ga_den * 100 if ga_den else 0,
        "MAPQ_threshold": args.mapq,
        "BQ_threshold": args.baseq,
        "duplicates_included": args.include_duplicates,
    }
    df = pd.DataFrame([row])
    df.to_csv(os.path.join(args.outdir, "sample_summary.csv"), index=False)
    return df


def plot_all12(sub_df, outdir):
    x = sub_df["substitution"]
    y = sub_df["frequency_percent"]
    plt.figure(figsize=(10, 5))
    bars = plt.bar(x, y)
    for b, s in zip(bars, sub_df["substitution"]):
        if s in ("C>T", "G>A"):
            b.set_color("#d62728")
        else:
            b.set_color("#7f7f7f")
    plt.ylabel("frequency (%)")
    plt.title("All 12 substitution types, global frequency")
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig(os.path.join(outdir, "substitution_all12.png"), dpi=300)
    plt.close()


def plot_cg_profiles(cg5, cg3, outdir):
    for label in sorted(set(cg5["read"]) | set(cg3["read"])):
        for sub, col5, col3, prefix in (
                ("C>T", "C_to_T_frequency_percent",
                 "C_to_T_frequency_percent", "CtoT"),
                ("G>A", "G_to_A_frequency_percent",
                 "G_to_A_frequency_percent", "GtoA")):
            for side, df in (("5prime", cg5), ("3prime", cg3)):
                d = df[df["read"] == label]
                if d.empty:
                    continue
                pos_col = ("position_5prime" if side == "5prime"
                           else "position_3prime")
                col = col5 if side == "5prime" else col3
                plt.figure(figsize=(8, 5))
                plt.plot(d[pos_col], d[col], marker=".", label=f"{label} {sub}")
                plt.xlabel(f"position in read ({side})")
                plt.ylabel(f"{sub} frequency (%)")
                plt.title(f"FFPE damage profile — {label} — {sub} — {side}")
                plt.legend()
                plt.tight_layout()
                suffix = "" if side == "5prime" else "_3prime"
                plt.savefig(os.path.join(
                    outdir, f"{prefix}_{label}{suffix}.png"), dpi=300)
                plt.close()


def main():
    args = parse_args()
    os.makedirs(args.outdir, exist_ok=True)

    bam = pysam.AlignmentFile(args.bam, "rb")
    if not bam.has_index():
        bam.close()
        raise RuntimeError("BAM index (.bai) not found.")
    fasta = pysam.FastaFile(args.reference)

    d = process_bam(bam, fasta, args)
    bam.close()
    fasta.close()

    sub_df = write_substitution_summary(d, args)
    write_read_position_profiles(d, args)
    cg5, cg3 = write_cg_profiles(d, args)
    write_strand_profiles(d, args)
    write_end_enrichment(d, args)
    write_beds_and_variants(d, args)
    write_summary(d, args)
    plot_all12(sub_df, args.outdir)
    plot_cg_profiles(cg5, cg3, args.outdir)

    print(f"total reads     : {d['total_reads']}")
    print(f"used reads      : {d['used_reads']}")
    print(f"usable bases    : {d['usable_bases']}")
    print(f"C>T global      : "
          f"{d['sub_counts']['C>T']} / {d['ref_base_counts']['C']} = "
          f"{d['sub_counts']['C>T'] / d['ref_base_counts']['C'] * 100:.3f}%")
    print(f"G>A global      : "
          f"{d['sub_counts']['G>A']} / {d['ref_base_counts']['G']} = "
          f"{d['sub_counts']['G>A'] / d['ref_base_counts']['G'] * 100:.3f}%")
    print(f"outputs         : {args.outdir}/")


if __name__ == "__main__":
    main()