Concoct.snake

import os.path
#import numpy as np
from collections import Counter
#from Bio.SeqIO.FastaIO import SimpleFastaParser as sfp


# ---- get the bed file for the cut up contigs-------------------------------------------
rule cut_contigs:
    input:  fa="{group}/contigs/contigs.fa",
            gff="{group}/annotation/contigs.gff"
    output: contig="{group}/contigs/contigs_C10K.fa",
            Contig_bed=temp("{group}/annotation/contigs_C10K.bedtemp")
    priority: 50
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    message:"Use orfs annotation to cut contigs"
    conda : CONDA_ENV + "/pythonenv.yaml"
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    shell:  """{SCRIPTS}/Use_orf_to_cut.py {input.fa} {input.gff} {output.contig} {output.Contig_bed}"""


#------- extract scg ------------------------------------------------------------------------
rule extract_SCG_sequences:
    input:  annotation="{filename}_cogs_best_hits.tsv",
            gff="{filename}.gff",
            fna="{filename}.fna"
    output: "{filename}_SCG.fna"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    conda : CONDA_ENV + "/pythonenv.yaml"
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    shell:  "{SCRIPTS}/Extract_SCG.py {input.fna} {input.annotation} {SCG_DATA}/scg_cogs_min0.97_max1.03_unique_genera.txt {input.gff}>{output}"

#--------- Concoct ------------------------------------------------------------------------
#def get_initial_number_of_bins(wildcards):
#    SCF_faa_handle=open(checkpoints.extract_SCG_sequences.get(filename=wildcards.group+"/annotation/contigs").output[0])
#    nb_bin=int(2*np.median(list(Counter([header.split(" ")[1] for header,seq in SimpleFastaParser(SCF_faa_handle)]).values())))
#    return min(nb_bin,MAX_BIN_NB)

#def get_initial_number_of_bins(file):
#    nb_bin=int(2*np.median(list(Counter([header.split(" ")[1] for header,seq in sfp(open(file))]).values())))
#    return min(nb_bin,MAX_BIN_NB)


rule concoct:
    input:   cov="{group}/profile/coverage_contigs_C10K.tsv",
             fasta="{group}/contigs/contigs_C10K.fa",
             SCG="{group}/annotation/contigs_SCG.fna"
    output:  cluster="{group}/binning/concoct/clustering_gt"+str(MIN_CONTIG_SIZE)+".csv",
             Data="{group}/binning/concoct/original_data_gt%d.csv"%MIN_CONTIG_SIZE
    #log: "{group}/binning/concoct/log.txt"
    params:  min_contig_size=MIN_CONTIG_SIZE,
             max_bin_nb=MAX_BIN_NB,
             log="{group}/binning/concoct/log.txt",
             error="Not enough contigs pass the threshold"
    conda : CONDA_ENV + "/concoct.yaml"
    singularity : "docker://quay.io/annacprice/concoct:1.1.0"
    threads: 20
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    # I don't want to use checkpoints, (buggy and increase DAG resolution time?) so I'll run the code inside a python run
    shell : """
            string=$({SCRIPTS}/concoct_get_bins.py -s {input.SCG} -m {params.max_bin_nb}) 
            set +e
            concoct --coverage_file {input.cov} -i 1000 --composition_file {input.fasta} -b {wildcards.group}/binning/concoct -c $string -l {params.min_contig_size} -t {threads}
            if [ $? -ne 0 ]; then
                if grep -q '{params.error}' {params.log}; then
                    touch {output}
                    exit 0
                fi
                exit 1
            fi
            """

#--------- Get SCG table out of clustering --------------------------------------------------
rule SCG_table:
    input  : bins="{group}/binning/{binner}/clustering_{name}.csv",
             SCG="{group}/annotation/contigs_SCG.fna",
             orf_bed="{group}/annotation/orf.bed",
             split_bed="{group}/annotation/contigs_C10K.bed"
    output : "{group}/binning/{binner}/{name}_SCG_table.csv"
    conda : CONDA_ENV + "/pythonenv.yaml"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    shell  : "{SCRIPTS}/SCG_in_Bins.py {input.bins} {input.SCG} {input.orf_bed} {input.split_bed} -t {output}"


