Snakefile

import os
from pathlib import Path
from typing import Dict, Union, List

import pandas as pd
from snakemake.utils import min_version

min_version("5.14.0")

GB = 1_024
PathLike = Union[str, Path, os.PathLike]


# ======================================================
# Config files
# ======================================================
configfile: "config.yaml"


inclusion_expr = "excluded == 0"
samplesheet = pd.read_csv(config["samplesheet"]).query(inclusion_expr)
samplesheet.set_index("sample", drop=False, inplace=True, verify_integrity=True)
containers: Dict[str, PathLike] = config["containers"]
envs: Dict[str, PathLike] = config["envs"]
scripts: Dict[str, PathLike] = config["scripts"]
rule_log_dir = Path("logs/stderr")
decontam_db = Path("decontam_db")
data_dir = Path(config["data_dir"])
nfs_data_dir = Path(config["nfs_data_dir"])
report_dir = Path("report")
plot_dir = report_dir / "plotting"
filtered_dir = Path("filtered")
subsample_dir = Path("subsampled")
compass_vcf_dir = Path(config["compass_vcf_dir"])

# ======================================================
# Global functions and variables
# ======================================================
MM2_PRESET = "map-ont"
BWA_EXTNS = [".amb", ".ann", ".bwt", ".pac", ".sa"]
BASECALL_VER = config["basecall_model"]
technologies = ["illumina", "nanopore"]

output_files = set()
samples = set()
filter_logfiles = set()
subsample_logfiles = set()
for idx, row in samplesheet.iterrows():
    site = row["site"]
    sample = row["sample"]
    run = row["run"]
    samples.add(sample)

    for tech in technologies:
        filter_logfiles.add(
            rule_log_dir / "filter_contamination" / f"{site}" / f"{sample}.{tech}.log"
        )
        subsample_logfiles.add(
            rule_log_dir / f"subsample_{tech}_reads" / f"{site}" / f"{sample}.log"
        )
        output_files.add(plot_dir / "krona" / f"{site}" / f"{sample}.{tech}.krona.html")

output_files.add(report_dir / "composition.html")
output_files.add(report_dir / "coverage.html")
output_files.add(report_dir / "coverage.csv")


# ======================================================
# Rules
# ======================================================
localrules:
    all,


report: report_dir / "workflow.rst"


rule all:
    input:
        output_files,


rule basecall:
    input:
        fast5=nfs_data_dir / "{site}/nanopore/raw_data/{run}/multi_fast5s",
    output:
        summary=(
            data_dir
            / f"{{site}}/nanopore/raw_data/{{run}}/{BASECALL_VER}/basecall/sequencing_summary.txt"
        ),
    threads: 2
    resources:
        mem_mb=int(4 * GB),
    container:
        containers["guppy-gpu"]
    params:
        extras=" ".join(
            [
                "--recursive",
                "--compress_fastq",
                "--device 'cuda:all:100%'",
                "--disable_qscore_filtering",
                f"-c {config['basecall_config']}",
            ]
        ),
        save_path=lambda wildcards, output: Path(output.summary).parent,
    log:
        rule_log_dir / f"basecall/{BASECALL_VER}/{{site}}/{{run}}.log",
    shell:
        """
        guppy_basecaller {params.extras} \
            --num_callers {threads} \
            -i {input.fast5} \
            -s {params.save_path} > {log} 2>&1
        """


def infer_barcode_kit(wildcards) -> str:
    kits = set(samplesheet.query("run == @wildcards.run")["barcode_kit"])
    return f"--barcode_kits \"{' '.join(kits)}\""


