T2T基因组组装-端粒填补

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Step1: 从 PAF 文件筛选候选端粒 reads
条件：
  1. mapq > min_mapq
  2. (比对长度 / query 长度) >= min_ratio
  3. 比对位置落在参考基因组首尾（前/后 flank 范围内）
  4. reads 的延伸部分 >= min_extend（默认 1000bp）
输出：候选 read 名列表 (txt)
"""
import argparse
def parse_paf(line: str) -> dict:
    """解析一行 PAF，返回主要字段"""
    fields = line.strip().split("\t")
    if len(fields) < 12:  # 确保至少有前12列
        return None
    return {
        "query_name": fields[0],
        "query_len": int(fields[1]),
        "query_start": int(fields[2]),
        "query_end": int(fields[3]),
        "strand": fields[4],
        "target_name": fields[5],
        "target_len": int(fields[6]),
        "target_start": int(fields[7]),
        "target_end": int(fields[8]),
        "matching bases": int(fields[9]),
        "matching lenth": int(fields[10]),
        "mapq": int(fields[11])
    }

def is_at_chrom_end(aln, flank=10000):
    """
    判断比对是否落在参考序列的首端或末端
    返回: "5prime" / "3prime" / None
    """
    if aln["target_start"] < flank:
        return "5prime"
    elif aln["target_len"] - aln["target_end"] < flank:
        return "3prime"
    return None

def check_extension(aln, chrom_end, min_extend=1000):
    """
    判断 read 是否有足够的延伸序列
    参数:
      aln: paf记录
      chrom_end: "5prime" or "3prime"
      min_extend: 最小延伸长度
    """
    qlen = aln["query_len"]
    qs, qe = aln["query_start"], aln["query_end"]

    if aln["strand"] == "+":  # 正链
        if chrom_end == "5prime":
            extend = qs  # read 左端未比对长度
            ref_anchor = aln["target_start"]
            return (extend - ref_anchor) >= min_extend
        else:  # 3prime
            extend = qlen - qe  # read 右端未比对长度
            ref_extend = aln["target_len"] - aln["target_end"]
            return (extend - ref_extend) >= min_extend
    else:  # 负链
        if chrom_end == "5prime":
            extend = qlen - qe  # read 右端未比对长度
            ref_anchor = aln["target_start"]
            return (extend - ref_anchor) >= min_extend
        else:  # 3prime
            extend = qs  # read 左端未比对长度
            ref_extend = aln["target_len"] - aln["target_end"]
            return (extend - ref_extend) >= min_extend
          
def filter_paf(paf_file, min_mapq=50, min_ratio=0.6, min_extend=1000, flank=10000, out_txt="candidate_reads.paf"):
    """
    筛选 PAF 文件，输出候选 read 的完整 PAF 行
    """
    selected_lines = []

    with open(paf_file) as f:
        for line in f:
            if not line.strip():
                continue
            aln = parse_paf(line)
            if aln is None:
                continue

            # 过滤条件1: mapq
            if aln["mapq"] < min_mapq:
                continue

            # 过滤条件2: 比对比例
            ratio = aln["matching bases"] / aln["matching lenth"]
            if ratio < min_ratio:
                continue

            # 过滤条件3: 在染色体端部
            chrom_end = is_at_chrom_end(aln, flank)
            if not chrom_end:
                continue

            # 过滤条件4: 延伸长度
            if not check_extension(aln, chrom_end, min_extend):
                continue

            # 保留完整 PAF 行
            selected_lines.append(line.strip())

    # 写结果
    with open(out_txt, "w") as out:
        for l in selected_lines:
            out.write(l + "\n")

    print(f"✅ 候选比对数量: {len(selected_lines)}")
    print(f"结果已写入: {out_txt}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="筛选 PAF 文件中的端粒候选 reads")
    parser.add_argument("-i", "--input", required=True, help="输入 PAF 文件")
    parser.add_argument("-o", "--output", default="step1_candidate_reads.txt", help="输出 txt 文件")
    parser.add_argument("--min_mapq", type=int, default=40, help="最小比对质量")
    parser.add_argument("--min_ratio", type=float, default=0.5, help="比对比例阈值 (align_len/query_len)")
    parser.add_argument("--min_extend", type=int, default=1000, help="最小延伸长度 (bp)")
    parser.add_argument("--flank", type=int, default=100000, help="参考序列端部检测范围 (bp)")
    args = parser.parse_args()

    filter_paf(args.input, args.min_mapq, args.min_ratio, args.min_extend, args.flank, args.output)

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Step2: Filter candidate reads for telomere motifs with sufficient extension

功能说明：
1. 输入：
    - step1_candidate_reads.txt : Step1 输出的 PAF 格式候选 reads，每行是完整 PAF 行
    - hifi_reads_all.fa : 原始 reads FASTA
2. 调用 seqtk telo 检测端粒 motif（CCCTAA / TTAGGG）
3. 对每个 motif occurrence 判断：
    a) 是否位于比对 overhang 区
    b) 从 motif 起始位置到参考端是否真实延伸 >= min_extend
4. 输出：
    - step2_telo_reads.txt : Step1 原始 PAF 行 + "\tmotif_len:<N>\tmotif:<NAME>\textend:<Left|Right>"
    - step2_telo_reads.fa : 被保留 reads 的原始 fasta 序列
"""

