Snakefile

if not config:

    configfile: "config/config.yaml"


build_names = ["lineage-b.1", "clade-iib", "all-clades", "clade-i"]
# build_names = ["lineage-b.1", "clade-iib", "all-clades"]
# build_names = ["all-clades"]

PREALIGN_REFERENCE = "resources/lineage-b.1/reference.fasta"
PREMASK_BED = "resources/lineage-b.1/mask.bed"

dataset_files = [
    "pathogen.json",
    "sequences.fasta",
    "genome_annotation.gff3",
    "README.md",
    "CHANGELOG.md",
]


# Helper functions to allow build-specific annotations
def annotation_path(wildcards):
    if "annotation" in config[wildcards.build_name]:
        return config[wildcards.build_name]["annotation"]
    return "resources/genome_annotation.gff3"


def gene_path(wildcards):
    if "genes" in config[wildcards.build_name]:
        return config[wildcards.build_name]["genes"]
    return "resources/genes.txt"


def genes(wildcards):
    gene_file = gene_path(wildcards)
    genes = [line.strip() for line in open(gene_file)]
    return genes


rule all:
    input:
        datasets=expand(
            "datasets/{dataset}/{file}", dataset=build_names, file=dataset_files
        ),
        tests=expand("test_output/{dataset}", dataset=build_names),


rule deploy_to_nextstrain_staging:
    input:
        rules.all.input,
    shell:
        """
        nextstrain deploy s3://nextstrain-staging {input}
        """


rule download:
    "Downloading ingested sequences and metadata from data.nextstrain.org"
    output:
        sequences="data/sequences.fasta.xz",
        metadata="data/metadata.tsv.gz",
    params:
        sequences_url="https://data.nextstrain.org/files/workflows/mpox/sequences.fasta.xz",
        metadata_url="https://data.nextstrain.org/files/workflows/mpox/metadata.tsv.gz",
    shell:
        """
        curl -fsSL --compressed {params.sequences_url:q} --output {output.sequences}
        curl -fsSL --compressed {params.metadata_url:q} --output {output.metadata}
        """


rule decompress:
    "Decompressing sequences and metadata"
    input:
        sequences="data/sequences.fasta.xz",
        metadata="data/metadata.tsv.gz",
    output:
        sequences="data/sequences.fasta",
        metadata="data/metadata.tsv",
    shell:
        """
        gzip --decompress --keep {input.metadata}
        xz --decompress --keep {input.sequences}
        """


rule remove_known_duplicates:
    input:
        known_duplicates="resources/known_duplicates.txt",
        sequences="data/sequences.fasta",
    output:
        sequences_without_known_duplicates="results/sequences_without_known_duplicates.fasta",
    shell:
        """
        seqkit grep -w0 -v -f {input.known_duplicates} {input.sequences} > {output.sequences_without_known_duplicates}
        """


rule index:
    "Indexing sequences"
    input:
        sequences="results/sequences_without_known_duplicates.fasta",
    output:
        index="results/index.tsv",
    shell:
        """
        augur index \
            --sequences {input.sequences} \
            --output {output.index}
        """


rule prealign:
    input:
        sequences="results/sequences_without_known_duplicates.fasta",
        reference=PREALIGN_REFERENCE,
        annotation=annotation_path,
    output:
        alignment="results/prealigned.fasta",
    shell:
        """
        nextclade3 run \
            {input.sequences} \
            --retry-reverse-complement \
            --excess-bandwidth 100 \
            --terminal-bandwidth 300 \
            --allowed-mismatches 8 \
            --window-size 40 \
            --min-seed-cover 0.1 \
            --input-ref {input.reference} \
            --gap-alignment-side left \
            --input-annotation {input.annotation} \
            --output-fasta /dev/stdout \
            --include-reference \
        | sed 's/ |(reverse complement)$//' > {output.alignment}
        """


