pickref.py

#!/usr/bin/env python3
"""pickref.py — Rank genomes to pick a reference genome.

Replicates the dRep ``d_choose.py`` scoring formula and combines it with a
Mash all-vs-all centrality (mean Mash-ANI to every other genome in the set,
candidate set "Y").  The highest-scoring genome is a good reference for
analyses such as core-SNP calling, read mapping, or pangenomics.

Scoring formula (dRep d_choose.py):

    score = com*W_com - con*W_con + W_str*con*(het/100)
            + log10(N50)*W_n50 + log10(Length)*W_size
            + (centrality - S_ani)*W_cent + extra

Default weights (identical to dRep): W_com=1, W_con=1, W_str=1, W_n50=1,
W_size=1, W_cent=1, S_ani=0.99, extra=0.

Inputs (all table-separated):
  --list     one FASTA path per line
  --quality  table with columns: assembly, N50, Length  (any order;
             assembly must match the GCA/GCF prefix of the FASTA basenames)
  --checkm   CheckM output table (first column = bin_id, plus columns
             Completeness, Contamination, Strain heterogeneity)

Output: ranked table, highest score first.
"""
import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile

import numpy as np
import pandas as pd

DEFAULT_SKETCH_SIZE = 10000
DEFAULT_THREADS = 8
ID_REGEX = re.compile(r"^(GC[AF]_\d+\.\d+)")


def parse_args(argv=None):
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--list", required=True, metavar="FILE",
                   help="text file with one FASTA path per line")
    p.add_argument("--quality", required=True, metavar="FILE",
                   help="TSV with columns: assembly, N50, Length")
    p.add_argument("--checkm", required=True, metavar="FILE",
                   help="CheckM output TSV (first column bin_id)")
    p.add_argument("-o", "--out", default="pickref_scores.tsv", metavar="FILE",
                   help="output ranked table (default: pickref_scores.tsv)")
    p.add_argument("--mash", default=shutil.which("mash"), metavar="PATH",
                   help="mash executable (default: mash on PATH)")
    p.add_argument("--sketch-size", type=int, default=DEFAULT_SKETCH_SIZE,
                   metavar="N", help=f"mash sketch size (default: {DEFAULT_SKETCH_SIZE})")
    p.add_argument("-t", "--threads", type=int, default=DEFAULT_THREADS,
                   metavar="N", help=f"threads per mash dist (default: {DEFAULT_THREADS})")
    p.add_argument("--sketch-dir", metavar="DIR",
                   help="cache dir for .msh sketches (default: <out>/../sketch)")
    p.add_argument("--pairwise-out", metavar="FILE",
                   help="also write per-genome mean-ANI table")
    p.add_argument("--limit", type=int, metavar="N",
                   help="only score the first N genomes from --list (testing)")
    g = p.add_argument_group("dRep weights")
    g.add_argument("--w-com", type=float, default=1.0, help="completeness (1.0)")
    g.add_argument("--w-con", type=float, default=1.0, help="contamination (1.0)")
    g.add_argument("--w-str", type=float, default=1.0, help="strain heterogeneity (1.0)")
    g.add_argument("--w-n50", type=float, default=1.0, help="log10(N50) (1.0)")
    g.add_argument("--w-size", type=float, default=1.0, help="log10(Length) (1.0)")
    g.add_argument("--w-cent", type=float, default=1.0, help="centrality (1.0)")
    g.add_argument("--s-ani", type=float, default=0.99, help="centrality offset (0.99)")
    g.add_argument("--extra", type=float, default=0.0, help="bonus term (0.0)")
    return p.parse_args(argv)


def sketch_genomes(fna_list, sketch_dir, limit, mash, sketch_size):
    """Build a mash sketch per FASTA; returns list of (bin_id, fna, msh)."""
    os.makedirs(sketch_dir, exist_ok=True)
    sketches = []
    with open(fna_list) as fh:
        for line in fh:
            fna = line.strip()
            if not fna:
                continue
            if limit and len(sketches) >= limit:
                break
            base = os.path.basename(fna)
            bin_id = re.sub(r"\.(fna|fa|fasta)$", "", base, flags=re.I)
            msh = os.path.join(sketch_dir, bin_id + ".msh")
            if not (os.path.exists(msh) and os.path.getsize(msh) > 0):
                subprocess.run(
                    [mash, "sketch", "-s", str(sketch_size), "-o", msh, fna],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    check=True,
                )
            sketches.append((bin_id, fna, msh))
    return sketches


def all_vs_all_centrality(sketches, mash, threads, tmpdir):
    """Mean Mash-ANI of each genome to all others (candidate set Y)."""
    total = {m[0]: 0.0 for m in sketches}
    count = {m[0]: 0 for m in sketches}
    listf = os.path.join(tmpdir, "all.msh.list")
    with open(listf, "w") as fh:
        fh.write("\n".join(m[2] for m in sketches) + "\n")
    for i, (bin_id, fna, msh) in enumerate(sketches):
        p = subprocess.run(
            [mash, "dist", "-p", str(threads), msh, "-l", listf],
            capture_output=True, text=True,
        )
        for line in p.stdout.splitlines():
            f = line.split("\t")
            if len(f) < 5:
                continue
            q_id = re.sub(r"\.(fna|fa|fasta)$", "", os.path.basename(f[1]), flags=re.I)
            if q_id != bin_id:
                total[bin_id] += 1.0 - float(f[2])
                count[bin_id] += 1
        if (i + 1) % 500 == 0:
            print(f"  centrality {i+1}/{len(sketches)}", flush=True)
    return {m: (total[m] / count[m]) if count[m] else np.nan
            for m in total}


