Snakefile_without_linking

SEGMENTS = ["M", "L", "S"]


rule all:
    input:
        expand("auspice/CCHF_{Segment}.json", Segment=SEGMENTS),


dropped_strains = ("config/dropped_strains.txt",)
reference = ("config/outgroup_{Segment}.gb",)
reference_fasta = ("config/reference_{Segment}.fasta",)
reference_gff3 = ("config/genome_annotation_{Segment}.gff3",)
pathogen_json = ("config/pathogen_{Segment}.json",)
colors = ("config/colors.tsv",)
lat_longs = ("config/lat_longs.tsv",)
auspice_config = "config/auspice_config.json"
TAXON_ID = 3052518

if os.uname().sysname == "Darwin":
    # Don't use conda-forge unzip on macOS
    # Due to https://github.com/conda-forge/unzip-feedstock/issues/16
    unzip = "/usr/bin/unzip"
else:
    unzip = "unzip"


rule fetch_ncbi_dataset_package:
    output:
        dataset_package="data/ncbi_dataset.zip",
    shell:
        """
        datasets download virus genome taxon {TAXON_ID} \
            --no-progressbar \
            --filename {output.dataset_package} \
        """


rule extract_ncbi_dataset_sequences:
    input:
        dataset_package="data/ncbi_dataset.zip",
    output:
        ncbi_dataset_sequences="data/sequences.fasta",
    params:
        unzip=unzip,
    shell:
        """
        {params.unzip} -o {input.dataset_package} -d data 
        seqkit seq -w0 data/ncbi_dataset/data/genomic.fna \
        > {output.ncbi_dataset_sequences}
        """


rule format_ncbi_dataset_report:
    input:
        dataset_package="data/ncbi_dataset.zip",
    output:
        ncbi_dataset_tsv="data/metadata_post_extract.tsv",
    shell:
        """
        dataformat tsv virus-genome \
            --package {input.dataset_package} \
            > {output.ncbi_dataset_tsv}
        """


rule filter_rawdata:
    input:
        sequences="data/sequences.fasta",
        ncbi_dataset_tsv="data/metadata_post_extract.tsv",
    output:
        sequences=expand("data/sequences_{Segment}.fasta", Segment=SEGMENTS),
        metadata=expand("data/metadata_{Segment}.tsv", Segment=SEGMENTS),
    shell:
        """
        python scripts/filter_rawdata.py --no-group True \
            --sequences {input.sequences} \
            --metadata {input.ncbi_dataset_tsv}
        """


rule index_sequences:
    message:
        """
        Creating an index of sequence composition for filtering.
        """
    input:
        sequences="data/sequences_{Segment}.fasta",
    output:
        sequence_index="results/sequence_index_{Segment}.tsv",
    shell:
        """
        augur index \
            --sequences {input.sequences} \
            --output {output.sequence_index}
        """


rule filter:
    message:
        """
        Filtering to
          - {params.sequences_per_group} sequence(s) per {params.group_by!s}
          - from {params.min_date} onwards
          - excluding strains in {input.exclude}
        """
    input:
        sequences="data/sequences_{Segment}.fasta",
        sequence_index="results/sequence_index_{Segment}.tsv",
        metadata="data/metadata_{Segment}.tsv",
        exclude=dropped_strains,
    output:
        sequences="results/filtered_{Segment}.fasta",
    params:
        group_by="country year",
        sequences_per_group=20,
        id_column="accession",
        min_date=1900,
    shell:
        """
        augur filter \
            --sequences {input.sequences} \
            --sequence-index {input.sequence_index} \
            --metadata {input.metadata} \
            --exclude {input.exclude} \
            --output {output.sequences} \
            --group-by {params.group_by} \
            --sequences-per-group {params.sequences_per_group} \
            --metadata-id-columns {params.id_column} \
            --min-date {params.min_date} \
        """


rule align:
    message:
        """
        Aligning sequences to {input.reference_fasta}
        """
    input:
        sequences=rules.filter.output.sequences,
        reference_fasta=reference_fasta,
        pathogen_json=pathogen_json,
    output:
        alignment="results/aligned_{Segment}.fasta",
    shell:
        """
        nextclade run \
        --input-pathogen-json={input.pathogen_json} \
        --input-ref={input.reference_fasta} \
        --output-fasta={output.alignment} \
        --output-tsv={output.alignment}.tsv \
        --include-reference=true \
        --replace-unknown=true \
        {input.sequences}
        """