import subprocess
import os
from collections import defaultdict
from typing import List, Dict, Tuple
from Bio import SeqIO

# =============================
# 参数配置区（可根据项目调整）
# =============================
CANDIDATE_PAF = "step1_candidate_reads.txt"  # Step1 输出
READS_FASTA = "../ont.ul.fasta"            # 原始 reads FASTA
TMP_FASTA = "step2_tmp_candidate_reads.fa"  # 临时 FASTA 文件
MOTIFS = ["ccctaa", "ttaggg"]               # 需要检测的端粒 motif
MIN_TELOMERE_LEN = 1000                      # motif stretch 最小连续长度（bp）
MIN_EXTEND = 1000                             # motif 起始到参考端真实可延伸最小长度（bp）
OUT_TXT = "step2_telo_reads.txt"            # 输出 TXT 文件
OUT_FA = "step2_telo_reads.fa"              # 输出 FASTA 文件
FLANK_FOR_END = 10000                        # 判断比对是否在染色体首尾的 window（bp）

# =============================
# 工具函数模块
# =============================

def parse_paf_fields(line: str) -> Dict:
    """
    解析 PAF 格式的一行，提取关键字段并返回字典

    参数:
        line : str, PAF 文件中的一行
    返回:
        dict, 包含 qname, qlen, qstart, qend, strand, tname, tlen, tstart, tend, aln_len, mapq, 原始行
    """
    f = line.rstrip("\n").split("\t")
    parsed = {
        "line": line.rstrip("\n"),                # 原始 PAF 行
        "qname": f[0],                             # read id
        "qlen": int(f[1]),                         # read 总长度
        "qstart": int(f[2]),                        # read 比对起始
        "qend": int(f[3]),                          # read 比对结束
        "strand": f[4],                             # '+' 或 '-'
        "tname": f[5],                              # 参考染色体名
        "tlen": int(f[6]),                          # 参考染色体长度
        "tstart": int(f[7]),                        # 参考比对起始
        "tend": int(f[8]),                          # 参考比对结束
        "aln_len": int(f[10]) if len(f) > 10 and f[10].isdigit() else None,  # 对齐长度
        "mapq": int(f[11]) if len(f) > 11 and f[11].isdigit() else None      # mapping quality
    }
    return parsed

def extract_candidate_reads(read_ids: set, fasta_file: str, out_fasta: str) -> None:
    """
    从原始 fasta 中提取候选 reads 到临时 fasta 文件

    参数:
        read_ids : set[str], 候选 read id
        fasta_file : str, 原始 fasta 文件
        out_fasta : str, 输出 fasta 文件
    """
    with open(out_fasta, "w") as outfh:
        for rec in SeqIO.parse(fasta_file, "fasta-pearson"):
            if rec.id in read_ids:
                # 写入 fasta 格式
                # print("\n从原始 fasta 中提取候选 reads：{}".format(rec.id), flush=True)
                outfh.write(f">{rec.id}\n{str(rec.seq)}\n")