def load_quality(path):
    df = pd.read_csv(path, sep="\t", dtype=str)
    for col in ("assembly", "N50", "Length"):
        if col not in df.columns:
            raise SystemExit(f"--quality missing column '{col}': {path}")
    df["N50"] = pd.to_numeric(df["N50"], errors="coerce")
    df["Length"] = pd.to_numeric(df["Length"], errors="coerce")
    return df


def load_checkm(path):
    df = pd.read_csv(path, sep="\t", dtype=str)
    df = df.rename(columns={df.columns[0]: "bin_id"})
    for col in ("Completeness", "Contamination", "Strain heterogeneity"):
        if col not in df.columns:
            raise SystemExit(f"--checkm missing column '{col}': {path}")
    df = df.rename(columns={
        "Completeness": "completeness",
        "Contamination": "contamination",
        "Strain heterogeneity": "strain_het",
    })
    df["completeness"] = pd.to_numeric(df["completeness"], errors="coerce")
    df["contamination"] = pd.to_numeric(df["contamination"], errors="coerce")
    df["strain_het"] = pd.to_numeric(df["strain_het"], errors="coerce")
    # Match quality rows by either full bin_id or its GCA/GCF prefix.
    short = df["bin_id"].str.extract(ID_REGEX.pattern)[0]
    df["assembly"] = short.fillna(df["bin_id"])
    return df


def main(argv=None):
    args = parse_args(argv)
    if not args.mash:
        raise SystemExit("mash not found; install mash or pass --mash PATH")
    if not os.path.exists(args.list):
        raise SystemExit(f"--list file not found: {args.list}")

    # 1) quality + checkm
    qual = load_quality(args.quality)
    chk = load_checkm(args.checkm)
    # Also allow --quality `assembly` to be a GCA/GCF short prefix.
    short = qual["assembly"].str.extract(ID_REGEX.pattern)[0]
    qual["_key"] = short.fillna(qual["assembly"])
    chk["_key"] = chk["assembly"]
    df = qual.merge(chk.drop_duplicates(subset="assembly")[["_key", "bin_id",
                    "completeness", "contamination", "strain_het"]],
                    on="_key", how="left").drop(columns=["_key"])
    df = df[df["completeness"].notna()]
    n_missing = len(qual) - len(df)
    print(f"quality rows: {len(qual)}, matched checkM: {len(df)}"
          f" (missing {n_missing})", flush=True)
    if df.empty:
        raise SystemExit("no genomes matched between --quality and --checkm")

    # 2) mash sketches
    sketch_dir = args.sketch_dir or os.path.join(
        os.path.dirname(os.path.abspath(args.out)), "sketch")
    print("mash sketch...", flush=True)
    sketches = sketch_genomes(args.list, sketch_dir, args.limit,
                              args.mash, args.sketch_size)
    print(f"sketch done: {len(sketches)}", flush=True)

    # 3) all-vs-all centrality
    bin_ids = {m[0] for m in sketches}
    df = df[df["bin_id"].isin(bin_ids)].copy()
    if df.empty:
        raise SystemExit("no overlap between --list FASTA basenames and "
                         "--checkm bin_id")
    print(f"all-vs-all mash dist...", flush=True)
    tmpdir = tempfile.mkdtemp(prefix="pickref_")
    try:
        cent = all_vs_all_centrality(sketches, args.mash, args.threads, tmpdir)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
    df["centrality"] = df["bin_id"].map(cent)
    df["centrality"] = df["centrality"].fillna(0.0)

    # 4) dRep score
    df["score"] = (
        df["completeness"] * args.w_com
        - df["contamination"] * args.w_con
        + args.w_str * df["contamination"]
          * (df["strain_het"].fillna(0.0) / 100.0)
        + np.log10(df["N50"]) * args.w_n50
        + np.log10(df["Length"]) * args.w_size
        + (df["centrality"] - args.s_ani) * args.w_cent
        + args.extra
    )
    df = df.sort_values("score", ascending=False)

    cols = ["bin_id", "assembly", "completeness", "contamination",
            "strain_het", "N50", "Length", "centrality", "score"]
    out = df[cols]
    out.to_csv(args.out, sep="\t", index=False)
    print(f"\ndone! ranked table -> {args.out}", flush=True)

    if args.pairwise_out:
        pd.DataFrame(
            sorted(cent.items()), columns=["bin_id", "mean_ani"]
        ).to_csv(args.pairwise_out, sep="\t", index=False)
        print(f"mean-ANI table -> {args.pairwise_out}", flush=True)

    print("\nTop 10:")
    print(out.head(10).to_string(index=False))


if __name__ == "__main__":
    sys.exit(main())