rule tree:
    message:
        "Building tree"
    input:
        alignment=rules.align.output.alignment,
        constraint_tree="config/constraint_tree_S.nwk",
    output:
        tree="results/tree_raw_{Segment}.nwk",
    shell:
        """
        if [[ "{wildcards.Segment}" == "S" ]]; then
        augur tree \
            --alignment {input.alignment} \
            --tree-builder-args="-g {input.constraint_tree}" \
            --output {output.tree}
        else
        augur tree \
            --alignment {input.alignment} \
            --output {output.tree}
        fi
        """


rule refine:
    message:
        """
        Refining tree
          - estimate timetree
          - use {params.coalescent} coalescent timescale
          - estimate {params.date_inference} node dates
          - filter tips more than {params.clock_filter_iqd} IQDs from clock expectation
        """
    input:
        tree=rules.tree.output.tree,
        alignment=rules.align.output,
        metadata="data/metadata_{Segment}.tsv",
    output:
        tree="results/tree_{Segment}.nwk",
        node_data="results/branch_lengths_{Segment}.json",
    params:
        coalescent="opt",
        date_inference="marginal",
        clock_filter_iqd=4,
        id_column="accession",
    shell:
        """
        augur refine \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --metadata {input.metadata} \
            --output-tree {output.tree} \
            --output-node-data {output.node_data} \
            --timetree \
            --coalescent {params.coalescent} \
            --date-confidence \
            --date-inference {params.date_inference} \
            --clock-filter-iqd {params.clock_filter_iqd} \
            --metadata-id-columns {params.id_column}
        """


rule ancestral:
    message:
        "Reconstructing ancestral sequences and mutations"
    input:
        tree=rules.refine.output.tree,
        alignment=rules.align.output,
        reference=reference,
    output:
        node_data="results/nt_muts_{Segment}.json",
    params:
        inference="joint",
    shell:
        """
        augur ancestral \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --output-node-data {output.node_data} \
            --root-sequence {input.reference} \
            --inference {params.inference}
        """


rule translate:
    message:
        "Translating amino acid sequences"
    input:
        tree=rules.refine.output.tree,
        node_data=rules.ancestral.output.node_data,
        reference=reference_gff3,
    output:
        node_data="results/aa_muts_{Segment}.json",
    shell:
        """
        augur translate \
            --tree {input.tree} \
            --ancestral-sequences {input.node_data} \
            --reference-sequence {input.reference} \
            --output-node-data {output.node_data} \
        """


rule clades:
    input:
        tree=rules.refine.output.tree,
        aa_muts=rules.translate.output.node_data,
        nuc_muts=rules.ancestral.output.node_data,
        clades="config/clades.tsv",
    output:
        node_data="results/clades_{Segment}.json",
    shell:
        """
        if [[ "{wildcards.Segment}" == "S" ]]; then
            augur clades \
            --tree {input.tree} \
            --mutations {input.nuc_muts} {input.aa_muts} \
            --clades {input.clades} \
            --output-node-data {output.node_data}
        else
            python scripts/get_clades.py \
            --tree {input.tree} \
            --node-data {output.node_data} \
            --clade-name {wildcards.Segment}
        fi
        """


rule traits:
    message:
        "Inferring ancestral traits for {params.columns!s}"
    input:
        tree=rules.refine.output.tree,
        metadata="data/metadata_{Segment}.tsv",
    output:
        node_data="results/traits_{Segment}.json",
    params:
        columns="region country",
        id_column="accession",
    shell:
        """
        augur traits \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --output-node-data {output.node_data} \
            --columns {params.columns} \
            --confidence \
            --metadata-id-columns {params.id_column}
        """


rule export:
    message:
        "Exporting data files for for auspice"
    input:
        tree=rules.refine.output.tree,
        metadata="data/metadata_{Segment}.tsv",
        clades=rules.clades.output.node_data,
        branch_lengths=rules.refine.output.node_data,
        traits=rules.traits.output.node_data,
        nt_muts=rules.ancestral.output.node_data,
        aa_muts=rules.translate.output.node_data,
        colors=colors,
        lat_longs=lat_longs,
        auspice_config=auspice_config,
    output:
        auspice_json="auspice/CCHF_{Segment}.json",
    params:
        id_column="accession",
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --node-data {input.branch_lengths} {input.traits} {input.nt_muts} {input.aa_muts} {input.clades} \
            --lat-longs {input.lat_longs} \
            --colors {input.colors} \
            --auspice-config {input.auspice_config} \
            --output {output.auspice_json} \
            --metadata-id-columns {params.id_column}
        """


rule clean:
    message:
        "Removing directories: {params}"
    params:
        "results ",
        "auspice",
    shell:
        "rm -rfv {params}"