Snakefile

""" Benchmarking the OptiFit algorithm using an external reference database """
import itertools
import shutil

wildcard_constraints:
    seed='\d+',
    region='bact_v4'

configfile: 'config/config.yaml'

subworkflow prep_db:
    workdir:
        "../0_prep_db/"
    configfile:
        config['configpath']
subworkflow prep_samples:
    workdir:
        "../1_prep_samples/"
    configfile:
        config['configpath']

regions = ['bact_v4']
printrefs = ['f'] # do not include the reference sequences in the final MCC score since we only care about the dataset of interest, not the public database.
dist_thresh = config['dist_thresh']
datasets = config['datasets']
refs = config['references']
methods = config['methods']
seeds = range(1, config['seeds'] + 1)

rule rbind_optifit_dbs:
    input:
        fcns='code/R/functions.R',
        R='code/R/rbind_tsv.R',
        tsv=expand('results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/results.tsv',
                    dataset = datasets,
                    ref = refs,
                    region = regions,
                    method = methods,
                    printref = printrefs,
                    seed = seeds),
        xtra=['results/ref_sizes.tsv', # because this is the target rule
              'results/denovo_dbs.tsv']
    output:
        tsv='results/optifit_dbs_results.tsv'
    log:
        'log/rbind_optifit_dbs.txt'
    script:
        'code/R/rbind_tsv.R'

rule copy_preclust_inputs:  # because filter.seqs doesn't like hyphens in filepaths
    input:
        filter=prep_samples("data/{dataset}/processed/{dataset}.filter"),
        fasta=prep_samples("data/{dataset}/processed/{dataset}.fasta"),
        count_table=prep_samples("data/{dataset}/processed/{dataset}.count_table"),
        fasta_db=prep_db("data/{ref}/{ref}.{region}.fasta"),
        tax=prep_db("data/{ref}/{ref}.bacteria.tax")
    output:
        filter=temp("data/{dataset}_{ref}_{region}/inputs/{dataset}.filter"),
        fasta=temp("data/{dataset}_{ref}_{region}/inputs/{dataset}.fasta"),
        count_table=temp("data/{dataset}_{ref}_{region}/inputs/{dataset}.count_table"),
        fasta_db=temp("data/{dataset}_{ref}_{region}/inputs/{ref}.{region}.fasta"),
        tax=temp("data/{dataset}_{ref}_{region}/inputs/{ref}.bacteria.tax")
    params:
        outdir="data/{dataset}_{ref}_{region}/inputs/"
    shell:
        """
        for f in {input}; do
            cp $f {params.outdir}
        done
        """

rule preclust_db:
    input:
        fasta_db=rules.copy_preclust_inputs.output.fasta_db,
        fasta_sample=rules.copy_preclust_inputs.output.fasta,
        tax=rules.copy_preclust_inputs.output.tax,
        hard_filter=rules.copy_preclust_inputs.output.filter
    output:
        fasta=temp("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.filter.unique.precluster.fasta"),
        count_table=temp("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.filter.unique.precluster.pick.count_table"),
        tax=temp("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.pick.tax"),
        dist=temp("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.filter.unique.precluster.dist")
    params:
        tax="data/{dataset}_{ref}_{region}/preclust_db/{ref}.bacteria.pick.tax",
        outdir="data/{dataset}_{ref}_{region}/preclust_db/",
        cutoff=dist_thresh
    resources:
        procs=8,
        walltime_hrs=48
    log:
        "log/{dataset}_{ref}_{region}/preclust_db.log"
    shell:
        """
        mothur '#set.logfile(name={log}); set.dir(output={params.outdir});

        filter.seqs(fasta={input.fasta_db}, hard={input.hard_filter});
        unique.seqs(fasta=current, format=count);
        pre.cluster(fasta=current, count=current, diffs=2, processors={resources.procs});
        dist.seqs(fasta=current, cutoff={params.cutoff}, processors={resources.procs});

        list.seqs(fasta=current);
        get.seqs(accnos=current, taxonomy={input.tax});
        get.seqs(accnos=current, count=current);
        summary.seqs(fasta=current, count=current);
        rename.file(input={params.tax}, new={output.tax}) '
        """