#--------- Concoct refine -------------------------------------------------------------------

rule refine:
    input:  bins="{group}/{path}clustering_gt%d.csv"%MIN_CONTIG_SIZE,
            table="{group}/{path}gt%d_SCG_table.csv"%MIN_CONTIG_SIZE,
            SCG="{group}/annotation/contigs_SCG.fna",
            Data="{group}/binning/concoct/original_data_gt%d.csv"%MIN_CONTIG_SIZE
    output: bins="{group}/{path,.*/.*/}clustering_refine.csv",
    params: temp="{group}/{path}refine.temp" 
    log:    temp("{group}/{path,.*/.*/}clustering.log")
    conda : CONDA_ENV + "/concoct.yaml"
    threads: 20
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    singularity : "docker://quay.io/annacprice/concoct:1.1.0"
    shell:  """
            if [ ! -s {input.bins} ]; then
                cp {input.bins} {output.bins}
                exit 0
            fi
            CPATH=`/usr/bin/which concoct_refine`
            if [ -w $CPATH ]
            then
              sed -i 's/values/to_numpy/g' $CPATH
              sed -i 's/as_matrix/to_numpy/g' $CPATH
              sed -i 's/int(NK), args.seed, args.threads)/ int(NK), args.seed, args.threads, 500)/g' $CPATH
            fi
            ROOTDIR=$(pwd)
            sed '1d' {input.bins}  > {params}
            cd {wildcards.group}/binning/concoct/
            concoct_refine $ROOTDIR/{params} $ROOTDIR/{input.Data} $ROOTDIR/{input.table} -t {threads} &>$ROOTDIR/{log}
            rm $ROOTDIR/{params}
            cd $ROOTDIR
            """

#--------- merge back contigs -------------------------------------------------------------------
rule merge_contigs:
    input:   refine="{path}clustering_refine.csv",
             table ="{path}refine_SCG_table.csv" 
    output:  "{path}clustering_concoct.csv"
    log:     "{path}clustering_consensus.log"
    conda : CONDA_ENV + "/pythonenv.yaml"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    shell:   "{SCRIPTS}/Consensus.py {input.refine} >{output} 2>{log}"

#--------- estimate the number of mag ------------------------------------------------------------

rule output_number_of_mag:
    input:   table="{path}/{binner}_SCG_table.csv"
    output:  mag_nb="{path}/{binner}_MAG_nb.txt",
             mag_list="{path}/{binner}_MAG_list.txt"
    run:
        with open(output["mag_nb"],"w") as handle_nb:
            with open(output["mag_list"],"w") as handle_list:
                nb=0
                mags=[]
                for index,line in enumerate(open(input["table"])) :
                    if index==0:
                        continue
                    split_line=line.rstrip().split(',')
                    if sum([element=="1" for element in split_line[1:]])>=(0.75*len(SCG)) :
                        nb+=1
                        mags.append(split_line[0])
                handle_nb.write(str(nb)+"\n")
                handle_list.write("\n".join([nb for nb in mags]))

#--------- produce a contig file for each bins be they good or not (metabat2 style) ---------------
ruleorder : metabat2>output_bins
rule output_bins:
    input:   contigs="{group}/contigs/contigs.fa",
             clustering="{group}/binning/{binner}/clustering_{binner}.csv"
    output:  "{group}/binning/{binner}/bins/done"
    conda : CONDA_ENV + "/pythonenv.yaml"
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    shell :"""
    {SCRIPTS}/Split_fasta_by_bin.py {input.clustering} $(dirname {output}) --fasta {input.contigs}
    touch {output}
    """

#--------- create a consensus binning between concoct and metabat2 ----------------

