Snakefile

# AMS genome assembly Snakemake pipeline

# this is based on the VGA pipeline, my own experiments with assembly, and https://doi.org/10.1186/s12864-024-10747-8.

# config
import os
import yaml

# functions and setup
with open("config-ref.yaml", "r") as f:
    config = yaml.safe_load(f)

# variable definitions

samples = config["read_loc"]
WORK_DIR = config["workdir"]

#######################################################################################

rule all:
    input:
        expand(WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa", sample=samples.keys()),
        expand(WORK_DIR + "/{sample}/full-stat-report.txt", sample=samples.keys()),
        expand(WORK_DIR + "/{sample}/contam-report.txt", sample=samples.keys()),
        expand(WORK_DIR + "/{sample}/quast/report.pdf", sample=samples.keys()),
        expand(WORK_DIR + "/{sample}/nucmer/{sample}-nucmer.png", sample=samples.keys()),
        expand("logs/{sample}.done.log", sample=samples.keys())

rule genomescope:
    message: "take {input.hifi_zip} and convert to fasta format, then run GenomeScope to determine heterozygosity etc."
    input:
        hifi_zip=lambda wildcards: samples[wildcards.sample],
    output:
        genomescope=WORK_DIR + "/{sample}/raw_genomescope_{sample}/linear_plot.png",
        fastq="hifi_reads/{sample}-reads.fastq",
        stats=WORK_DIR + "/{sample}/raw_genomescope_{sample}/SIMULATED_testing.tsv"
    log: "logs/{sample}.genomescope.log"
    params:
        dir=WORK_DIR
    conda:
        "envs/assembly-2.yaml"
    shell:
        """
        mkdir -p hifi_reads
        mkdir -p {params.dir}
        scripts/pre-qc.sh {input.hifi_zip} {wildcards.sample}
        """

rule purge_params:
    message: "using genomescope2 parameters from {input.stats}, calculate bounds for aggressive purging (if required)."
    input:
        stats=WORK_DIR + "/{sample}/raw_genomescope_{sample}/SIMULATED_testing.tsv"
    output:
        bounds=WORK_DIR + "/{sample}/raw_genomescope_{sample}/bound_stats.txt"
    log: "logs/{sample}.purge_params.log"
    shell:
        """
        scripts/purge_dup_stat.sh {input.stats} {output.bounds}
        """

rule draft_hifiasm:
    message: "build the first draft assembly from {input.hifi_fastq} with light purging and index the genome."
    input:
        hifi_fastq="hifi_reads/{sample}-reads.fastq",
        stats=WORK_DIR + "/{sample}/raw_genomescope_{sample}/bound_stats.txt",
    output:
        draft=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/{sample}.p_utg.fa",
        index= WORK_DIR + "/{sample}/hifiasm-draft_{sample}/{sample}.p_utg.fa.fai",
        dir=directory(WORK_DIR + "/{sample}/hifiasm-draft_{sample}/")
    log: "logs/{sample}.draft_hifiasm.log"
    conda:
        "envs/assembly-2.yaml"
    params:
        assembly=config["assembly_type"]
    shell:
        """
        mkdir -p {output.dir}
        scripts/hifiasm_assemble.sh {input.stats} {input.hifi_fastq} {params.assembly} {output.dir} {wildcards.sample}
        """

rule draft_qc:
    message: "first round of assembly statistics on the {input.draft} assembly."
    input:
        draft=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/{sample}.p_utg.fa",
        hifi_fastq="hifi_reads/{sample}-reads.fastq"
    output:
        genometools=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/{sample}.seqstat",
        merqury=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/merqury.qv",
        busco=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/busco.summary",
        allstats=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/aggregate.txt"   
    params:
        draft_dir=WORK_DIR + "/{sample}/hifiasm-draft_{sample}"
    log: "logs/{sample}.draft_qc.log"
    conda:
        "envs/genome-qc.yaml"
    shell:
        """
        meryl --version
        meryl k=31 count {input.hifi_fastq} output {WORK_DIR}/{wildcards.sample}/raw_genomescope_{wildcards.sample}/{wildcards.sample}-hifi_reads.meryl
        meryl histogram {WORK_DIR}/{wildcards.sample}/raw_genomescope_{wildcards.sample}/{wildcards.sample}-hifi_reads.meryl > {log}

        scripts/genome-qc.sh {wildcards.sample} {input.draft} {params.draft_dir} {input.hifi_fastq} {WORK_DIR}
        scripts/merqury.sh {wildcards.sample} {params.draft_dir}
        cat {output.genometools} {output.merqury} {output.busco} > {output.allstats}
        """

rule purge_dups:
    message: "use minimap2 and purge_dups to purge the primary assembly."
    input:
        draft=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/{sample}.p_utg.fa",
        hifi_fastq="hifi_reads/{sample}-reads.fastq",
        allstats=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/aggregate.txt"
    output:
        purged=WORK_DIR + "/{sample}/purged-draft_{sample}/{sample}_p-utg.purged.fa"
    log: "logs/{sample}.purge_dups.log"
    conda:
        "envs/purgedups.yml"
    shell:
        """
        scripts/purgedups.sh {wildcards.sample} {input.draft} {input.hifi_fastq}
        """

rule second_qc:
    message: "second round of assembly statistics on the {input.purged} assembly."
    input:
        purged=WORK_DIR + "/{sample}/purged-draft_{sample}/{sample}_p-utg.purged.fa",
        hifi_fastq="hifi_reads/{sample}-reads.fastq"
    output:
        genometools=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/{sample}.seqstat",
        merqury=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/merqury.qv",
        busco=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/busco.summary",
        allstats=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/aggregate.txt"
    params:
        purged_dir=WORK_DIR + "/{sample}/purged-draft_{sample}"
    log: "logs/{sample}.second_qc.log"
    conda:
        "envs/genome-qc.yaml"
    shell:
        """
        cp {input.purged} {params.purged_dir}/{wildcards.sample}.p_utg.fa
        scripts/genome-qc.sh {wildcards.sample} {input.purged} {params.purged_dir} {input.hifi_fastq} {WORK_DIR}
        scripts/merqury.sh {wildcards.sample} {params.purged_dir}
        cat {output.genometools} {output.merqury} {output.busco} > {output.allstats}
        """

rule inspector:
    message: "check purged assembly for errors using raw reads using Inspector."
    input:
        draft=WORK_DIR + "/{sample}/purged-draft_{sample}/{sample}_p-utg.purged.fa",
        reads="hifi_reads/{sample}-reads.fastq",
        allstats=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/aggregate.txt"
    output:
        inspector=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa"
    conda:
        "envs/corrections.yaml"
    log: "logs/inspector.{sample}.log"
    shell:
        """
        scripts/inspector.sh {wildcards.sample} {input.draft} {input.reads}
        """

rule ntlink:
    message: "fill gaps after Inspector correction with raw reads using ntLink. Note that this step will fail if the pythonpath is not set to the current conda env!"
    input:
        draft=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa",
        reads="hifi_reads/{sample}-reads.fastq"
    output:
        ntlink=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa"
    log: "logs/ntlink.{sample}.log"
    conda:
        "envs/corrections.yaml"
    params:
        dir=WORK_DIR + "/{sample}/corrected_{sample}/"
    shell:
        """
        cp {input.reads} {params.dir}{wildcards.sample}-reads.fastq
        cd {params.dir}
        ntLink_rounds run_rounds_gaps clean rounds=3 target=contig_corrected_inspectorFINAL.fa reads={wildcards.sample}-reads.fastq
        cd ../../../
        """

rule fcs_toolkit:
    message: "screen the draft genome for adapter contamination with ncbi fcs-adaptor. NOTE that this rule can be modified to also use ncbi-gx to screen for foreign contaminants, although this should be taken care of with Kraken2 in the previous step."
    input:
        fasta=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa",
    params:
        path_to_contam_dir=WORK_DIR + "/{sample}/decontam/",
        adaptor_sif="scripts/fcs-adaptor.sif"
    output:
        cleaned_seq=WORK_DIR + "/{sample}/decontam/fcs_adaptor_report.txt",
    conda:
        "envs/contam.yaml"
    log: "logs/fcs_toolkit.{sample}.log"
    shell:
        """
        scripts/run_fcsadaptor.sh --fasta-input {input.fasta} --output-dir {params.path_to_contam_dir} --euk --container-engine singularity --image {params.adaptor_sif}
        cat {output.cleaned_seq}
        """

rule mito:
    message: "search for the mitochondrial genome in the {input.draft} assembly using Heliconius clysonymus as the closest relative complete mitochondrial genome, using MitoFinder and MitFi."
    input:
        draft=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa",
        mito_genbank=config["mito_genbank"]
    output:
        mito_assembled=directory(WORK_DIR + "/{sample}/mito/{sample}_MitoFinder_mitfi_Final_Results/"),
        mito_list=WORK_DIR + "/{sample}/mito/mito-candidates.txt"        
    params:
        singularity_mitofinder="scripts/remiallio_default_mitofinder.sif"
    log: "logs/mito.{sample}.log"
    shell:
        """
        mkdir -p {WORK_DIR}/{wildcards.sample}/mito
        singularity run {params.singularity_mitofinder} -j {wildcards.sample} -a {input.draft} -r {input.mito_genbank} -o 5 -p 8 -m 8 --rename-contig yes 
        cp -r {wildcards.sample}/{wildcards.sample}_MitoFinder_mitfi_Final_Results/ {WORK_DIR}/{wildcards.sample}/mito
        echo "Mito Candidates" >> {output.mito_list}
        ls {output.mito_assembled} >> {output.mito_list}
        """

rule kraken:
    message: "check corrected and gap-filled assembly for contaminants using Kraken2."
    input:
        draft=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa",
        reads="hifi_reads/{sample}-reads.fastq"
    output:
        kraken=WORK_DIR + "/{sample}/decontam/kraken_sequences.report"         
    conda:
        "envs/contam.yaml"
    log: "logs/kraken.{sample}.log"
    shell:
        """
        mkdir -p {WORK_DIR}/{wildcards.sample}/decontam
        tr a-z A-Z < {input.draft} | dustmasker -level 40 -out {WORK_DIR}/{wildcards.sample}/decontam/tmp.fa -outfmt 'fasta'
        tr [:lower:] 'N' < {WORK_DIR}/{wildcards.sample}/decontam/tmp.fa > {WORK_DIR}/{wildcards.sample}/decontam/out.masked
        
        kraken2 --db scripts/k2_core_nt --threads 8 --conf 0.30 --report {WORK_DIR}/{wildcards.sample}/decontam/kraken_sequences.report --use-names \
        --output {WORK_DIR}/{wildcards.sample}/decontam/kraken_sequences.kraken {WORK_DIR}/{wildcards.sample}/decontam/out.masked
        """

rule third_qc:
    message: "second round of assembly statistics on the {input.decon} assembly after correction and decontamination."
    input:
        decon=WORK_DIR + "/{sample}/corrected_{sample}/contig_corrected_inspectorFINAL.fa.k32.w100.z1000.ntLink.gap_fill.3rounds.fa",
        hifi_fastq="hifi_reads/{sample}-reads.fastq"
    output:
        genometools=WORK_DIR + "/{sample}/decontam/stats/{sample}.seqstat",
        merqury=WORK_DIR + "/{sample}/decontam/stats/merqury.qv",
        busco=WORK_DIR + "/{sample}/decontam/stats/busco.summary",
        allstats=WORK_DIR + "/{sample}/decontam/stats/aggregate.txt",
        draft=WORK_DIR + "/{sample}/decontam/{sample}.p_utg.fa"
    params:
        decontam_dir=WORK_DIR + "/{sample}/decontam"
    log: "logs/third_qc.{sample}.log"
    conda:
        "envs/genome-qc.yaml"
    shell:
        """
        cp {input.decon} {params.decontam_dir}/{wildcards.sample}.p_utg.fa
        scripts/genome-qc.sh {wildcards.sample} {input.decon} {params.decontam_dir} {input.hifi_fastq} {WORK_DIR}
        scripts/merqury.sh {wildcards.sample} {params.decontam_dir}
        cat {output.genometools} {output.merqury} {output.busco} > {output.allstats}
        """

rule compile_stats:
    message: "Aggregate genome assembly statistics on raw reads and on all three QC checkpoints in the assembly process - draft, after purging, and after decontamination/correction"
    input:
        genomescope=WORK_DIR + "/{sample}/raw_genomescope_{sample}/summary.txt",
        draft=WORK_DIR + "/{sample}/hifiasm-draft_{sample}/stats/aggregate.txt",
        purged=WORK_DIR + "/{sample}/purged-draft_{sample}/stats/aggregate.txt",
        decontam=WORK_DIR + "/{sample}/decontam/stats/aggregate.txt"
    output:
        stats_report=WORK_DIR + "/{sample}/full-stat-report.txt"
    shell:
        """
        cat {input.genomescope} {input.draft} {input.purged} {input.decontam} > {output.stats_report}
        """

rule compile_contams:
    message: "Aggregate mitochondrial and contaminant information from each draft assembly."
    input:
        kraken= WORK_DIR + "/{sample}/decontam/kraken_sequences.report",
        mito= WORK_DIR + "/{sample}/mito/mito-candidates.txt",
        adapters= WORK_DIR + "/{sample}/decontam/fcs_adaptor_report.txt"
    output:
        contam_report= WORK_DIR + "/{sample}/contam-report.txt"
    shell:
        """
        touch {output.contam_report}
        cat {input.kraken} {input.adapters} >> {output.contam_report}
        if ls {WORK_DIR}/{wildcards.sample}/mito/{wildcards.sample}_MitoFinder_mitfi_Final_Results/*.infos ; then
            find {WORK_DIR}/{wildcards.sample}/mito/{wildcards.sample}_MitoFinder_mitfi_Final_Results/ -maxdepth 1 -type f -name "*infos" -exec grep -w "Initial contig name" {{}} \; >> {output.contam_report}
        else
            echo "no mitochondrial contigs identified" >> {output.contam_report}
        fi
        """

rule quast:
    message: "run QUAST on the final assembly relative to iHelSar1.2, or other reference genome"
    input:
        draft=WORK_DIR + "/{sample}/decontam/{sample}.p_utg.fa",
        reference=config["reference"]
    output:
        report= WORK_DIR + "/{sample}/quast/report.pdf",
        contigs= directory(WORK_DIR + "/{sample}/quast/contigs_reports"),
        aligned= directory(WORK_DIR + "/{sample}/quast/aligned_stats")
    conda:
        "envs/quast.yaml"
    shell:
       """
       quast.py -t 8 --eukaryote --large -r {input.reference} -o {WORK_DIR}/{wildcards.sample}/quast/ {input.draft}
       """

rule plot_nucmer:
    message: "generate nucmer plot of draft genome aligned to reference or iHelSar1.2."
    input:
        draft=WORK_DIR + "/{sample}/decontam/{sample}.p_utg.fa",
        reference=config["reference"]
    output:
        out= WORK_DIR + "/{sample}/nucmer/{sample}-nucmer.png"
    params:
        script="scripts/mummerCoordsDotPlotly.R"
    conda:
        "envs/nucmer.yaml"
    shell:
        """
        mkdir -p {WORK_DIR}/{wildcards.sample}/nucmer
        samtools faidx {input.reference}
        cp {input.reference}* {WORK_DIR}/{wildcards.sample}/nucmer && cd {WORK_DIR}/{wildcards.sample}/nucmer
        nucmer -l 100 -p {wildcards.sample} ../decontam/{wildcards.sample}.p_utg.fa {input.reference}
        delta-filter -l 10000 -q -r {wildcards.sample}.delta > {wildcards.sample}.filt.delta
        show-coords -c {wildcards.sample}.filt.delta > {wildcards.sample}.coords
        Rscript ../../../{params.script} -i {wildcards.sample}.coords -o {wildcards.sample}-nucmer -s -t -m 500 -q 10000 -l -x
        """

rule cleanup:
    message: "clean up extra files into intermediate directory."
    input:
        WORK_DIR + "/{sample}/full-stat-report.txt"
    output:
        "logs/{sample}.done.log"
    shell:
        """
        if [ -d "{wildcards.sample}" ]; then
            mkdir -p intermediate_files/
            mv {wildcards.sample}/ busco* intermediate_files
            mv *.log intermediate_files
            touch logs/{wildcards.sample}.done.log
        else
            touch logs/{wildcards.sample}.done.log
        fi
        """

#repeatmasker, workdir, mask

#z/w chromosome identification

# print done

################ REFERENCE GENOMES ONLY: #############

#minimap 2 for old assembly to new

#compare to ihelsar/nucmer/R

#genomescope of old genome?

# print done