mbin_nmdc.wdl

version 1.0

import "https://code.jgi.doe.gov/metagenome-science-program/mbin-docker/-/raw/v2.1.0/v2-wdl/mbin_v2.wdl" as mbin_v2

workflow nmdc_mags {
    input {
        String proj
        File contig_file
        File sam_file
        File gff_file
        File proteins_file
        File cog_file
        File ec_file
        File ko_file
        File pfam_file
        File tigrfam_file
        File crispr_file
        File product_names_file
        File gene_phylogeny_file
        File lineage_file
        File? map_file
        Int threads=64
        String seqtype
        Boolean use_gpu = true  
        String gtdbtk_db="/refdata/GTDBTK_DB/release220"
        String gtdbtk_mash_db = '/refdata/GTDBTK_DB/mash_sketch_db_r220.msh'
        String checkm2_db="/refdata/CheckM2_database/uniref100.KO.1.dmnd"
        String eukcc2_db="/refdata/EUKCC2_DB/eukcc2_db_ver_1.2"
        String package_container = "microbiomedata/nmdc-mbin_vis:0.8.4"
        #String prok_bin_methods = 'SemiBin2:v2.2.0, CheckM2:v1.0.2, GTDB-Tk:v2.4.0, GTDB-Tk-database:release220'
        #String euk_bin_methods = 'SemiBin2:v2.2.0, EukCC:v2.1.2'
    }
    call stage {
        input:
            container=package_container,
            contig_file=contig_file,
            sam_file=sam_file,
            gff_file=gff_file,
            proteins_file=proteins_file,
            cog_file=cog_file,
            ec_file=ec_file,
            ko_file=ko_file,
            pfam_file=pfam_file,
            tigrfam_file=tigrfam_file,
            crispr_file=crispr_file,
            product_names_file=product_names_file,
            gene_phylogeny_file=gene_phylogeny_file,
            lineage_file=lineage_file,
            map_file=map_file
    }

    call check_id_map {
        input:
            container=package_container,
            contig_file=stage.contig,
            proteins_file=stage.proteins,
            map_file=stage.map_tsv
    }

    call mbin_v2.mbin {
        input:  
                contigs_file = check_id_map.contig,
                seqtype = seqtype,
                aln_file = stage.sam,
                img_gff = stage.gff,
                img_lin_tsv=stage.lineage_tsv,
                img_map = stage.map_tsv,
                threads =  threads,
                gtdbtk_db = gtdbtk_db,
                gtdbtk_mash_db = gtdbtk_mash_db, 
                checkm2_db = checkm2_db,
                eukcc_db = eukcc2_db,
                use_gpu = use_gpu
    }

    call mbin_v2_stats{
        input:
            mbin_sdb = mbin.mbin_sdb,
            contigs_file = check_id_map.contig,
            container = package_container
    }

    call package {
        input:  
                proj = proj,
                bins_tar=mbin.bins_tar,
                json_stats=mbin_v2_stats.stats_json,
                gff_file=stage.gff,
                proteins_file=stage.proteins,
                cog_file=stage.cog,
                ec_file=stage.ec,
                ko_file=stage.ko,
                pfam_file=stage.pfam,
                tigrfam_file=stage.tigrfam,
                crispr_file=stage.crispr,
                gene_phylogeny_file=stage.gene_phylogeny,
                product_names_file=stage.product_names,
                container=package_container
    }

    call finish_mags {
        input:  
            container=package_container,
            proj=proj,
            gtdbtk_json= mbin.gtdbtk_json,
            checkm = mbin.checkm_out,
            mbin_sdb = mbin.mbin_sdb,
            mbin_version = mbin.methods_file,
            stats_json = package.stats_json,
            stats_tsv = mbin_v2_stats.stats_tsv,
            hqmq_bin_tarfiles = package.hqmq_bin_tarfiles,
            lq_bin_tarfiles = package.lq_bin_tarfiles,
            barplot = package.barplot,
            heatmap = package.heatmap,
            kronaplot = package.kronaplot,
            eukcc_file=mbin.eukcc_out,
            ko_matrix = package.ko_matrix
    }

    output {
        File final_hqmq_bins_zip = finish_mags.final_hqmq_bins_zip
        File final_lq_bins_zip = finish_mags.final_lq_bins_zip
        File? final_gtdbtk_json = finish_mags.final_gtdbtk_json
        File? final_checkm = finish_mags.final_checkm
        File mags_version = finish_mags.final_version
        File final_stats_json = finish_mags.final_stats_json
        File barplot = finish_mags.final_barplot
        File heatmap = finish_mags.final_heatmap
        File kronaplot = finish_mags.final_kronaplot
    }


}

