FilterMafftAlignment.py

#!/usr/bin/env python3
"""
filter_mafft_alignment.py

对 MAFFT 生成的比对文件按位点缺失/模糊比例过滤（多阈值支持），内存友好实现。

用法示例：
  python filter_mafft_alignment.py input_aln.fasta outprefix \
      --thresholds 10 30 50 \
      --gap-chars "Nn-.?RYSMKW" \
      --mode "gt" \
      --format fasta

说明:
 - thresholds: 百分比阈值（整数或浮点），例如 10 30 50
 - gap-chars: 被视为缺失/模糊的字符集合（默认 "Nn-.?"）
 - mode: "gt" 表示严格大于 (> x%) 才删除；"ge" 表示大于等于 (>= x%) 才删除
 - 输出: 为每个阈值创建:
     <outprefix>_filtered_{t}pct.fasta
     <outprefix>_removed_{t}pct.txt   (移除位点，1-based，按区间压缩)
"""

import argparse
import sys
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from typing import List, Tuple
import math

def parse_args():
    p = argparse.ArgumentParser(description="Filter alignment columns by gap/ambiguous fraction (memory-friendly).")
    p.add_argument("infile", help="Input alignment (MAFFT) file (FASTA or other Biopython-supported format).")
    p.add_argument("outprefix", help="Output prefix for filtered files.")
    p.add_argument("--thresholds", nargs="+", type=float, default=[10,30,50],
                   help="Threshold percentages. Example: --thresholds 10 30 50")
    p.add_argument("--gap-chars", type=str, default="Nn-.?",
                   help='Characters counted as gap/ambiguous (default "Nn-.?"). Add IUPAC letters here if desired.')
    p.add_argument("--mode", choices=("gt","ge"), default="gt",
                   help='"gt": delete when fraction > threshold (default). "ge": delete when fraction >= threshold.')
    p.add_argument("--format", type=str, default="fasta",
                   help="Input file format for Biopython SeqIO (default fasta).")
    p.add_argument("--line-width", type=int, default=80,
                   help="Wrap sequence lines to this width in output FASTA (default 80). Use 0 to disable wrapping.")
    p.add_argument("--show-progress", action="store_true", help="Print progress summary lines during processing.")
    return p.parse_args()

def compute_gap_counts(infile: str, fmt: str, gap_chars: set, show_progress=False) -> Tuple[List[int], int]:
    """
    First pass: compute per-column gap counts and return (counts_list, nseq).
    Ensures all sequences have the same length; raises ValueError otherwise.
    """
    nseq = 0
    counts = None
    seq_len = None
    with open(infile, 'r') as fh:
        for rec in SeqIO.parse(fh, fmt):
            s = str(rec.seq)
            if seq_len is None:
                seq_len = len(s)
                counts = [0] * seq_len
            else:
                if len(s) != seq_len:
                    raise ValueError(f"Sequence length mismatch: record {rec.id} length {len(s)} != expected {seq_len}.")
            # increment counts
            for i, ch in enumerate(s):
                if ch in gap_chars:
                    counts[i] += 1
            nseq += 1
            if show_progress and (nseq % 200 == 0):
                print(f"[counting] processed {nseq} sequences...", file=sys.stderr)
    if nseq == 0:
        raise ValueError("No sequences found in input.")
    if show_progress:
        print(f"[counting] done. {nseq} sequences, {seq_len} columns.", file=sys.stderr)
    return counts, nseq

def decide_columns(counts: List[int], nseq: int, threshold: float, mode: str) -> Tuple[List[int], List[int]]:
    """
    Given per-column counts and threshold percentage, return (keep_indices, removed_indices),
    indices are 0-based for internal use; removed_indices returned also.
    """
    keep = []
    removed = []
    denom = nseq
    cutoff = threshold / 100.0
    for i, c in enumerate(counts):
        frac = c / denom
        if mode == "gt":
            if frac > cutoff:
                removed.append(i)
            else:
                keep.append(i)
        else:  # ge
            if frac >= cutoff:
                removed.append(i)
            else:
                keep.append(i)
    return keep, removed

