korusEncode

#!/usr/bin/env python3

"""
Driver script for korusEncode
Pipeline:
    ntCard -> ntStat -> korusMatrix.py

Builds windowed k-mer signal matrices for KORUS ancestry inference.
"""

import argparse
import os
import shlex
import subprocess
from packaging import version
import snakemake

KORUSENCODE_VERSION = "v0.1.0"


def set_up_parser():
    parser = argparse.ArgumentParser(
        description="korusEncode: build KORUS feature matrices from sequencing reads",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="Input reads are identified using a shared prefix."
    )

    parser.add_argument(
        "--reads",
        required=True,
        help="Prefix of input read files (FASTQ/FASTA/gz/bz2 compatible)"
    )

    parser.add_argument(
        "--kmer_file",
        required=True,
        help="SNV k-mer definition file used by korusMatrix.py"
    )

    parser.add_argument(
        "-k",
        required=True,
        type=int,
        help="k-mer size"
    )

    parser.add_argument(
        "--tile",
        default=500000,
        type=int,
        help="Window/tile size for matrix aggregation [default=500000]"
    )

    parser.add_argument(
        "-t",
        default=4,
        type=int,
        help="Number of threads [default=4]"
    )

    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        help="Print commands without executing"
    )

    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Force rerun of all rules"
    )

    parser.add_argument(
        "-V",
        "--version",
        action="version",
        version=KORUSENCODE_VERSION
    )

    return parser


def main():

    parser = set_up_parser()
    args = parser.parse_args()

    if not os.path.isfile(args.kmer_file):
        raise FileNotFoundError(
            f"k-mer definition file not found: {args.kmer_file}"
        )

    base_dir = os.path.dirname(os.path.realpath(__file__))

    reads_prefix = os.path.basename(args.reads)

    intro = [
        "Running korusEncode...",
        "Parameter settings:",
        f"\t--reads {args.reads}",
        f"\t--kmer_file {args.kmer_file}",
        f"\t-k {args.k}",
        f"\t--tile {args.tile}",
        f"\t-t {args.t}"
    ]

    print("\n".join(intro), flush=True)

    command = (
        f"snakemake "
        f"-s {base_dir}/korusEncode_pipeline.smk "
        f"--nolock -p "
        f"--cores {args.t} "
        f"--config "
        f"reads={args.reads} "
        f"kmer_file={args.kmer_file} "
        f"kmer={args.k} "
        f"tile_size={args.tile} "
        f"threads={args.t}"
    )

    if version.parse(snakemake.__version__) >= version.parse("7.8.0"):
        command += " --rerun-trigger mtime"

    if args.dry_run:
        command += " -n"

    if args.force:
        command += " -F"

    print(f"Running {command}", flush=True)

    command = shlex.split(command)

    ret = subprocess.call(command)

    if ret != 0:
        raise subprocess.SubprocessError(
            "korusEncode failed — check logs for details."
        )

    print("Done korusEncode!")


if __name__ == "__main__":
    main()