rule demux:
    input:
        summary=rules.basecall.output.summary,
    output:
        summary=(
            data_dir
            / f"{{site}}/nanopore/raw_data/{{run}}/{BASECALL_VER}/demux/barcoding_summary.txt"
        ),
        save_path=directory(
            data_dir / f"{{site}}/nanopore/raw_data/{{run}}/{BASECALL_VER}/demux"
        ),
    threads: 8
    resources:
        mem_mb=int(8 * GB),
    container:
        containers["guppy-cpu"]
    log:
        rule_log_dir / f"demux/{BASECALL_VER}/{{site}}/{{run}}.log",
    params:
        extras="--trim_barcodes --compress_fastq",
        input_dir=lambda wildcards, input: Path(input.summary).parent,
        kit=infer_barcode_kit,
    shell:
        """
        guppy_barcoder -i {params.input_dir} -s {output.save_path} -t {threads} \
            {params.extras} {params.kit} > {log} 2>&1
        """


def infer_demux_for_sample(wildcards) -> str:
    run_for_sample = samplesheet.at[wildcards.sample, "run"]
    site_for_sample = samplesheet.at[wildcards.sample, "site"]
    return str(
        data_dir
        / f"{site_for_sample}/nanopore/raw_data/{run_for_sample}/{BASECALL_VER}/demux/barcoding_summary.txt"
    )


def infer_barcode_dir(wildcards) -> str:
    barcode = int(samplesheet.at[wildcards.sample, "barcode"])
    return f"barcode{barcode:02}"  # zero-pad to two digits


rule combine_demux_fastqs:
    input:
        demux_summary=infer_demux_for_sample,
    output:
        fastq=(
            data_dir
            / f"{{site}}/nanopore/raw_data/{{run}}/{BASECALL_VER}/{{sample}}.nanopore.fq.gz"
        ),
    threads: 1
    resources:
        mem_mb=int(GB),
    container:
        containers["base"]
    params:
        barcode_dir=infer_barcode_dir,
        demux_dir=lambda wildcards, input: Path(input.demux_summary).parent,
    log:
        rule_log_dir / f"combine_demux_fastqs/{BASECALL_VER}/{{site}}/{{run}}/{{sample}}.log",
    shell:
        "cat {params.demux_dir}/{params.barcode_dir}/*.fastq* > {output.fastq} 2> {log}"


rule build_decontamination_db:
    output:
        fasta=decontam_db / "remove_contam.fa.gz",
        metadata=decontam_db / "remove_contam.tsv",
    threads: 1
    resources:
        mem_mb=GB,
    params:
        script=scripts["download_decontam"],
        outdir=lambda wildcards, output: Path(output.fasta).parent,
    conda:
        envs["decontam_db"]
    log:
        rule_log_dir / "build_decontamination_db.log",
    shell:
        """
        perl {params.script} {params.outdir} &> {log}
        tmpfile=$(mktemp)
        sed 's/NTM\t0/NTM\t1/g' {output.metadata} > "$tmpfile"
        mv "$tmpfile" {output.metadata}
        """


rule index_decontam_db:
    input:
        fasta=rules.build_decontamination_db.output.fasta,
    output:
        bwa_index=multiext(str(decontam_db / "remove_contam.fa.gz"), *BWA_EXTNS),
        mm2_index=decontam_db / "remove_contam.fa.gz.map-ont.mmi",
    threads: 4
    resources:
        mem_mb=lambda wildcards, attempt: attempt * int(32 * GB),
    conda:
        envs["aln_tools"]
    params:
        preset=MM2_PRESET,
        extras="-I 12G",
    log:
        rule_log_dir / "index_decontam_db.log",
    shell:
        """
        bwa index {input.fasta} 2> {log}
        minimap2 {params.extras} \
            -x {params.preset} \
            -t {threads} \
            -d {output.mm2_index} \
            {input.fasta} 2>> {log}
        """