rule exclude_bad:
    # Removing strains that do not satisfy certain requirements.
    input:
        sequences="results/sequences_without_known_duplicates.fasta",
        metadata="data/metadata.tsv",
        index="results/index.tsv",
    output:
        sequences="results/decent_sequences.fasta",
        metadata="results/decent_metadata.tsv",
        log="results/decent_filter.log",
    shell:
        """
        augur filter \
            --sequences {input.sequences} \
            --sequence-index {input.index} \
            --metadata-id-columns accession \
            --metadata {input.metadata} \
            --output-sequences {output.sequences} \
            --output-metadata {output.metadata} \
            --min-length 150000 \
            --query "(QC_rare_mutations == 'good')" \
            --output-log {output.log}
        """


rule premask:
    input:
        sequences="results/prealigned.fasta",
        mask=PREMASK_BED,
    output:
        "results/premasked.fasta",
    shell:
        """
        augur mask \
            --sequences {input.sequences} \
            --mask {input.mask} \
            --output {output}
        """


rule deduplicate:
    """
    Remove identical sequences (even if they have differing Ns)
    Keep those sequences with fewer Ns
    Focus for Nextclade is on diversity, not on representativeness
    """
    input:
        sequences="results/premasked.fasta",
    output:
        "results/duplicates.txt",
    shell:
        """
        python3 scripts/deduplicate.py {input.sequences} {output}
        """


rule subsample:
    input:
        sequences="results/decent_sequences.fasta",
        metadata="results/decent_metadata.tsv",
        exclude="resources/exclude_accessions.txt",
        specific_exclude="resources/{build_name}/exclude_accessions.txt",
        deduplicate_exclude="results/duplicates.txt",
        specific_include="resources/{build_name}/include_accessions.txt",
        include=lambda w: config[w.build_name]["general_include"],
        index="results/index.tsv",
    output:
        sequences="results/{build_name}/filtered_raw.fasta",
        metadata="results/{build_name}/filtered_raw_metadata.tsv",
        log="results/{build_name}/filtered_raw.log",
    params:
        min_date=lambda w: config[w.build_name]["min_date"],
        min_length=config["min_length"],
        exclude_where=lambda w: config[w.build_name]["exclude_where"],
        subsampling=lambda w: config[w.build_name]["subsampling"],
    shell:
        """
        augur filter \
            --sequences {input.sequences} \
            --metadata {input.metadata} \
            --sequence-index {input.index} \
            --metadata-id-columns accession \
            --exclude {input.exclude} {input.specific_exclude} {input.deduplicate_exclude} \
            {params.exclude_where} \
            {params.min_date} \
            {params.subsampling} \
            --include {input.include} {input.specific_include} \
            --probabilistic-sampling \
            --output {output.sequences} \
            --output-metadata {output.metadata} \
            --min-length {params.min_length} \
            --output-log {output.log}
        """


rule postsample_exclude:
    input:
        sequences="results/{build_name}/filtered_raw.fasta",
        metadata="results/{build_name}/filtered_raw_metadata.tsv",
        exclude="resources/{build_name}/post_filter_exclude.txt",
        index="results/index.tsv",
    output:
        sequences="results/{build_name}/filtered.fasta",
        metadata="results/{build_name}/metadata.tsv",
        log="results/{build_name}/filtered.log",
    shell:
        """
        augur filter \
            --sequences {input.sequences} \
            --metadata {input.metadata} \
            --sequence-index {input.index} \
            --metadata-id-columns accession \
            --exclude {input.exclude} \
            --output {output.sequences} \
            --output-metadata {output.metadata} \
            --output-log {output.log}
        """


rule align:
    input:
        sequences="results/{build_name}/filtered.fasta",
        reference="resources/{build_name}/reference.fasta",
        root="resources/{build_name}/root.fasta",
        annotation=annotation_path,
    output:
        alignment="results/{build_name}/aligned.fasta",
        translations=directory("results/{build_name}/translations"),
    params:
        translation_template=lambda w: f"results/{w.build_name}/translations/{{cds}}.fasta",
        genes=lambda w: ",".join(genes(w)),
    shell:
        """
        cat {input.root} {input.reference} {input.sequences} \
        | seqkit seq -w0 -i \
        | seqkit rmdup -n \
        | nextclade3 run \
            --retry-reverse-complement \
            --input-ref {input.reference} \
            --excess-bandwidth 100 \
            --terminal-bandwidth 300 \
            --allowed-mismatches 8 \
            --window-size 40 \
            --min-seed-cover 0.1 \
            --input-annotation {input.annotation} \
            --gap-alignment-side left \
            --output-fasta /dev/stdout \
            --output-translations {params.translation_template} \
        | seqkit seq -w0 -i > {output.alignment}
        """