rule cluster_db:
    input:
        dist=rules.preclust_db.output.dist,
        count=rules.preclust_db.output.count_table
    output:
        list=temp("results/{dataset}_{ref}_{region}/cluster_db/seed_{seed}/{ref}.{region}.filter.unique.precluster.opti_mcc.list"),
        sensspec=temp("results/{dataset}_{ref}_{region}/cluster_db/seed_{seed}/{ref}.{region}.filter.unique.precluster.opti_mcc.sensspec")
    params:
        outdir="results/{dataset}_{ref}_{region}/cluster_db/seed_{seed}/",
        cutoff=dist_thresh
    log:
        "log/{dataset}_{ref}_{region}/cluster_db.seed_{seed}.log"
    benchmark:
        "benchmarks/{dataset}_{ref}_{region}/cluster_db.seed_{seed}.txt"
    resources:
        procs=8
    shell:
        """
        mothur '#set.logfile(name={log}); set.dir(output={params.outdir});
        set.seed(seed={wildcards.seed});
        set.current(processors={resources.procs});
        cluster(column={input.dist}, count={input.count}, cutoff={params.cutoff}) '
        """

rule mutate_db_sensspec:
    input:
        tsv=rules.cluster_db.output.sensspec,
        code='code/R/mutate_db_sensspec.R',
        fcns='code/R/functions.R'
    output:
        tsv=temp("results/{dataset}_{ref}_{region}/cluster_db/seed_{seed}/{ref}.{region}.filter.unique.precluster.opti_mcc.mod.sensspec")
    script:
        'code/R/mutate_db_sensspec.R'

rule sum_denovo_dbs:
    input:
        tsv=expand("results/{dataset}_{ref}_{region}/cluster_db/seed_{seed}/{ref}.{region}.filter.unique.precluster.opti_mcc.mod.sensspec",
            dataset = datasets, ref = refs, region = regions, seed = seeds),
        code='code/R/rbind_tsv.R',
        fcns='code/R/functions.R'
    output:
        tsv="results/denovo_dbs.tsv"
    script:
        'code/R/rbind_tsv.R'

rule pick_best_cluster_seed2:
    input:
        sensspec=expand("results/{{dataset}}_{{ref}}_{{region}}/cluster_db/seed_{seed}/{{ref}}.{{region}}.filter.unique.precluster.opti_mcc.sensspec",
            seed=seeds),
        list=expand("results/{{dataset}}_{{ref}}_{{region}}/cluster_db/seed_{seed}/{{ref}}.{{region}}.filter.unique.precluster.opti_mcc.list",
            seed=seeds),
        code='code/R/pick_best_cluster_seed.R',
        fcns='code/R/functions.R'
    output:
        list="results/{dataset}_{ref}_{region}/cluster_db/best_seed.opti_mcc.list"
    script:
        'code/R/pick_best_cluster_seed.R'


rule mutate_ref_count:
    input:
        count_table=rules.preclust_db.output.count_table,
        code="code/R/mutate_ref_count.R"
    output:
        count_table=temp("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.filter.unique.precluster.pick.mod.count_table")
    script:
        "code/R/mutate_ref_count.R"


rule combine_seqs:
    input:
        fasta_sample=rules.copy_preclust_inputs.output.fasta,
        fasta_db=rules.preclust_db.output.fasta,
        count_sample=rules.copy_preclust_inputs.output.count_table,
        count_db=rules.mutate_ref_count.output.count_table
    output:
        fasta_all=temp("data/{dataset}_{ref}_{region}/combined_seqs/{dataset}_{ref}.{region}.all.fasta"),
        count_all=temp("data/{dataset}_{ref}_{region}/combined_seqs/{dataset}_{ref}.{region}.all.count_table"),
        accnos_refs=temp("data/{dataset}_{ref}_{region}/combined_seqs/{ref}.{region}.filter.unique.precluster.accnos")
    params:
        outdir="data/{dataset}_{ref}_{region}/combined_seqs/"
    log:
        "log/{dataset}_{ref}_{region}/combine_seqs.log"
    shell:
        """
        cat {input.fasta_sample} {input.fasta_db} > {output.fasta_all}
        mothur '#set.logfile(name={log}); set.dir(output={params.outdir});
            list.seqs(fasta={input.fasta_sample});
            list.seqs(fasta={input.fasta_db});
            merge.count(count={input.count_sample}-{input.count_db}, output={output.count_all})
        '
        """