rule map_illumina_to_decontam_db:
    input:
        index=rules.index_decontam_db.output.bwa_index,
        ref=rules.build_decontamination_db.output.fasta,
        r1=data_dir / "{site}" / "illumina" / "{sample}" / "{sample}.R1.fastq.gz",
        r2=data_dir / "{site}" / "illumina" / "{sample}" / "{sample}.R2.fastq.gz",
    output:
        bam="mapped/{site}/{sample}.illumina.sorted.bam",
        index="mapped/{site}/{sample}.illumina.sorted.bam.bai",
    threads: 8
    resources:
        mem_mb=lambda wildcards, attempt: attempt * int(12 * GB),
    params:
        map_extras="-M",
    conda:
        envs["aln_tools"]
    log:
        rule_log_dir / "map_illumina_to_decontam_db" / "{site}" / "{sample}.log",
    shell:
        """
        (bwa mem {params.map_extras} -t {threads} {input.ref} {input.r1} {input.r2} | \
            samtools sort -@ {threads} -o {output.bam}) 2> {log}
        samtools index -@ {threads} {output.bam} &>> {log}
        """


def infer_nanopore_fastq(wildcards) -> str:
    run_for_sample = samplesheet.at[wildcards.sample, "run"]
    return str(
        data_dir
        / f"{wildcards.site}/nanopore/raw_data/{run_for_sample}/{BASECALL_VER}/{wildcards.sample}.nanopore.fq.gz"
    )


rule map_nanopore_to_decontam_db:
    input:
        index=rules.index_decontam_db.output.mm2_index,
        query=infer_nanopore_fastq,
    output:
        bam="mapped/{site}/{sample}.nanopore.sorted.bam",
        index="mapped/{site}/{sample}.nanopore.sorted.bam.bai",
    threads: 8
    resources:
        mem_mb=lambda wildcards, attempt: attempt * int(16 * GB),
    params:
        map_extras="-aL2",
        preset=MM2_PRESET,
    conda:
        envs["aln_tools"]
    log:
        rule_log_dir / "map_nanopore_to_decontam_db" / "{site}" / "{sample}.log",
    shell:
        """
        (minimap2 {params.map_extras} -x {params.preset} -t {threads} {input.index} {input.query} | \
            samtools sort -@ {threads} -o {output.bam}) 2> {log}
        samtools index -@ {threads} {output.bam} &>> {log}
        """


rule filter_contamination:
    input:
        bam="mapped/{site}/{sample}.{tech}.sorted.bam",
        metadata=rules.build_decontamination_db.output.metadata,
    output:
        keep_ids=filtered_dir / "{site}" / "{tech}" / "{sample}" / "keep.reads",
        contam_ids=(
            filtered_dir / "{site}" / "{tech}" / "{sample}" / "contaminant.reads"
        ),
        unmapped_ids=(
            filtered_dir / "{site}" / "{tech}" / "{sample}" / "unmapped.reads"
        ),
    threads: 1
    resources:
        mem_mb=lambda wildcards, attempt: attempt * GB,
    conda:
        envs["filter"]
    params:
        script=scripts["filter"],
        extra="--verbose --ignore-secondary",
        outdir=lambda wildcards, output: Path(output.keep_ids).parent,
    log:
        rule_log_dir / "filter_contamination" / "{site}" / "{sample}.{tech}.log",
    shell:
        """
        python {params.script} {params.extra} \
            -i {input.bam} \
            -m {input.metadata} \
            -o {params.outdir} 2> {log}
        """


rule extract_decontaminated_illumina_reads:
    input:
        r1=rules.map_illumina_to_decontam_db.input.r1,
        r2=rules.map_illumina_to_decontam_db.input.r2,
        read_ids=filtered_dir / "{site}" / "illumina" / "{sample}" / "keep.reads",
    output:
        r1=(
            filtered_dir
            / "{site}"
            / "illumina"
            / "{sample}"
            / "{sample}.filtered.R1.fastq.gz"
        ),
        r2=(
            filtered_dir
            / "{site}"
            / "illumina"
            / "{sample}"
            / "{sample}.filtered.R2.fastq.gz"
        ),
    threads: 1
    resources:
        mem_mb=lambda wildcards, attempt: int(2 * GB) * attempt,
    log:
        rule_log_dir / "extract_decontaminated_illumina_reads" / "{site}" / "{sample}.log",
    container:
        containers["seqkit"]
    shell:
        """
        seqkit grep -o {output.r1} -f {input.read_ids} {input.r1} 2> {log}
        seqkit grep -o {output.r2} -f {input.read_ids} {input.r2} 2>> {log}
        """