def run_seqtk_telo(fasta_file: str, motif: str, out_file: str) -> None:
    """
    调用 seqtk telo 检测指定 motif

    参数:
        fasta_file : str, 输入候选 reads fasta
        motif : str, motif 序列（小写或大写均可）
        out_file : str, 输出 seqtk telo 文件
    """
    cmd = ["seqtk", "telo", "-m", motif, fasta_file]
    with open(out_file, "w") as outf:
        subprocess.run(cmd, check=True, stdout=outf)

def parse_seqtk_telo_output(telo_file: str) -> Dict[str, List[Tuple[str,int,int,int]]]:
    """
    解析 seqtk telo 输出
    返回 motif 匹配信息

    参数:
        telo_file : str, seqtk telo 输出文件
    返回:
        dict: read_id -> list of (motif_name, start, end, stretch_len)
    """
    matches = defaultdict(list)
    with open(telo_file) as fh:
        for ln in fh:
            if not ln.strip() or ln.startswith("#"):
                continue
            parts = ln.strip().split("\t")
            if len(parts) < 4:
                continue
            rid = parts[0]
            try:
                start = int(parts[1])
                end = int(parts[2])
            except ValueError:
                continue
            stretch = end - start
            if stretch > 0:
                # motif_name 使用文件名标记
                motif_name = os.path.basename(telo_file).split("_")[-1].replace(".txt","").upper()
                matches[rid].append((motif_name, start, end, stretch))
    return matches

def is_at_chrom_end(aln: Dict, flank: int = FLANK_FOR_END) -> Tuple[bool,bool]:
    """
    判断一条比对是否在参考首端或尾端

    参数:
        aln : dict, parse_paf_fields 输出
        flank : int, 首尾端点判定 window
    返回:
        tuple (is_5prime, is_3prime)
    """
    return aln["tstart"] < flank, aln["tlen"] - aln["tend"] < flank

def check_extension(aln: dict, motif_start: int, motif_end: int, side: str, min_extend: int = MIN_EXTEND) -> tuple:
    """
    判断 motif 在 read 上是否属于 overhang 并且延伸长度 >= min_extend
    支持 motif 与比对端点部分重叠的情况

    参数:
        aln : dict, parse_paf_fields 输出的比对信息
        motif_start : int, motif 在 read 上的起始位置 (0-based)
        motif_end : int, motif 在 read 上的结束位置 (0-based, 不含)
        side : str, "LeftExtend" 或 "RightExtend"，表示想判断 read 的哪一端
        min_extend : int, 最小真实延伸长度

    返回:
        tuple: (bool 是否满足条件, 实际延伸长度)
    """
    qlen = aln["qlen"]
    qstart = aln["qstart"]
    qend = aln["qend"]
    tstart = aln["tstart"]
    tend = aln["tend"]
    tlen = aln["tlen"]
    strand = aln["strand"]

    ext = None

    if strand == "+":
        if side == "LeftExtend":
            # motif 在 read 左端 overhang 上的长度
            overhang_len = max(0, min(motif_end, qstart) - motif_start)
            if overhang_len > 0:
                ext = overhang_len - tstart
        elif side == "RightExtend":
            overhang_len = max(0, motif_end - max(motif_start, qend))
            if overhang_len > 0:
                ext = overhang_len - (tlen - tend)
    else:  # strand == "-"
        if side == "LeftExtend":
            overhang_len = max(0, motif_end - max(motif_start, qend))
            if overhang_len > 0:
                ext = overhang_len - tstart
        elif side == "RightExtend":
            overhang_len = max(0, min(motif_end, qstart) - motif_start)
            if overhang_len > 0:
                ext = overhang_len - (tlen - tend)

    if ext is None:
        return False, 0
    else:
        return ext >= min_extend, ext

