Snakefile_subtype

import os

# To run locally:
# snakemake --snakefile Snakefile_subtype --keep-going --cores 4 --use-singularity --singularity-prefix ~/.singularity --singularity-args "--home ~"

# To visualise the pipeline:
# snakemake --snakefile Snakefile_subtype --dag | dot -Tsvg > pipeline_subtype.svg

localrules: all

folder = os.path.abspath(config.get("folder", '..'))
data_dir = os.path.join(folder, 'data')

# The files below are the ones provided by the UK Drug Resistance Database
sequence_data = os.path.join(data_dir, 'input', 'sequencetable.csv')
sequence_fa = os.path.join(data_dir, 'input', 'sequences.fasta.xz')

# The alignment will be split N chunks, which will be jpHMM-subtyped in parallel, and the results recombined.
# N should be smaller than the number of sequences in the alignment
N = 10000

rule all:
    input:
        os.path.join(data_dir, 'metadata', 'metadata.uk.tab'),
        os.path.join(data_dir, 'aln', 'aln.uk.fa.xz'),


rule metadata:
    '''Reformats input UKHIVDB metadata.'''
    input:
        data = sequence_data,
    output:
        data = temp(os.path.join(data_dir, 'metadata', 'metadata.uk.no_subtype.tab'))
    params:
        mem = 2000,
        name = 'metadata',
        qos = 'fast'
    threads: 1
    singularity: "docker://evolbioinfo/python-evol:v3.6"
    shell:
        """
        python3 {folder}/py/prepare_metadata.py --input_data {input.data} --output_data {output.data}
        """

rule sequences:
    '''
    Unaligns input sequences.
    '''
    input:
        aln = sequence_fa,
    output:
        fa = temp(os.path.join(data_dir, 'aln', 'sequences.uk.fa'))
    singularity:
        "docker://evolbioinfo/goalign:v0.3.6"
    params:
        mem = 2000,
        name = 'unalign',
        qos = 'fast',
    threads: 1
    shell:
        """
        goalign unalign -i {input.aln} > {output.fa}
        """


rule xzip:
    '''
    Xzips the alignments.
    '''
    input:
        aln = os.path.join(data_dir, 'aln', '{aln}.fa'),
    output:
        aln = os.path.join(data_dir, 'aln', '{aln}.fa.xz'),
    params:
        mem = 120000,
        name = 'xz_{aln}',
        qos = 'fast',
    threads: 1
    shell:
        """
        xz -zc {input.aln} > {output.aln}
        """


rule unxzip_LA_aln:
    '''
    Unxzips an input LA alignment.
    '''
    input:
        fa = os.path.join(data_dir, 'input', 'HIV1_REF_2010_pol_DNA.fa.xz')
    output:
        fa = temp(os.path.join(data_dir, 'input', 'HIV1_REF_2010_pol_DNA.fa'))
    params:
        mem = 120000,
        name = 'unxz_ref',
        qos = 'fast',
        partition = 'common,dedicated,bioevo'
    threads: 1
    shell:
        """
        xz -dc {input.fa} > {output.fa}
        """


rule ref_jpHMM:
    '''Transforms an official LA alignment into a reference file for jpHMM'''
    input:
        fa = os.path.join(data_dir, 'input', 'HIV1_REF_2010_pol_DNA.fa')
    output:
        fa = temp(os.path.join(data_dir, 'input', 'jpHMM_files', 'input', 'HIV1_REF_2010_pol_DNA.fa')),
    params:
        mem = 2000,
        name = 'jphmm_ref',
        qos = 'fast',
    threads: 1
    singularity: "docker://evolbioinfo/jphmm_tools:v0.1.2"
    shell:
        """
        jphmm_ref --la_alignment {input.fa} --jphmm_ref_alignment {output.fa}
        """

rule aln_length:
    '''
    Calculates alignment length.
    '''
    input:
        aln = os.path.join(data_dir, '{prefix}.fa.xz'),
    output:
        length = os.path.join(data_dir, '{prefix}.length'),
    singularity:
        "docker://evolbioinfo/goalign:v0.3.6"
    params:
        mem = 2000,
        name = 'aln_len',
        qos = 'fast',
    threads: 1
    shell:
        """
        goalign stats length -i {input.aln} > {output.length}
        """