rule mask:
    input:
        sequences="results/{build_name}/aligned.fasta",
        mask="resources/{build_name}/mask.bed",
    output:
        "results/{build_name}/masked.fasta",
    shell:
        """
        augur mask \
            --sequences {input.sequences} \
            --mask {input.mask} \
            --output {output}
        """


rule tree:
    input:
        alignment="results/{build_name}/masked.fasta",
    output:
        tree="results/{build_name}/tree_raw.nwk",
    shell:
        """
        augur tree \
            --alignment {input.alignment} \
            --output {output.tree} \
            --nthreads 3 \
            --tree-builder-args '--redo --pathogen-force --tree-search EXHAUSTIVE  --output-multifurcating-tree --make-consistent'
        """


rule collapse:
    input:
        script="scripts/collapse-zero-branches.py",
        tree="results/{build_name}/tree_raw.nwk",
    output:
        tree="results/{build_name}/tree_collapsed.nwk",
    shell:
        """
        python {input.script} \
            --threshold 0.0000001 \
            --verbose \
            --input-tree {input.tree} \
            --output-tree {output.tree}
        """


rule refine:
    input:
        tree="results/{build_name}/tree_collapsed.nwk",
        alignment="results/{build_name}/masked.fasta",
        metadata="results/{build_name}/metadata.tsv",
    output:
        tree="results/{build_name}/tree.nwk",
        node_data="results/{build_name}/branch_lengths.json",
    params:
        root=lambda w: config[w.build_name]["root"],
    shell:
        """
        augur refine \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --metadata {input.metadata} \
            --output-tree {output.tree} \
            --root {params.root} \
            --divergence-unit mutations \
            --keep-polytomies \
            --output-node-data {output.node_data}
        """


rule ancestral:
    # "Reconstructing ancestral sequences and mutations"
    input:
        tree="results/{build_name}/tree.nwk",
        alignment="results/{build_name}/aligned.fasta",
        annotation=annotation_path,
        genes=lambda w: gene_path(w),
        translations="results/{build_name}/translations",
        reference="resources/{build_name}/reference.fasta",
    output:
        node_data="results/{build_name}/nt_muts.json",
    params:
        inference="joint",
        translation_template="results/{build_name}/translations/%GENE.fasta",
        genes=lambda w: " ".join(genes(w)),
    shell:
        """
        augur ancestral \
            --tree {input.tree} \
            --annotation {input.annotation} \
            --alignment {input.alignment} \
            --root-sequence {input.reference} \
            --infer-ambiguous \
            --translations {params.translation_template} \
            --genes {params.genes} \
            --output-node-data {output.node_data} \
            --inference joint
        """


rule clades:
    # "Adding internal clade labels"
    input:
        tree="results/{build_name}/tree.nwk",
        nuc_muts="results/{build_name}/nt_muts.json",
        clades="resources/{build_name}/clades.tsv",
    output:
        node_data="results/{build_name}/clades_raw.json",
    shell:
        """
        augur clades \
            --tree {input.tree} \
            --mutations {input.nuc_muts} \
            --clades {input.clades} \
            --output-node-data {output.node_data} 2>&1 | tee {log}
        """


rule rename_clades:
    input:
        "results/{build_name}/clades_raw.json",
    output:
        node_data="results/{build_name}/clades.json",
    shell:
        """
        python scripts/clades_renaming.py \
        --input-node-data {input} \
        --output-node-data {output.node_data}
        """


rule colors:
    input:
        ordering="resources/color_ordering.tsv",
        color_schemes="resources/color_schemes.tsv",
        metadata="results/{build_name}/metadata.tsv",
    output:
        colors="results/{build_name}/colors.tsv",
    shell:
        """
        python3 scripts/assign-colors.py \
            --ordering {input.ordering} \
            --color-schemes {input.color_schemes} \
            --output {output.colors} \
            --metadata {input.metadata} 2>&1
        """