rule get_consensus_binning :
    # only support 2 binner as of now
    input : c_bin_def = "{group}/binning/concoct/clustering_concoct.csv",
            m_bin_def = "{group}/binning/metabat2/clustering_metabat2.csv",
            c_mag_list = "{group}/binning/concoct/concoct_MAG_list.txt",
            m_mag_list = "{group}/binning/metabat2/metabat2_MAG_list.txt",
            scg = "{group}/annotation/contigs_SCG.fna",
            contig_profiles = "{group}/binning/concoct/original_data_gt%s.csv"%MIN_CONTIG_SIZE,
            contig_bed = "{group}/annotation/contigs.bed"
    output : "{group}/binning/consensus/concoct_vs_metabat2/clustering_consensus.csv"
    conda : CONDA_ENV + "/pythonenv.yaml"
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    shell :"""
    {SCRIPTS}/consensus_binning.py -c_bin_def {input.c_bin_def} -m_bin_def {input.m_bin_def} -c_mag_list {input.c_mag_list} -m_mag_list {input.m_mag_list} -scg {input.scg} -contig_profiles {input.contig_profiles} -contig_bed {input.contig_bed} -o {output}
    """

#--------- run checkm as an other mag assessment option ----------------
rule output_bins_faa:
    input:   contigs="{group}/annotation/contigs.faa",
             clustering="{group}/binning/{binner}/clustering_{binner}.csv"
    output:  "{group}/binning/{binner}/quality/bins_aa/done"
    conda : CONDA_ENV + "/pythonenv.yaml"
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    shell :"""
    {SCRIPTS}/Split_fasta_by_bin.py {input.clustering} $(dirname {output}) --fasta {input.contigs}
    touch {output}
    """


#--------- Meren marker + CPR markers ----------------
# use checkm on meren hmm markers, super fast/easy
rule checkm_analyse:
    # only support 2 binner as of now
    input : done = "{group}/binning/{binner}/quality/bins_aa/done"
    output : "{group}/binning/{binner}/quality/{type}/storage/bin_stats.analyze.tsv"
    params : folder = "{group}/binning/{binner}/quality/{type}",
             tmp = "{group}/binning/{binner}/quality/{type}/tmp",
             hmm = lambda w:HMM[w.type]
    threads: 32
    conda : "%s/checkm.yaml"%CONDA_ENV
    singularity : "docker://quay.io/annacprice/gtdbtk:1.4.0"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    shell :"""
           mkdir -p {params.folder}/tmp/
           checkm analyze {params.hmm} $(dirname {input}) {params.folder} -g -x .faa --tmpdir {params.tmp} -t {threads}
           """  

rule checkm_quality_assessment :
    # only support 2 binner as of now
    input : done = "{group}/binning/{binner}/quality/{type}/storage/bin_stats.analyze.tsv"
    output : "{group}/binning/{binner}/quality/checkm_{type}.tsv"
    params : folder = "{group}/binning/{binner}/quality/{type}",
             tmp = "{group}/binning/{binner}/quality/{type}/tmp",
             hmm = lambda w:HMM[w.type]
    threads: 32
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    conda : "%s/checkm.yaml"%CONDA_ENV
    singularity : "docker://quay.io/annacprice/gtdbtk:1.4.0"
    shell :"""
           checkm qa {params.hmm} {params.folder} --tmpdir {params.tmp} -f {output} --tab_table -t {threads} -o 2
           """  

rule experimental_quality :
    input: marker = expand("{{group}}/binning/{{binner}}/quality/checkm_{type}.tsv",type=["cpr43","ar76","bac71"])
    output: out = "{group}/binning/{binner}/quality/quality_test.tsv"
    conda : CONDA_ENV + "/pythonenv.yaml"
    resources:
        slurm_partition = get_resource("partition"),
        mem_mb = get_resource("mem")
    singularity : "docker://quay.io/annacprice/pythonenv:3.9"
    shell:"{SCRIPTS}/experimental_quality.py {output.out} {input.marker}"