rule calc_dists_combined:
    input:
        fasta=rules.combine_seqs.output.fasta_all
    output:
        column=temp("data/{dataset}_{ref}_{region}/combined_seqs/{dataset}_{ref}.{region}.all.dist")
    params:
        outdir="data/{dataset}_{ref}_{region}/combined_seqs/",
        cutoff=dist_thresh
    log:
        "log/{dataset}_{ref}_{region}/calc_dists_combined.log"
    resources:
        procs=16
    shell:
        """
        mothur '#set.logfile(name={log}); set.dir(output={params.outdir});
            dist.seqs(fasta={input.fasta}, cutoff={params.cutoff}, processors={resources.procs}) '
        """

rule fit_to_ref_db:
    input:
        fasta=rules.combine_seqs.output.fasta_all,
        count=rules.combine_seqs.output.count_all,
        column=rules.calc_dists_combined.output.column,
        reflist=rules.pick_best_cluster_seed2.output.list
    output:
        sensspec=temp('results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/{dataset}_{ref}.{region}.all.optifit_mcc.sensspec'),
        list=temp('results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/{dataset}_{ref}.{region}.all.optifit_mcc.list'),
    params:
        outdir="results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/"
    benchmark:
        "benchmarks/{dataset}_{ref}_{region}/optifit.method_{method}.printref_{printref}.seed_{seed}.txt"
    log:
        "log/{dataset}_{ref}_{region}/optifit.method_{method}.printref_{printref}.seed_{seed}.log"
    resources:
        procs=8
    shell:
        """
        mothur "#set.logfile(name={log}); set.dir(output={params.outdir});
        set.seed(seed={wildcards.seed});
        set.current(processors={resources.procs});
        cluster.fit(reflist={input.reflist}, fasta={input.fasta}, count={input.count}, column={input.column},  printref={wildcards.printref}, method={wildcards.method}) "
        """

rule list_seqs2:
    input:
        list=rules.fit_to_ref_db.output.list,
        count=rules.copy_preclust_inputs.output.count_table
    output:
        list_accnos=temp('results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/{dataset}_{ref}.{region}.all.optifit_mcc.accnos')
    params:
        outdir="results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/"
    log:
        "log/{dataset}_{ref}_{region}/list_seqs.method_{method}.printref_{printref}.seed_{seed}.log"
    shell:
        """
        mothur "#set.logfile(name={log}); set.dir(output={params.outdir});
        list.seqs(list={input.list}) "
        """

rule fraction_reads_mapped2:
    input:
        code="code/py/fraction_reads_mapped.py",
        mapped=rules.list_seqs2.output.list_accnos,
        query=rules.copy_preclust_inputs.output.count_table,
        ref=rules.preclust_db.output.count_table
    output:
        txt=temp('results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/fraction_reads_mapped.txt')
    script:
        "code/py/fraction_reads_mapped.py"

def count_seqs(filename):
    with open(filename, 'r') as infile:
        num_seqs = sum(1 for line in infile if line.startswith('>'))
    return num_seqs

rule count_ref_sizes:
    input:
        fasta=expand("data/{dataset}_{ref}_{region}/preclust_db/{ref}.{region}.filter.unique.precluster.fasta", dataset=datasets, ref=refs, region=regions)
    output:
        txt='results/ref_sizes.tsv'
    run:
        with open(output.txt, 'w') as outfile:
            outfile.write('reference\tregion\tnum_seqs\tdataset_filter\n')
            for filename in input.fasta:
                parameters = filename.split('/')[1].split('_')
                num_seqs = count_seqs(filename)
                dataset = parameters[0]
                ref = parameters[1]
                region = parameters[3]
                outfile.write(f"{ref}\t{region}\t{num_seqs}\t{dataset}\n")

rule cbind_optifit_seed2:
    input:
        tsv=[rules.fit_to_ref_db.output.sensspec,
            rules.fit_to_ref_db.benchmark,
            rules.fraction_reads_mapped2.output.txt],
        code='code/R/cbind_optifit_seed.R',
        fcns='code/R/functions.R'
    output:
        tsv='results/{dataset}_{ref}_{region}/optifit/method_{method}/printref_{printref}/seed_{seed}/results.tsv'
    script:
        'code/R/cbind_optifit_seed.R'