task mbin_v2_stats{
    input{
        File? mbin_sdb
        File contigs_file
        String container
    }
    command<<<
        ~{if defined(mbin_sdb)
            then "mbin_stats.py --sdb " + mbin_sdb + " --fasta " + contigs_file
            else "mbin_stats.py --fasta " + contigs_file
        }
        
        touch MAGs_stats.json
        touch MAGs_stats.tsv
    >>>

    runtime{
        docker: container
        memory: "2 GiB"
        runtime_minutes: 100
        cpu: 1
    }
    output {
        File stats_json = "MAGs_stats.json"
        File stats_tsv = "MAGs_stats.tsv"
    }

}

task mbin_nmdc {
    input{
        String name
        File   fna
        File   aln
        File   gff
        File   lineage
        File?  map_file
        Int?   threads
        Int?   pthreads
        String gtdbtk_env
        String checkm_env
        String? eukcc2_env
        String  mbin_container
    }

    command<<<
        set -euo pipefail
        export GTDBTK_DATA_PATH=~{gtdbtk_env}
        export CHECKM_DATA_PATH=~{checkm_env}
        mbin.py ~{"--threads " + threads} ~{"--pthreads " + pthreads} ~{"--map " + map_file} ~{"--eukccdb " + eukcc2_env} --fna ~{fna} --gff ~{gff} --aln ~{aln} --lintsv ~{lineage}
        mbin_stats.py $PWD
        mbin_versions.py > mbin_nmdc_versions.log
        touch MAGs_stats.tsv
    
        if [ -f  gtdbtk-output/gtdbtk.bac120.summary.tsv ]; then
            echo "bacterial summary exists."
        else
            mkdir -p gtdbtk-output
            echo "No Bacterial Results for ~{name}" > gtdbtk-output/gtdbtk.bac120.summary.tsv
        fi

        if [ -f  gtdbtk-output/gtdbtk.ar122.summary.tsv ]; then
            echo "archaeal summary exists."
        else
            mkdir -p gtdbtk-output
            echo "No Archaeal Results for ~{name}" > gtdbtk-output/gtdbtk.ar122.summary.tsv
        fi

        if [ -f checkm-qa.out ]; then
            echo "checkm summary exists."
        else
            mkdir -p gtdbtk-output
            echo "No Checkm Results for ~{name}" > checkm-qa.out
        fi

        if [ -f mbin.sdb ]; then
            echo "mbin.sdb exists."
        else
            mkdir -p gtdbtk-output
            echo "Mbin Sdb Could not be created for ~{name}" && touch mbin.sdb
        fi

        if [ -f eukcc_output/eukcc.csv.final ]; then
            echo "eukcc.csv.final exists."
        else
            mkdir -p eukcc_output
            echo "No EUKCC2 result for ~{name}" > eukcc_output/eukcc.csv.final
        fi
    >>>

    runtime{
        docker: mbin_container
        memory: "120 G"
        runtime_minutes: 100
        cpu: threads
    }

    output{
        File short = "bins.tooShort.fa"
        File low = "bins.lowDepth.fa"
        File unbinned = "bins.unbinned.fa"
        File checkm = "checkm-qa.out"
        File stats_json = "MAGs_stats.json"
        File stats_tsv = "MAGs_stats.tsv"
        File mbin_sdb = "mbin.sdb"
        File mbin_version = "mbin_nmdc_versions.log"
        File bacsum = "gtdbtk-output/gtdbtk.bac120.summary.tsv"
        File arcsum = "gtdbtk-output/gtdbtk.ar122.summary.tsv"
        File eukcc_csv = "eukcc_output/eukcc.csv.final"
        Array[File] hqmq_bin_fasta_files = glob("hqmq-metabat-bins/*fa")
        Array[File] lq_bin_fasta_files = glob("filtered-metabat-bins/*fa")
    }    
}