def write_filtered_outputs(infile: str, fmt: str, outprefix: str, thresholds_info: dict,
                           wrap_width: int = 80, show_progress=False):
    """
    Second pass: stream sequences and for each threshold write filtered FASTA.
    thresholds_info: dict threshold->(keep_indices, removed_indices)
    """
    # Prepare output file handles and SeqIO writers by threshold
    writers = {}
    outfiles = {}
    for t in thresholds_info:
        outf = f"{outprefix}_filtered_{int(t)}pct.fasta"
        outfiles[t] = outf
        writers[t] = open(outf, 'w')

    # Iterate input sequences and write filtered variants
    total = 0
    with open(infile, 'r') as fh:
        for rec in SeqIO.parse(fh, fmt):
            s = str(rec.seq)
            for t, (keep_idx, rem_idx) in thresholds_info.items():
                if len(keep_idx) == 0:
                    new_seq_str = ""
                elif len(keep_idx) == len(s):  # keep all
                    new_seq_str = s
                else:
                    # build new seq via join on indices (this is memory per sequence only)
                    # faster approach: list comprehension
                    new_seq_str = ''.join(s[i] for i in keep_idx)
                # write fasta entry with wrapping
                writers[t].write(f">{rec.id}\n")
                if wrap_width and wrap_width > 0:
                    for i in range(0, len(new_seq_str), wrap_width):
                        writers[t].write(new_seq_str[i:i+wrap_width] + "\n")
                else:
                    writers[t].write(new_seq_str + "\n")
            total += 1
            if show_progress and (total % 200 == 0):
                print(f"[writing] processed {total} sequences...", file=sys.stderr)
    # close writers
    for fh in writers.values():
        fh.close()
    if show_progress:
        print(f"[writing] done. processed {total} sequences. Files written: {', '.join(outfiles.values())}", file=sys.stderr)

def write_removed_positions(outprefix: str, threshold: float, removed_indices: List[int]):
    """
    Write removed positions as 1-based indices, compressed into ranges for readability.
    """
    def compress(ints: List[int]) -> str:
        if not ints:
            return ""
        ints_sorted = sorted(i+1 for i in ints)  # to 1-based
        ranges = []
        start = prev = ints_sorted[0]
        for v in ints_sorted[1:]:
            if v == prev + 1:
                prev = v
            else:
                if start == prev:
                    ranges.append(str(start))
                else:
                    ranges.append(f"{start}-{prev}")
                start = prev = v
        # finish
        if start == prev:
            ranges.append(str(start))
        else:
            ranges.append(f"{start}-{prev}")
        return ','.join(ranges)

    outname = f"{outprefix}_removed_{int(threshold)}pct.txt"
    with open(outname, 'w') as fh:
        fh.write("# removed positions (1-based). Ranges like start-end or single positions, comma-separated.\n")
        fh.write(compress(removed_indices) + "\n")

def main():
    args = parse_args()
    gap_chars = set(args.gap_chars)
    try:
        counts, nseq = compute_gap_counts(args.infile, args.format, gap_chars, show_progress=args.show_progress)
    except Exception as e:
        print("Error during counting:", e, file=sys.stderr)
        sys.exit(1)

    # decide columns for each threshold
    thresholds_info = {}
    for t in args.thresholds:
        keep, removed = decide_columns(counts, nseq, t, args.mode)
        thresholds_info[t] = (keep, removed)
        print(f"[decide] threshold {t}% => kept {len(keep)} cols, removed {len(removed)} cols")

    # write removed position files
    for t, (_, removed) in thresholds_info.items():
        write_removed_positions(args.outprefix, t, removed)
        print(f"[output] removed list written: {args.outprefix}_removed_{int(t)}pct.txt")

    # write filtered FASTA files in one pass
    write_filtered_outputs(args.infile, args.format, args.outprefix, thresholds_info,
                           wrap_width=args.line_width, show_progress=args.show_progress)

    print("All done.")

if __name__ == "__main__":
    main()