# =============================
# 主流程函数
# =============================
def main():
    # 1) 读取 Step1 PAF 文件
    cands = defaultdict(list)  # read_id -> list of alignment dicts
    with open(CANDIDATE_PAF) as fh:
        for ln in fh:
            if not ln.strip():
                continue
            parsed = parse_paf_fields(ln)
            cands[parsed["qname"]].append(parsed)
    print(f"[INFO] 读取候选 reads: {len(cands)} 个")

    # 2) 提取候选 reads FASTA
    extract_candidate_reads(set(cands.keys()), READS_FASTA, TMP_FASTA)
    print(f"[INFO] 候选 reads 写入临时 fasta: {TMP_FASTA}")

    # 3) seqtk telo 检测 motif
    telo_matches = defaultdict(list)
    for motif in MOTIFS:
        telo_out = f"tmp_seqtk_telo_{motif}.txt"
        print(f"[INFO] 运行 seqtk telo: motif={motif} -> {telo_out}")
        run_seqtk_telo(TMP_FASTA, motif, telo_out)
        parsed = parse_seqtk_telo_output(telo_out)
        for rid, lst in parsed.items():
            telo_matches[rid].extend(lst)

    # 4) 过滤 reads：motif stretch + overhang + extension
    selected = {}
    for rid, aln_list in cands.items():
        if rid not in telo_matches:
            continue
        for aln in aln_list:
            is_5p, is_3p = is_at_chrom_end(aln)
            if not (is_5p or is_3p):
                continue  # 非端点比对，跳过
            occurrences = telo_matches[rid]
            for motif_name, s, e, stretch in occurrences:
                if stretch < MIN_TELOMERE_LEN:
                    continue
                side = None
                # 判断 motif 是否位于 overhang
                if aln["strand"] == "+":
                    if is_5p and s < aln["qstart"]:
                        side = "LeftExtend"
                    elif is_3p and s >= aln["qend"]:
                        side = "RightExtend"
                else:
                    if is_5p and s >= aln["qend"]:
                        side = "LeftExtend"
                    elif is_3p and s < aln["qstart"]:
                        side = "RightExtend"
                if side:
                    ok, ext_len = check_extension(aln, s, e, side)
                    if ok:
                        selected[rid] = {
                            "paf_line": aln["line"],
                            "motif": motif_name,
                            "motif_len": stretch,
                            "motif_start": s,
                            "motif_end": e,
                            "extend_side": side,
                            "extend_bp": ext_len
                        }
                        break
            if rid in selected:
                break

    print(f"[INFO] 满足条件的 reads 数: {len(selected)}")

    # 5) 输出 TXT 文件（改造版）
    with open(OUT_TXT, "w") as outf:
        for rid, info in selected.items():
            paf_cols = info["paf_line"].strip().split("\t")[:12]  # 保留前12列
            # 拼接新的信息: motif延伸长度、motif起止位置、motif序列、左右端信息
            new_cols = [
                f"motif_seq:{info['motif']}",
                f"motif_start:{info.get('motif_start', 'NA')}",
                f"motif_end:{info.get('motif_end', 'NA')}",
                f"motif_extend:{info['extend_bp']}",
                f"extend_side:{info['extend_side']}"
            ]
            outf.write("\t".join(paf_cols + new_cols) + "\n")
    # 输出对应 FASTA
    with open(OUT_FA, "w") as outfh:
        for rec in SeqIO.parse(READS_FASTA, "fasta-pearson"):
            if rec.id in selected:
                outfh.write(f">{rec.id}\n{str(rec.seq)}\n")

    print(f"[DONE] 结果写入: {OUT_TXT} 和 {OUT_FA}")

if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
端粒修补脚本
================
功能：
1) 读取 step2_telo_reads.txt，选择每条染色体最佳端粒 read（motif_extend 最大）。
2) 根据比对信息提取 read 的延伸序列，正负链区分，左端/右端分别处理。
3) 将延伸序列补充到参考基因组首尾。
4) 输出修补后的参考序列，同时生成详细修补日志和汇总统计。

