基因组测序数据覆盖度分析

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
plot_coverage.py

对大文件分箱 coverage（每行: chr <tab> bin_start <tab> cov）按染色体绘图。
改进：对极端离群点进行鲁棒处理，避免把全染色体压扁为一条线；并可选择单染色体单文件输出。

用法示例:
    python plot_coverage.py -i hifi.depth.100bp.txt -o hifi.coverage.100bp -c 4 --smooth 3
"""

import os
import argparse
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def parse_args():
    p = argparse.ArgumentParser(description="按染色体绘制 binned coverage（chr, bin_start, cov）并做鲁棒 y 轴缩放")
    p.add_argument("-i", "--input", required=True, help="输入文件（tsv/空白分隔），每行: chr bin_start cov")
    p.add_argument("-o", "--outprefix", default="coverage", help="输出文件前缀 (生成 outprefix.pdf/outprefix.png)")
    p.add_argument("--bin-size", type=int, default=1, 
               help="每个 bin 的实际长度(bp)，默认 1 表示 bin 已是 bp。100 表示每个 bin=100bp，10000 表示10kb等")
    p.add_argument("-c", "--cols", type=int, default=4, help="每行子图个数 (默 4)")
    p.add_argument("--width", type=float, default=5.0, help="每列子图宽度 inch（默认5）")
    p.add_argument("--height", type=float, default=2.5, help="每行子图高度 inch（默认2.5）")
    p.add_argument("--dpi", type=int, default=300, help="输出 DPI（默认300）")
    p.add_argument("--smooth", type=int, default=0, help="平滑窗口（滑动平均），0 为不平滑")
    p.add_argument("--log", action="store_true", help="对 coverage 使用 log10(cov+1) 绘图")
    p.add_argument("--per-chr", action="store_true", help="是否为每个染色体单独输出图文件（放入 outprefix_per_chr 目录）")
    p.add_argument("--maxchroms", type=int, default=9999, help="限制绘制的染色体数量（用于调试）")
    return p.parse_args()

def robust_ylim(y_values):
    """
    计算一个鲁棒的 y 上限，避免极端离群点把小幅度特征压扁。
    返回 (y_top, outlier_mask)：
      - y_top: 子图应该使用的 y 上限
      - outlier_mask: 布尔数组，True 表示为被认为的极端离群点（会标记显示）
    逻辑：
      - 计算 p50, p75, p95, p99, max
      - 若 max / p95 > 10 -> 判定存在极端点，y_top = max(p95 * 1.2, p75 + 3*IQR)
      - 否则 y_top = max(max * 1.05, p99 * 1.2)
    """
    y = np.asarray(y_values)
    y = y[~np.isnan(y)]
    if y.size == 0:
        return 1.0, np.array([], dtype=bool)
    p50 = np.percentile(y, 50)
    p75 = np.percentile(y, 75)
    p95 = np.percentile(y, 95)
    p99 = np.percentile(y, 99)
    ymax = np.max(y)

    iqr = p75 - np.percentile(y, 25)
    upper_whisker = p75 + 3 * iqr

    # 防止 p95 为 0 时除法报错
    if p95 <= 0:
        ratio = float('inf') if ymax > 0 else 1.0
    else:
        ratio = ymax / float(p95)

    if ratio > 10 and p95 > 0:
        # 存在极端离群点，y_top 用 p95 的放大值或者 upper_whisker
        y_top = max(p95 * 1.2, upper_whisker if upper_whisker > 0 else p95 * 1.2)
        # 将被认为 outlier 的点定义为 > y_top
        outlier_mask = y > y_top
    else:
        # 无极端离群点或离群不严重 -> 显示真实 max 的一点空间
        y_top = max(ymax * 1.05, p99 * 1.2, 1.0)
        outlier_mask = y > y_top
    # 如果 y_top 太小（例如全0），至少显示 1
    y_top = max(y_top, 1.0)
    return float(y_top), outlier_mask

def natural_chr_sort_key(name):
    """对 chr 名称做自然排序（chr1, chr2, ..., chr10, chrX...）"""
    s = str(name)
    import re
    m = re.match(r'^(chr|Chr)?(\d+)$', s)
    if m:
        return (0, int(m.group(2)))
    # X Y MT 等排后面
    if s.lower().endswith('x'):
        return (1, 23)
    if s.lower().endswith('y'):
        return (1, 24)
    if 'mt' in s.lower() or 'mito' in s.lower():
        return (2, 0)
    return (3, s)

def plot_one_chrom(chrom_df, ax, args, chr_name):
    """在指定 ax 上绘制单个染色体曲线并做鲁棒 y 轴设置与离群点标记"""
    # 将 bin 转为 Mb 作为 x 轴
    x = chrom_df["bin"].values.astype(float) * args.bin_size / 1e6  # 转换成 Mb
    y = chrom_df["cov"].values.astype(float)

    if args.smooth and args.smooth > 1:
        y = pd.Series(y).rolling(window=args.smooth, min_periods=1, center=True).mean().values

    if args.log:
        y_plot = np.log10(y + 1.0)
    else:
        y_plot = y

    # 计算 y 上限与 outlier
    y_top, outlier_mask = robust_ylim(y_plot)

    # 画主曲线
    # linewidth 根据点数自动调整：点越多越细
    lw = 0.6 if len(x) > 2000 else 1.0
    ax.plot(x, y_plot, linewidth=lw, alpha=0.8, color="#1f77b4")

    # 标记 outliers（在 y_top 位置用红三角显示，并注释数量与最大值）
    if outlier_mask.size > 0:
        # outlier_mask 是相对于 y_plot 的 mask，但若 y_plot 被 log 化，
        # 我们也已用同一 y_plot 进行判断
        if np.any(outlier_mask):
            # 把这些点投影到 y_top*0.98 (靠近上限处) 以示意
            x_out = x[outlier_mask]
            # 把显示位置定在 y_top * 0.98（如果 log，则是 log-space）
            y_marker = np.full_like(x_out, y_top * 0.98)
            ax.scatter(x_out, y_marker, marker="^", color="red", s=6, label="outlier")
            # 在子图角落标注：outlier_count, max
            max_out = np.max(y_plot[outlier_mask])
            ax.text(0.98, 0.95, f"out:{x_out.size},max:{max_out:.0f}",
                    transform=ax.transAxes, ha="right", va="top", fontsize=7, color="red")

    # 标题、坐标、范围
    mean_cov = np.nanmean(y_plot) if y_plot.size>0 else 0.0
    ax.set_title(f"{chr_name} (bins={len(x)}, mean={mean_cov:.1f})", fontsize=9)
    ax.set_xlabel("Position (Mb)", fontsize=7)
    ax.set_ylabel("log10(cov+1)" if args.log else "Coverage", fontsize=7)
    ax.tick_params(axis="both", which="major", labelsize=6)

    # y 下限设 0 或最小值
    ymin = 0.0 if not args.log else 0.0
    ax.set_ylim(ymin, y_top)

    # 美化：自动设置 x 轴刻度数量（避免过度拥挤）
    try:
        nlabels = 5
        xmin = np.min(x)
        xmax = np.max(x)
        if np.isfinite(xmin) and np.isfinite(xmax) and xmin != xmax:
            ticks = np.linspace(xmin, xmax, min(nlabels, 6))
            ax.set_xticks(ticks)
    except Exception:
        pass

def main():
    args = parse_args()

    if not os.path.exists(args.input):
        raise SystemExit(f"[ERROR] 输入文件不存在: {args.input}")

    # 读取前三列
    print(f"[INFO] 读取 {args.input} ... （请保证每行：chr bin_start cov）")
    try:
        df = pd.read_csv(args.input, sep=r'\s+', header=None, usecols=[0,1,2], engine='python',
                         names=["chr","bin","cov"], dtype={"chr":str})
    except Exception as e:
        raise SystemExit(f"[ERROR] 读取文件失败: {e}")

    # 清理并转换类型
    df = df.dropna(subset=["chr","bin","cov"])
    df["bin"] = pd.to_numeric(df["bin"], errors="coerce").astype("Int64")
    df["cov"] = pd.to_numeric(df["cov"], errors="coerce")
    df = df.dropna(subset=["bin","cov"])
    df["bin"] = df["bin"].astype(int)
    df["cov"] = df["cov"].astype(float)

    # 染色体列表（自然排序）
    chroms = sorted(df["chr"].unique(), key=natural_chr_sort_key)
    if len(chroms) == 0:
        raise SystemExit("[ERROR] 未发现染色体信息，请检查输入文件格式。")

    # 限制数量
    if len(chroms) > args.maxchroms:
        chroms = chroms[:args.maxchroms]
        print(f"[WARN] 染色体数量超过 --maxchroms，已只处理前 {args.maxchroms} 条染色体")

    # 如果 per-chr 输出目录
    per_chr_dir = None
    if args.per_chr:
        per_chr_dir = args.outprefix + "_per_chr"
        os.makedirs(per_chr_dir, exist_ok=True)

    # 主绘图（多子图）
    n = len(chroms)
    ncols = min(max(1, args.cols), n)
    nrows = math.ceil(n / ncols)
    fig_w = ncols * args.width
    fig_h = nrows * args.height
    print(f"[INFO] 总绘图: {n} 条染色体，布局 {nrows}x{ncols}，图片尺寸 {fig_w}x{fig_h} inch")

    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(fig_w, fig_h), constrained_layout=True)
    # flatten axes
    if isinstance(axes, (list, np.ndarray)):
        axes_flat = np.array(axes).flatten()
    else:
        axes_flat = np.array([axes])

    for i, chr_name in enumerate(chroms):
        ax = axes_flat[i]
        chrom_df = df[df["chr"] == chr_name].sort_values("bin")
        if chrom_df.empty:
            ax.text(0.5, 0.5, "No data", ha="center", va="center")
            ax.set_title(chr_name)
            continue

        # 绘制主图
        plot_one_chrom(chrom_df, ax, args, chr_name)

        # 如果 per-chr 单图输出也需要单独保存
        if args.per_chr:
            # 生成单染色体图保存
            fig_single, ax_single = plt.subplots(figsize=(args.width*1.5, args.height*1.5))
            plot_one_chrom(chrom_df, ax_single, args, chr_name)
            out_png = os.path.join(per_chr_dir, f"{args.outprefix}_{chr_name}.png")
            fig_single.savefig(out_png, dpi=args.dpi, bbox_inches="tight")
            plt.close(fig_single)

    # 隐藏剩余子图
    for j in range(n, len(axes_flat)):
        try:
            axes_flat[j].axis("off")
        except Exception:
            pass

    # 总标题并保存
    fig.suptitle("Coverage across chromosomes (binned)", fontsize=14, y=1.02)
    out_pdf = args.outprefix + ".pdf"
    out_png = args.outprefix + ".png"
    print(f"[INFO] 保存总图: {out_pdf} / {out_png}")
    fig.savefig(out_pdf, dpi=args.dpi, bbox_inches="tight")
    fig.savefig(out_png, dpi=args.dpi, bbox_inches="tight")
    plt.close(fig)
    print("[DONE] 图像生成完成。")

if __name__ == "__main__":
    main()