task stage {
    input{
        String container
        File contig_file
        File sam_file
        File gff_file
        File proteins_file
        File cog_file
        File ec_file
        File ko_file
        File pfam_file
        File tigrfam_file
        File crispr_file
        File product_names_file
        File gene_phylogeny_file
        File lineage_file
        File? map_file
        String contigs_out="contigs.fasta"
        String bam_out="pairedMapped_sorted.bam"
        String gff_out="functional_annotation.gff"
        String proteins_out="proteins.faa"
        String cog_out="cog.gff"
        String ec_out="ec.tsv"
        String ko_out="ko.tsv"
        String pfam_out="pfam.gff"
        String tigrfam_out="tigrfam.gff"
        String crispr_out="crispr.tsv"
        String products_out="products.tsv"
        String gene_phylogeny_out="gene_phylogeny.tsv"
        String lineage_out="lineage.tsv"
        String map_out="map_file.tsv"
    }
    command<<<

        set -euo pipefail

        function stage() {
            in=$1
            out=$2
            if [ $( echo $in |egrep -c "https*:") -gt 0 ] ; then
                wget $in -O $out
            else
                ln $in $out || cp $in $out
            fi
        }
        function fail {
            printf '%s\n' "$1" >&2 ## Send message to stderr.
            exit "${2-1}" ## Return a code specified by $2, or 1 by default.
        }
        
        stage ~{contig_file} ~{contigs_out} &
        stage ~{sam_file} ~{bam_out} &
        stage ~{gff_file} ~{gff_out} &
        stage ~{proteins_file} ~{proteins_out} &
        stage ~{cog_file} ~{cog_out} &
        stage ~{ec_file} ~{ec_out} &
        stage ~{ko_file} ~{ko_out} &
        stage ~{pfam_file} ~{pfam_out} &
        stage ~{tigrfam_file} ~{tigrfam_out} &
        stage ~{crispr_file} ~{crispr_out} &
        stage ~{product_names_file} ~{products_out} &
        stage ~{gene_phylogeny_file} ~{gene_phylogeny_out} &
        stage ~{lineage_file} ~{lineage_out}
        ~{"stage " + map_file + " " + map_out}
        wait
        if grep -q "Retrying" stderr; then
            if [ ! -s ~{contigs_out} ]; then
                fail "Staged Contig file is empty."
            fi
            if [ ! -s ~{bam_out} ]; then 
                fail "Staged Bam file is empty."
            fi
            if [ ! -s ~{gff_out} ]; then
                fail "Staged gff file is empty."
            fi
            if [ ! -s ~{proteins_out} ]; then
                fail "Staged proteins file is empty."
            fi
            if [ ! -s ~{cog_out} ]; then
                fail "Staged cog file is empty."
            fi
            if [ ! -s ~{ec_out} ]; then
                fail "Staged ec file is empty."
            fi
            if [ ! -s ~{ko_out} ]; then
                fail "Staged ko file is empty."
            fi
            if [ ! -s ~{pfam_out} ]; then
                fail "Staged pfam file is empty."
            fi
            if [ ! -s ~{tigrfam_out} ]; then
                fail "Staged tigrfam file is empty."
            fi
            if [ ! -s ~{crispr_out} ]; then
                fail "Staged crispr file is empty."
            fi
            if [ ! -s ~{products_out} ]; then
                fail "Staged products file is empty."
            fi
            if [ ! -s ~{gene_phylogeny_out} ]; then
                fail "Staged gene_phylogeny file is empty."
            fi
            if [ ! -s ~{lineage_out} ]; then
                fail "Staged lineage file is empty."
            fi
            if [ ! -s ~{map_out} ]; then
                fail "Staged map file is empty."
            fi
        fi 
        if grep -q "ERROR" stderr; then
            fail "Staging task has Error. Please check the $PWD/stderr and make sure there are no incomplete/partial files."
        fi
        date --iso-8601=seconds > start.txt

    >>>

    output{
        File contig = contigs_out
        File sam = bam_out
        File gff = gff_out
        File proteins = proteins_out
        File cog = cog_out
        File ec = ec_out
        File ko = ko_out
        File pfam = pfam_out
        File tigrfam = tigrfam_out
        File crispr = crispr_out
        File product_names = products_out
        File gene_phylogeny = gene_phylogeny_out
        File lineage_tsv = lineage_out
        File? map_tsv = map_out
        String start = read_string("start.txt")
    }
    runtime {
        memory: "1 GiB"
        cpu:  2
        maxRetries: 1
        docker: container
        runtime_minutes: 1400
    }
}