输入：
- step2_telo_reads.txt
- reads.fasta
- ref.fa

输出：
- ref_with_telomere.fa          修补端粒后的参考序列
- telomere_fill.log.txt         逐条记录补齐信息
- telomere_fill_summary.txt     修补汇总（哪些 chr 被修补，补了哪端，长度）
"""

import sys
from Bio import SeqIO
from Bio.Seq import Seq

def load_genome(ref_fa):
    """
    读取参考基因组 fasta 文件
    返回字典 {chrom: sequence}
    """
    ref_dict = {}
    for record in SeqIO.parse(ref_fa, "fasta"):
        ref_dict[record.id] = str(record.seq)
    return ref_dict

def load_reads(read_fa):
    """
    读取 reads fasta 文件
    返回字典 {read_id: sequence}
    """
    read_dict = {}
    for record in SeqIO.parse(read_fa, "fasta"):
        read_dict[record.id] = str(record.seq)
    return read_dict

def choose_best_telo(step2_file):
    """
    从 step2_telo_reads.txt 中选择每条染色体最佳端粒 read
    标准：同一染色体 + 左/右端（side）下 motif_extend 最大
    返回字典:
        {(chrom, side): {chrom, strand, side, rid, qstart, qend, tstart, tend, qlen, tlen, extend}}
    """
    best_dict = {}
    with open(step2_file) as f:
        for line in f:
            if line.strip() == "" or line.startswith("#"):
                continue
            parts = line.strip().split()

            # 解析基本字段
            rid = parts[0]
            qlen, qstart, qend = map(int, parts[1:4])
            strand = parts[4]
            chrom = parts[5]
            tlen, tstart, tend = map(int, parts[6:9])

            # motif_extend 提取
            motif_extend = int([x for x in parts if x.startswith("motif_extend:")][0].split(":")[1])

            # extend_side 提取：LeftExtend / RightExtend → 'left' / 'right'
            extend_side = [x for x in parts if x.endswith("Extend")][0]
            extend_side = extend_side.split(":")[-1].replace("Extend","").lower()

            key = (chrom, extend_side)

            # 更新最佳 read（motif_extend 最大）
            if key not in best_dict or motif_extend > best_dict[key]["extend"]:
                best_dict[key] = {
                    "chrom": chrom,
                    "strand": strand,
                    "side": extend_side,
                    "rid": rid,
                    "qstart": qstart,
                    "qend": qend,
                    "tstart": tstart,
                    "tend": tend,
                    "qlen": qlen,
                    "tlen": tlen,
                    "extend": motif_extend
                }

    return best_dict

def extend_ref(ref_dict, read_dict, best_dict):
    """
    根据最佳 read 修补参考基因组
    左端/右端 + 正链/负链分别处理
    返回:
        - 新参考序列字典
        - 日志列表
        - 汇总列表
    """
    new_ref = ref_dict.copy()
    log_lines = []
    summary = []

    for key, info in best_dict.items():
        chrom = info["chrom"]
        strand = info["strand"]
        side = info["side"]
        rid = info["rid"]
        qstart, qend = info["qstart"], info["qend"]
        tstart, tend = info["tstart"], info["tend"]
        qlen, tlen = info["qlen"], info["tlen"]

        if rid not in read_dict:
            print(f"[警告] read {rid} 不在 reads.fasta 中，跳过 {chrom}")
            continue

        read_seq = read_dict[rid]
        extend_seq = ""

        # print("chrom:{},side:{},strand:{}".format(chrom,side,strand))
        # 左端修补
        # -----------------------------
        if side == "left":
            if strand == "+":
                # 正链左端：取 read 从 0 到 (qstart - tstart)
                extend_len = qstart - tstart
                if extend_len > 0:
                    extend_seq = read_seq[0:extend_len]
            else:
                # 负链左端：取 read 比对后未比对部分，反向互补
                # 需要保证长度不超 read 范围
                extend_len = (qlen - qend) - tstart
                if extend_len > 0:
                    raw_seq = read_seq[qend:qend+extend_len]
                    extend_seq = str(Seq(raw_seq).reverse_complement())

        # -----------------------------
        # 右端修补
        # -----------------------------
        elif side == "right":
            if strand == "+":
                # 正链右端：取 read 从 qend 到末端，减去参考已有末端 
                extend_start = qend+(tlen - tend)
                extend_seq = read_seq[extend_start:]
            else:
                # 负链右端：取 read 左端未比对部分，反向互补
                extend_len = tlen - tend
                raw_seq = read_seq[0:qstart-extend_len]
                extend_seq = str(Seq(raw_seq).reverse_complement())
        else:
            print(f"[错误] 未知 side {side}，跳过 {chrom}")
            continue

        if extend_seq == "":
            print(f"[提示] {chrom} {side}端没有需要修补的序列")
            continue

        # -----------------------------
        # 拼接序列
        # -----------------------------
        if side == "left":
            new_ref[chrom] = extend_seq + new_ref[chrom]
        else:
            new_ref[chrom] = new_ref[chrom] + extend_seq

        # 记录日志和汇总
        print(f"[修补] 染色体 {chrom} {side}端，方向 {strand}，补入 {len(extend_seq)} bp，read {rid}")
        log_lines.append(f"{chrom}\t{side}\t{strand}\t{rid}\t{len(extend_seq)}\n")
        summary.append(f"{chrom}\t{side}\t{len(extend_seq)}\n")

    return new_ref, log_lines, summary

def write_fasta(ref_dict, out_fa):
    """将序列写入 fasta 文件，每行 60bp"""
    with open(out_fa, "w") as f:
        for chrom, seq in ref_dict.items():
            f.write(f">{chrom}\n")
            for i in range(0, len(seq), 60):
                f.write(seq[i:i+60] + "\n")

# -----------------------------
# 主程序
# -----------------------------
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="端粒修补参考基因组")
    parser.add_argument("--step2_file", default="step2_telo_reads.txt",
                        help="step2_telo_reads.txt 文件，默认 step2_telo_reads.txt")
    parser.add_argument("--read_fa", default="step2_telo_reads.fa",
                        help="reads fasta 文件，默认 step2_telo_reads.fa")
    parser.add_argument("genome_fa",
                        help="参考基因组 fasta 文件（必须指定）")
    parser.add_argument("--out_fa", default="out_fix_telo_genome.fa",
                        help="修补后的参考基因组输出文件，默认 out_fix_telo_genome.fa")

    args = parser.parse_args()

    step2_file = args.step2_file
    read_fa = args.read_fa
    genome_fa = args.genome_fa
    out_fa = args.out_fa

    out_log = "step3_telomere_fill.log.txt"
    out_summary = "step3_telomere_fill_summary.txt"

    print("[步骤1] 读取参考基因组...")
    genome_dict = load_genome(genome_fa)

    print("[步骤2] 读取 reads...")
    read_dict = load_reads(read_fa)

    print("[步骤3] 筛选每条染色体最佳端粒 read...")
    best_dict = choose_best_telo(step2_file)

    print("[步骤4] 修补参考基因组...")
    new_ref, log_lines, summary = extend_ref(genome_dict, read_dict, best_dict)

    print(f"[步骤5] 输出结果到 {out_fa}")
    write_fasta(new_ref, out_fa)

    # 写入日志
    with open(out_log, "w") as f:
        f.write("chrom\tside\tstrand\tread_id\text_len\n")
        for ln in log_lines:
            f.write(ln)

    # 写入汇总
    with open(out_summary, "w") as f:
        f.write("chrom\tside\tfilled_len\n")
        for ln in summary:
            f.write(ln)

    print("========== 修补统计 ==========")
    if summary:
        for ln in summary:
            print("[汇总]", ln.strip())
    else:
        print("[汇总] 没有染色体被修补")
    print("[完成]")