From 0f40945d58ea230131213208a9a7e0c3de65891d Mon Sep 17 00:00:00 2001 From: Matiq Date: Sun, 6 Sep 2026 18:27:54 +0300 Subject: [PATCH] parallel: auto threads (n-1, capped at 8) - --threads/-t (default auto = detected cores -1, cap 8) - per-chromosome multiprocessing.Pool, FASTA caching - robust nested dict handling for merged profiles --- ffpe_damage_v2.py | 279 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 257 insertions(+), 22 deletions(-) diff --git a/ffpe_damage_v2.py b/ffpe_damage_v2.py index fdaad61..11cf99d 100644 --- a/ffpe_damage_v2.py +++ b/ffpe_damage_v2.py @@ -39,8 +39,9 @@ Outputs (in --outdir): """ import argparse +import multiprocessing import os -from collections import defaultdict +from collections import Counter, defaultdict import pysam import pandas as pd @@ -53,6 +54,200 @@ SUBSTITUTIONS = [f"{r}>{a}" for r in "ACGT" for a in "ACGT" if r != a] END_SIZES = [1, 3, 5, 10, 20] +def get_default_threads(cap=8): + try: + n = len(os.sched_getaffinity(0)) + except AttributeError: + n = os.cpu_count() or 1 + if n is None: + n = 1 + n = max(1, n - 1) + return min(n, cap) + + +def _process_shard(task): + bam_path, ref_path, chrom, mapq, baseq, include_dup = task + bam = pysam.AlignmentFile(bam_path, "rb") + fasta = pysam.FastaFile(ref_path) + try: + chrom_seq = fasta.fetch(chrom) + except Exception: + chrom_seq = None + + sub_counts = Counter() + ref_base_counts = Counter() + profile = {} + profile3 = {} + strand_profile = {} + strand_profile3 = {} + genomic = {} + total_reads = 0 + used_reads = 0 + usable_bases = 0 + + def _inc(store, k1, k2, k3, k4): + a = store.setdefault(k1, {}) + b = a.setdefault(k2, {}) + c = b.setdefault(k3, {}) + c[k4] = c.get(k4, 0) + 1 + + for read in bam.fetch(chrom): + total_reads += 1 + if read.is_secondary or read.is_supplementary: + continue + if read.mapping_quality < mapq: + continue + if read.is_duplicate and not include_dup: + continue + seq = read.query_sequence + quals = read.query_qualities + if seq is None or quals is None: + continue + used_reads += 1 + read_label = "R1" if read.is_read1 else "R2" if read.is_read2 else "single" + strand = "-" if read.is_reverse else "+" + read_len = read.query_length + chrom_name = bam.get_reference_name(read.reference_id) + seq_chrom = chrom_seq if chrom_name == chrom else None + if seq_chrom is None: + try: + seq_chrom = fasta.fetch(chrom_name) + except Exception: + seq_chrom = None + for query_pos, ref_pos in read.get_aligned_pairs(matches_only=True): + if quals[query_pos] < baseq: + continue + if seq_chrom is not None and 0 <= ref_pos < len(seq_chrom): + ref_base = seq_chrom[ref_pos].upper() + else: + try: + ref_base = fasta.fetch(chrom_name, ref_pos, ref_pos + 1).upper() + except Exception: + continue + 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 + key = (chrom_name, ref_pos) + g = genomic.get(key) + if g is None: + g = {"ref": ref_base, "depth": 0, "depth_plus": 0, "depth_minus": 0, "alts": {}} + genomic[key] = g + else: + g["ref"] = ref_base + g["depth"] += 1 + if strand == "+": + g["depth_plus"] += 1 + else: + g["depth_minus"] += 1 + _inc(profile, read_label, pos5, ref_base, "total") + _inc(profile3, read_label, pos3, ref_base, "total") + _inc(strand_profile, strand, pos5, ref_base, "total") + _inc(strand_profile3, strand, pos3, ref_base, "total") + if ref_base == alt_base: + continue + sub = f"{ref_base}>{alt_base}" + sub_counts[sub] += 1 + _inc(profile, read_label, pos5, ref_base, sub) + _inc(profile3, read_label, pos3, ref_base, sub) + _inc(strand_profile, strand, pos5, ref_base, sub) + _inc(strand_profile3, strand, pos3, ref_base, sub) + alts = g["alts"] + alt = alts.get(alt_base) + if alt is None: + alt = {"count": 0, "bq_sum": 0, "pos5_sum": 0, "pos3_sum": 0, "plus": 0, "minus": 0, "r1": 0, "r2": 0} + alts[alt_base] = alt + 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 + bam.close() + fasta.close() + return { + "total_reads": total_reads, + "used_reads": used_reads, + "usable_bases": usable_bases, + "sub_counts": dict(sub_counts), + "ref_base_counts": dict(ref_base_counts), + "profile": profile, + "profile3": profile3, + "strand_profile": strand_profile, + "strand_profile3": strand_profile3, + "genomic": genomic, + } + + +def _merge_nested(target, source): + for k1, v1 in source.items(): + t1 = target.setdefault(k1, {}) + for k2, v2 in v1.items(): + t2 = t1.setdefault(k2, {}) + for k3, v3 in v2.items(): + t3 = t2.setdefault(k3, {}) + for k4, cnt in v3.items(): + t3[k4] = t3.get(k4, 0) + cnt + + +def _merge_results(partials): + merged = { + "total_reads": 0, + "used_reads": 0, + "usable_bases": 0, + "sub_counts": Counter(), + "ref_base_counts": Counter(), + "profile": {}, + "profile3": {}, + "strand_profile": {}, + "strand_profile3": {}, + "genomic": {}, + } + for p in partials: + merged["total_reads"] += p["total_reads"] + merged["used_reads"] += p["used_reads"] + merged["usable_bases"] += p["usable_bases"] + merged["sub_counts"].update(p["sub_counts"]) + merged["ref_base_counts"].update(p["ref_base_counts"]) + _merge_nested(merged["profile"], p["profile"]) + _merge_nested(merged["profile3"], p["profile3"]) + _merge_nested(merged["strand_profile"], p["strand_profile"]) + _merge_nested(merged["strand_profile3"], p["strand_profile3"]) + for key, g in p["genomic"].items(): + mg = merged["genomic"].get(key) + if mg is None: + merged["genomic"][key] = g + else: + mg["depth"] += g["depth"] + mg["depth_plus"] += g["depth_plus"] + mg["depth_minus"] += g["depth_minus"] + for alt_base, alt in g["alts"].items(): + malt = mg["alts"].get(alt_base) + if malt is None: + mg["alts"][alt_base] = alt + else: + malt["count"] += alt["count"] + malt["bq_sum"] += alt["bq_sum"] + malt["pos5_sum"] += alt["pos5_sum"] + malt["pos3_sum"] += alt["pos3_sum"] + malt["plus"] += alt["plus"] + malt["minus"] += alt["minus"] + malt["r1"] += alt["r1"] + malt["r2"] += alt["r2"] + return merged + + def parse_args(): ap = argparse.ArgumentParser(description="FFPE damage analysis") ap.add_argument("--bam", required=True) @@ -74,6 +269,9 @@ def parse_args(): "counts as end-associated (default 10)") ap.add_argument("--filter-strand-p", type=float, default=0.05, help="Fisher p-value threshold for strand bias (default 0.05)") + ap.add_argument("--threads", "-t", type=int, default=None, + help="threads for BAM processing (default auto = " + "detected cores -1, capped at 8; 1 = single-thread)") return ap.parse_args() @@ -229,14 +427,14 @@ def write_read_position_profiles(d, args): for label in store: for pos in sorted(store[label]): for ref in "ACGT": - total = store[label][pos][ref]["total"] + total = _get_nested(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] + count = _get_nested(store, label, pos, ref, sub) out.append({ "read": label, pos_key: pos, @@ -261,14 +459,18 @@ def write_read_position_profiles(d, args): return df5, df3 +def _get_nested(store, label, pos, ref, field): + return store.get(label, {}).get(pos, {}).get(ref, {}).get(field, 0) + + 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"] + c_total = _get_nested(store, label, pos, "C", "total") + ct = _get_nested(store, label, pos, "C", "C>T") + g_total = _get_nested(store, label, pos, "G", "total") + ga = _get_nested(store, label, pos, "G", "G>A") out.append({ "read": label, pos_key: pos, @@ -310,10 +512,10 @@ def write_end_enrichment(d, args): 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"] + ct_den += _get_nested(store, label, pos, "C", "total") + ct_num += _get_nested(store, label, pos, "C", "C>T") + ga_den += _get_nested(store, label, pos, "G", "total") + ga_num += _get_nested(store, label, pos, "G", "G>A") rows.append({ "read": label, "end": side, @@ -592,27 +794,31 @@ def main(): args = parse_args() os.makedirs(args.outdir, exist_ok=True) - bam = pysam.AlignmentFile(args.bam, "rb") - if not bam.has_index(): + threads = args.threads + if threads is None: + threads = get_default_threads() + if threads < 1: + threads = 1 + + bam_probe = pysam.AlignmentFile(args.bam, "rb") + has_index = bam_probe.has_index() + bam_probe.close() + if not has_index: try: print(f"[info] BAM index not found, creating {args.bam}.bai ...") pysam.index(args.bam) - bam.close() - bam = pysam.AlignmentFile(args.bam, "rb") except Exception as e: - bam.close() raise RuntimeError( f"BAM index (.bai) not found and auto-creation failed: {e}. " f"Run: samtools index {args.bam}") from e try: - fasta = pysam.FastaFile(args.reference) + fasta_probe = pysam.FastaFile(args.reference) + fasta_probe.close() except Exception as e: - bam.close() if "fai" in str(e).lower() or "index" in str(e).lower(): try: print(f"[info] FASTA index not found, creating {args.reference}.fai ...") pysam.faidx(args.reference) - fasta = pysam.FastaFile(args.reference) except Exception as e2: raise RuntimeError( f"FASTA index (.fai) not found and auto-creation failed: {e2}. " @@ -620,9 +826,38 @@ def main(): else: raise - d = process_bam(bam, fasta, args) - bam.close() - fasta.close() + if threads == 1: + bam = pysam.AlignmentFile(args.bam, "rb") + fasta = pysam.FastaFile(args.reference) + d = process_bam(bam, fasta, args) + bam.close() + fasta.close() + else: + print(f"[info] using {threads} threads (auto, capped at 8, 1 reserved)") + bam_tmp = pysam.AlignmentFile(args.bam, "rb") + chroms = list(bam_tmp.references) + bam_tmp.close() + if not chroms: + bam = pysam.AlignmentFile(args.bam, "rb") + fasta = pysam.FastaFile(args.reference) + d = process_bam(bam, fasta, args) + bam.close() + fasta.close() + else: + tasks = [(args.bam, args.reference, c, args.mapq, args.baseq, + args.include_duplicates) for c in chroms] + with multiprocessing.Pool(threads) as pool: + partials = pool.map(_process_shard, tasks) + d = _merge_results(partials) + try: + bam_cnt = pysam.AlignmentFile(args.bam, "rb") + total_via_idx = sum(s.mapped for s in bam_cnt.get_index_statistics()) + total_via_idx += bam_cnt.unmapped + bam_cnt.close() + if total_via_idx > d["total_reads"]: + d["total_reads"] = total_via_idx + except Exception: + pass sub_df = write_substitution_summary(d, args) write_read_position_profiles(d, args)