task check_id_map{
    input{
        String container
        File   contig_file
        File   proteins_file
        File?   map_file
    }

    command<<<
    set -euo pipefail 

    if [[ ! -f "~{map_file}" ]]; then
        echo "no map file to check id mapping."
        exit;
    fi
    
    python <<CODE
    import os, sys
    map_file = "~{map_file}"
    contig_file = "~{contig_file}"
    proteins_file = "~{proteins_file}"

    map_file_name = os.path.basename(map_file)

    contigIDs = {}
    contig_anno_IDs = {}

    with open(map_file, "r") as m_file:
        for line in m_file:
            if line.startswith("#"):
                continue
            parts = line.rstrip().split("\t")
            if len(parts) < 2:
                print(f"Bad line in {map_file_name}: {line!r}", file=sys.stderr)
                sys.exit(1)
            contigIDs[parts[0]] = 1
            contig_anno_IDs[parts[1]] = 1

    with open(contig_file, "r") as c_file:
        for line in c_file:
            if line.startswith(">"):
                seq_id = line[1:].rstrip().split()[0]
                if seq_id not in contigIDs:
                    print(f"{seq_id} is not in {map_file_name}.", file=sys.stderr)
                    sys.exit(1)

    with open(proteins_file, "r") as p_file:
        for line in p_file:
            if line.startswith(">"):
                seq_id = line[1:].rstrip().split()[0]
                contig_anno_id = "_".join(seq_id.split("_")[0:-2])
                if contig_anno_id not in contig_anno_IDs:
                    print(f"{contig_anno_id} is not in {map_file_name}.", file=sys.stderr)
                    sys.exit(1)
    CODE
    >>>

    output{
        File contig = contig_file
    }
    runtime {
        memory: "1 GiB"
        cpu:  1
        docker: container
        runtime_minutes: 100
    }
}