rule extract_decontaminated_nanopore_reads:
    input:
        r1=infer_nanopore_fastq,
        read_ids=filtered_dir / "{site}" / "nanopore" / "{sample}" / "keep.reads",
    output:
        r1=(
            filtered_dir
            / "{site}"
            / "nanopore"
            / "{sample}"
            / "{sample}.filtered.fastq.gz"
        ),
    threads: 1
    resources:
        mem_mb=lambda wildcards, attempt: int(2 * GB) * attempt,
    log:
        rule_log_dir / "extract_decontaminated_nanopore_reads" / "{site}" / "{sample}.log",
    container:
        containers["seqkit"]
    shell:
        "seqkit grep -o {output.r1} -f {input.read_ids} {input.r1} 2> {log}"


rule generate_krona_input:
    input:
        bam="mapped/{site}/{sample}.{tech}.sorted.bam",
        metadata=rules.build_decontamination_db.output.metadata,
    output:
        krona_input=plot_dir / "krona" / "{site}" / "{sample}.{tech}.krona.tsv",
    threads: 1
    resources:
        mem_mb=lambda wildcards, attempt: int(0.5 * GB) * attempt,
    conda:
        envs["generate_krona_input"]
    params:
        script=scripts["generate_krona_input"],
        extras="--ignore-secondary",
    log:
        rule_log_dir / "generate_krona_input" / "{site}" / "{sample}.{tech}.log",
    shell:
        """
        python {params.script} {params.extras} \
            -i {input.bam} -m {input.metadata} -o {output.krona_input} 2> {log}
        """


rule plot_sample_composition:
    input:
        tsv=rules.generate_krona_input.output.krona_input,
    output:
        chart=report(
            plot_dir / "krona" / "{site}" / "{sample}.{tech}.krona.html",
            category="Krona",
            caption=report_dir / "krona.rst",
        ),
    threads: 1
    resources:
        mem_mb=lambda wildcards, attempt: attempt * GB,
    container:
        containers["krona"]
    log:
        rule_log_dir / "plot_sample_composition" / "{site}" / "{sample}.{tech}.log",
    shell:
        """
        ktImportText {input.tsv} -o {output.chart} &> {log}
        """


rule assign_lineages:
    input:
        vcf=compass_vcf_dir / "{sample}.compass.vcf.gz",
        panel=config["lineage_panel"],
    output:
        assignments=report_dir / "lineage_assignment" / "{sample}.lineage.csv",
    threads: 1
    resources:
        mem_mb=GB,
    conda:
        envs["assign_lineages"]
    params:
        script=scripts["assign_lineages"],
        default_lineage=config["default_lineage"],  # the name given to samples with no hits in the panel
        max_het=1,
        max_alt_lineages=1,
        ref_lineage_position=config["ref_lineage_position"],
        extras="--verbose",
    log:
        rule_log_dir / "assign_lineages" / "{sample}.log",
    shell:
        """
        python {params.script} --input {input.vcf} \
            --panel {input.panel} \
            --output {output.assignments} \
            --default-lineage {params.default_lineage} \
            --max-het {params.max_het} \
            --ref-lineage-position {params.ref_lineage_position} \
            --max-alt-lineages {params.max_alt_lineages} {params.extras} 2> {log}
        """