rule nice_strain_names:
    input:
        metadata="results/{build_name}/metadata.tsv",
    output:
        metadata="results/{build_name}/metadata_display.tsv",
    shell:
        """
        python3 scripts/make_nice_strain_names.py \
            --metadata {input.metadata} \
            --input-columns country accession date \
            --separator '/' \
            --output-column strain_display \
            --output {output.metadata}
        """


rule patch_auspice_config:
    input:
        auspice_config="resources/auspice_config.json",
        patch_json="resources/{build_name}/auspice_config_patch.json",
    output:
        auspice_config="results/{build_name}/auspice_config.json",
    shell:
        """
        jsonpatch {input.auspice_config} {input.patch_json} > {output.auspice_config}
        """


rule export:
    # "Exporting data files for for auspice"
    input:
        colors="results/{build_name}/colors.tsv",
        tree="results/{build_name}/tree.nwk",
        metadata="results/{build_name}/metadata_display.tsv",
        branch_lengths="results/{build_name}/branch_lengths.json",
        clades="results/{build_name}/clades.json",
        nt_muts="results/{build_name}/nt_muts.json",
        auspice_config="results/{build_name}/auspice_config.json",
    output:
        auspice_json="auspice/nextclade_mpox_{build_name}_pretty.json",
        root_sequence="auspice/nextclade_mpox_{build_name}_pretty_root-sequence.json",
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --node-data {input.branch_lengths} {input.nt_muts} {input.clades} \
            --colors {input.colors} \
            --auspice-config {input.auspice_config} \
            --include-root-sequence \
            --metadata-id-columns accession \
            --output {output.auspice_json}
        """


rule minify_tree:
    input:
        auspice="auspice/nextclade_mpox_{build_name}_pretty.json",
    output:
        auspice="auspice/nextclade_mpox_{build_name}.json",
    shell:
        """
        gojq -c '.' {input.auspice} > {output.auspice}
        """


rule assemble_dataset:
    input:
        reference="resources/{build_name}/reference.fasta",
        tree="auspice/nextclade_mpox_{build_name}.json",
        pathogen_json="resources/{build_name}/pathogen.json",
        sequences="resources/{build_name}/sequences.fasta",
        annotation=annotation_path,
        readme="resources/{build_name}/README.md",
        changelog="resources/{build_name}/CHANGELOG.md",
    output:
        reference="datasets/{build_name}/reference.fasta",
        tree="datasets/{build_name}/tree.json",
        pathogen_json="datasets/{build_name}/pathogen.json",
        sequences="datasets/{build_name}/sequences.fasta",
        annotation="datasets/{build_name}/genome_annotation.gff3",
        readme="datasets/{build_name}/README.md",
        changelog="datasets/{build_name}/CHANGELOG.md",
    shell:
        """
        cp {input.reference} {output.reference}
        cp {input.tree} {output.tree}
        cp {input.pathogen_json} {output.pathogen_json}
        cp {input.sequences} {output.sequences}
        cp {input.annotation} {output.annotation}
        cp {input.readme} {output.readme}
        cp {input.changelog} {output.changelog}
        """


rule test_dataset:
    input:
        tree="datasets/{build_name}/tree.json",
        pathogen_json="datasets/{build_name}/pathogen.json",
        sequences="datasets/{build_name}/sequences.fasta",
        annotation="datasets/{build_name}/genome_annotation.gff3",
        readme="datasets/{build_name}/README.md",
        changelog="datasets/{build_name}/CHANGELOG.md",
    output:
        outdir=directory("test_output/{build_name}"),
    params:
        dataset_dir="datasets/{build_name}",
    shell:
        """
        nextclade3 run \
            {input.sequences} \
            --input-dataset {params.dataset_dir} \
            --output-all {output.outdir}
        """


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


rule cleanall:
    # Removing directories: {params}"
    params:
        "results",
        "auspice",
        "data",
        "logs",
        "bin",
    shell:
        "rm -rfv {params}"