task package{
    input{
        String proj
        String prefix=sub(proj, ":", "_")
        File? bins_tar
        File json_stats
        File gff_file
        File proteins_file
        File cog_file
        File ec_file
        File ko_file
        File pfam_file
        File tigrfam_file
        File crispr_file
        File gene_phylogeny_file
        File product_names_file
        String container 
    }
    command<<<
        set -euo pipefail
        if [[ -f "~{bins_tar}" ]]; then
            tar xvf ~{bins_tar}
            create_tarfiles.py ~{prefix} \
                        ~{json_stats} ~{gff_file} ~{proteins_file} ~{cog_file} \
                        ~{ec_file} ~{ko_file} ~{pfam_file} ~{tigrfam_file} \
                        ~{crispr_file} ~{gene_phylogeny_file} \
                        ~{product_names_file} \
                        semibin2_out/output_bins/*
        fi

        if [ -f ~{prefix}_heatmap.pdf ]; then
            if [ -f ~{prefix}_barplot.pdf ]; then
                echo "KO analysis plot exists."
            else
                echo "There are no modules above 80% completeness. No barplot will be generated." && touch ~{prefix}_barplot.pdf
            fi
        else
            echo "No KO analysis result for ~{proj}" && touch ~{prefix}_heatmap.pdf
            echo "No KO analysis result for ~{proj}" && touch ~{prefix}_barplot.pdf
            echo "No KO analysis result for ~{proj}" > ~{prefix}_ko_krona.html
            echo "No KO analysis result for ~{proj}" > ~{prefix}_module_completeness.tab
        fi

        if [ ! -f ~{prefix}_stats.json ]; then
            cp ~{json_stats} ~{prefix}_stats.json
        fi
    >>>
    output {
        Array[File] hqmq_bin_tarfiles = flatten([glob("*_HQ.tar.gz"), glob("*_MQ.tar.gz")])
        Array[File] lq_bin_tarfiles = glob("*_LQ.tar.gz") 
        File barplot = prefix + "_barplot.pdf"
        File heatmap = prefix + "_heatmap.pdf"
        File kronaplot = prefix + "_ko_krona.html"
        File ko_matrix = prefix + "_module_completeness.tab"
        File stats_json = prefix + "_stats.json"
    }
    runtime {
        docker: container
        memory: "1 GiB"
        cpu:  1
        runtime_minutes: 100
    }
}

task finish_mags {
    input{
        String container
        File?   mbin_sdb
        File   mbin_version
        String proj
        String prefix=sub(proj, ":", "_")
        File?  gtdbtk_json
        File?  checkm
        Array[File] hqmq_bin_tarfiles
        Array[File] lq_bin_tarfiles
        File stats_json
        File stats_tsv
        Int  n_hqmq=length(hqmq_bin_tarfiles)
        Int  n_lq=length(lq_bin_tarfiles)
        File barplot
        File heatmap
        File kronaplot
        File ko_matrix
        File? eukcc_file
    }
    command<<<
        set -euo pipefail
        end=`date --iso-8601=seconds`

        if [[ -f "~{checkm}" ]]; then
            ln ~{checkm} ~{prefix}_checkm_qa.out
        fi
        if [[ -f "~{gtdbtk_json}" ]]; then
            ln ~{gtdbtk_json} ~{prefix}_gtdbtk.json
        fi
        ln ~{mbin_version} ~{prefix}_bin.info
        ln ~{barplot} ~{prefix}_barplot.pdf
        ln ~{heatmap} ~{prefix}_heatmap.pdf
        ln ~{kronaplot} ~{prefix}_kronaplot.html
        ln ~{ko_matrix} ~{prefix}_ko_matrix.txt
        if [[ -f "~{eukcc_file}" ]]; then
            ln ~{eukcc_file} ~{prefix}_eukcc.csv
        fi
        # cp all tarfiles, zip them under prefix, if empty touch no_mags.txt
        mkdir -p hqmq
        if [ ~{n_hqmq} -gt 0 ] ; then
            (cd hqmq && cp ~{sep=" " hqmq_bin_tarfiles} .)
            if [[ -f "~{mbin_sdb}" ]]; then
                (cd hqmq && cp ~{mbin_sdb} .)
            fi
            (cd hqmq && zip -j ../~{prefix}_hqmq_bin.zip * ../*pdf ../*kronaplot.html ../*ko_matrix.txt)
        else
            (cd hqmq && touch no_hqmq_mags.txt)
            if [[ -f "~{mbin_sdb}" ]]; then
                (cd hqmq && cp ~{mbin_sdb} .)
            fi
            (cd hqmq && zip -j ../~{prefix}_hqmq_bin.zip *)
        fi

        mkdir -p lq
        if [ ~{n_lq} -gt 0 ] ; then
            (cd lq && cp ~{sep=" " lq_bin_tarfiles} .)
            if [[ -f "~{mbin_sdb}" ]]; then
                (cd lq && cp ~{mbin_sdb} .)
            fi
            (cd lq && zip -j ../~{prefix}_lq_bin.zip * ../~{prefix}_eukcc.csv ../*pdf ../*kronaplot.html ../*ko_matrix.txt)
        else
            (cd lq && touch no_lq_mags.txt)
            if [[ -f "~{mbin_sdb}" ]]; then
                (cd lq && cp ~{mbin_sdb} .)
            fi
            (cd lq && zip -j ../~{prefix}_lq_bin.zip * ../~{prefix}_eukcc.csv )
        fi

        # Fix up attribute name
        cat ~{stats_json} | \
        sed 's/: null/: "null"/g' > ~{prefix}_mags_stats.json

    >>>

    output {
        File? final_checkm = "~{prefix}_checkm_qa.out"
        File final_hqmq_bins_zip = "~{prefix}_hqmq_bin.zip"
        File final_lq_bins_zip = "~{prefix}_lq_bin.zip"
        File final_stats_json = "~{prefix}_mags_stats.json"
        File? final_gtdbtk_json = "~{prefix}_gtdbtk.json"
        File final_version = "~{prefix}_bin.info"
        File final_kronaplot = "~{prefix}_kronaplot.html"
        File final_heatmap = "~{prefix}_heatmap.pdf"
        File final_barplot = "~{prefix}_barplot.pdf"
    }

    runtime {
        memory: "10 GiB"
        cpu:  4
        maxRetries: 1
        docker: container
        runtime_minutes: 1000
    }
}