rule composition_report:
    input:
        lineage=expand(
            str(report_dir / "lineage_assignment" / "{sample}.lineage.csv"),
            sample=samples,
        ),
        filter_logs=list(filter_logfiles),
    output:
        html=report(
            report_dir / "composition.html",
            category="Composition",
            caption=report_dir / "composition.rst",
        ),
    threads: 1
    resources:
        mem_mb=GB,
    conda:
        envs["composition_report"]
    params:
        script=scripts["composition_report"],
        assignment_dir=lambda wildcards, input: Path(input.lineage[0]).parent,
        logs_dir=lambda wildcards, input: Path(input.filter_logs[0]).parent.parent,
        template=config["composition_template"],
        contam_warning=5.0,
        unmapped_warning=5.0,
    log:
        rule_log_dir / "composition_report.log",
    shell:
        """
        python {params.script} --assignment-dir {params.assignment_dir} \
            --logs-dir {params.logs_dir} \
            --template {params.template} \
            -o {output.html} \
            --contam-warning {params.contam_warning} \
            --unmapped-warning {params.unmapped_warning} 2> {log}
        """


rule subsample_nanopore_reads:
    input:
        reads=rules.extract_decontaminated_nanopore_reads.output.r1,
    output:
        reads=(
            subsample_dir
            / "{site}"
            / "nanopore"
            / "{sample}"
            / "{sample}.subsampled.fastq.gz"
        ),
    threads: 1
    resources:
        mem_mb=int(0.5 * GB),
    container:
        containers["rasusa"]
    params:
        covg=config["max_nanopore_covg"],
        genome_size=config["genome_size"],
        seed=88,
    log:
        rule_log_dir / "subsample_nanopore_reads" / "{site}" / "{sample}.log",
    shell:
        """
        rasusa -c {params.covg} \
            -g {params.genome_size} \
            -i {input.reads} \
            -o {output.reads} \
            -s {params.seed} 2> {log}
        """


rule subsample_illumina_reads:
    input:
        r1=rules.extract_decontaminated_illumina_reads.output.r1,
        r2=rules.extract_decontaminated_illumina_reads.output.r2,
    output:
        r1=(
            subsample_dir
            / "{site}"
            / "illumina"
            / "{sample}"
            / "{sample}.subsampled.R1.fastq.gz"
        ),
        r2=(
            subsample_dir
            / "{site}"
            / "illumina"
            / "{sample}"
            / "{sample}.subsampled.R2.fastq.gz"
        ),
    threads: 1
    resources:
        mem_mb=int(0.5 * GB),
    container:
        containers["rasusa"]
    params:
        covg=config["max_illumina_covg"],
        genome_size=config["genome_size"],
        seed=88,
    log:
        rule_log_dir / "subsample_illumina_reads" / "{site}" / "{sample}.log",
    shell:
        """
        rasusa -c {params.covg} \
            -g {params.genome_size} \
            -i {input.r1} {input.r2} \
            -o {output.r1} {output.r2} \
            -s {params.seed} 2> {log}
        """


def infer_subsample_log_dirs(logfiles: List[PathLike]) -> List[str]:
    # rule_log_dir / f"subsample_{tech}_reads" / f"{site}" / f"{sample}.log"
    dirs = set()
    for p in logfiles:
        dirs.add(str(Path(p).parent.parent))
    return list(dirs)


rule coverage_report:
    input:
        lineage=expand(
            str(report_dir / "lineage_assignment" / "{sample}.lineage.csv"),
            sample=samples,
        ),
        subsample_logs=list(subsample_logfiles),
    output:
        html=report(
            report_dir / "coverage.html",
            category="Coverage",
            caption=report_dir / "coverage.rst",
        ),
        csv=report(report_dir / "coverage.csv", category="Coverage"),
    threads: 1
    resources:
        mem_mb=GB,
    conda:
        envs["coverage_report"]
    params:
        script=scripts["coverage_report"],
        assignment_dir=lambda wildcards, input: Path(input.lineage[0]).parent,
        logs_dir=lambda wildcards, input: infer_subsample_log_dirs(
            input.subsample_logs
        ),
    log:
        rule_log_dir / "coverage_report.log",
    shell:
        """
        python {params.script} --assignment-dir {params.assignment_dir} \
            -o {output.html} -c {output.csv} \
            {params.logs_dir} 2> {log}
        """