rule split_fa:
    '''Splits fasta into chunks.'''
    input:
        fa = os.path.join(data_dir, 'aln', 'sequences.uk.fa'),
    output:
        fa = temp(expand(os.path.join(data_dir, 'aln', 'split_aln', 'sequences.uk.{n}.fa'), n=range(N)))
    params:
        mem = 1000,
        name = 'split_fa',
        fa = os.path.join(data_dir, 'aln', 'split_aln', 'sequences.uk.{}.fa'),
        qos = 'fast'
    threads: 1
    singularity: "docker://evolbioinfo/jphmm_tools:v0.1.4"
    shell:
        """
        jphmm_split --out_fa_pattern {params.fa} --in_fa {input.fa} --n {N}
        """

rule subtype_jpHMM:
    '''Subtypes sequences with jpHMM'''
    input:
        fa = os.path.join(data_dir, 'aln', 'split_aln', 'sequences.uk.{n}.fa'),
        emission_priors = os.path.join(data_dir, 'input', 'jpHMM_files', 'priors', 'emissionPriors_HIV.txt'),
        transition_priors = os.path.join(data_dir, 'input', 'jpHMM_files', 'priors', 'transition_priors.txt'),
        aln = os.path.join(data_dir, 'input', 'jpHMM_files', 'input', 'HIV1_REF_2010_pol_DNA.fa'),
    output:
        msa = temp(os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', '{n}', 'alignment_to_msa.txt')),
        rec = temp(os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', '{n}', 'recombination.txt'))
    params:
        mem = 20000,
        name = 'subtype_{n}',
        output_dir = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', '{n}'),
        qos = 'normal'
    threads: 1
    singularity: "docker://evolbioinfo/jphmm:v03.2015"
    shell:
        """
        mkdir -p {params.output_dir}
        echo "jpHMM -s {input.fa} -a {input.emission_priors} -b {input.transition_priors} -i {input.aln} -v HIV -o {params.output_dir}"
        jpHMM -s {input.fa} -a {input.emission_priors} -b {input.transition_priors} -i {input.aln} -v HIV -o {params.output_dir}
        """

rule combine_jpHMM_output:
    '''Combines jpHMM output files'''
    input:
        msa = expand(os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', '{n}', 'alignment_to_msa.txt'), n=range(N)),
        rec = expand(os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', '{n}', 'recombination.txt'), n=range(N))
    output:
        msa = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', 'bk', 'alignment_to_msa.txt'),
        rec = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', 'bk', 'recombination.txt')
    params:
        mem = 1000,
        name = 'combine_uk',
        dir = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk'),
        qos = 'normal'
    threads: 1
    shell:
        """
        cd {params.dir}
        cat {params.dir}/*/recombination.txt > {params.dir}/bk/recombination.txt
        cat {params.dir}/*/alignment_to_msa.txt > {params.dir}/bk/alignment_to_msa.txt
        """

rule aln_jpHMM:
    '''Aligns sequences with jpHMM'''
    input:
        fa = os.path.join(data_dir, 'aln', 'sequences.uk.fa'),
        msa = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', 'bk', 'alignment_to_msa.txt'),
        length = os.path.join(data_dir, 'input', 'HIV1_REF_2010_pol_DNA.length'),
    output:
        aln = temp(os.path.join(data_dir, 'aln', 'aln.uk.fa')),
    params:
        mem = 10000,
        name = 'aln_jpHMM',
        qos = 'fast'
    threads: 1
    singularity: "docker://evolbioinfo/jphmm_tools:v0.1.4"
    shell:
        """
        jphmm_align --jphmm_msas {input.msa} --in_fas {input.fa} --out_aln {output.aln} \
         --aln_len `head -n 1 {input.length}`
        """

rule metadata_jpHMM:
    '''Adds jpHMM subtype to metadata.'''
    input:
        data = os.path.join(data_dir, 'metadata', 'metadata.uk.no_subtype.tab'),
        msa = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', 'bk', 'alignment_to_msa.txt'),
        rec = os.path.join(data_dir, 'aln', 'jpHMM_output', 'uk', 'bk', 'recombination.txt'),
        aln = os.path.join(data_dir, 'input', 'HIV1_REF_2010_pol_DNA.fa'),
    output:
        data = os.path.join(data_dir, 'metadata', 'metadata.uk.tab')
    params:
        mem = 2000,
        name = 'metadata_subtyped',
        qos = 'fast'
    threads: 1
    singularity: "docker://evolbioinfo/jphmm_tools:v0.1.4"
    shell:
        """
        python3 {folder}/py/add_jphmm_subtype.py --in_data {input.data} --out_data {output.data} \
        --in_msa {input.msa} --in_rec {input.rec} --ref_aln {input.aln}
        """