Snakefile_Priorities

# Does first few steps of a region build, then pastes the priorities
# calculated for each region into a JSON so that they can be viewed in Auspice.

include: "Snakefile"

ruleorder: export_priorities > export

REGIONS = ["_africa", "_asia", "_europe", "_north-america", "_oceania", "_south-america"]

# I recognise some rules here are a little redundant - happy to change file
# names so that they don't conflict with Region builds, but wanted this
# also to be stand-alone, so you don't have to have run region builds for
# this to work. And - can't include global build in list of regions!

rule priorities_for_all:
    input:
        expand("results/subsample_focus{region}.fasta", region=REGIONS),
        expand("results/subsampling_priorities{region}.tsv", region=REGIONS),
        "results/metadata_priorities.tsv",
        "auspice/ncov-with-priorities.json"

rule subsample_focus:
    message:
        """
        Subsample all sequences into a focal set
        """
    input:
        sequences = rules.mask.output.alignment,
        metadata = rules.download.output.metadata,
        include = files.include
    output:
        sequences = "results/subsample_focus{region}.fasta"
    params:
        group_by = "division year month",
        seq_per_group_global = 150, # i.e. if not regional build
        seq_per_group_regional = 300
    shell:
        """
        #Figure out what region being wanted
        rgn="{wildcards.region}"

        if [ -z "$rgn" ]; then
            seq_per_group={params.seq_per_group_global}
            regionarg="--exclude-where region="
            region="frog"   #I don't know! It wouldn't work without something!
            echo "Filtering for a global run - $seq_per_group per division"
        else
            seq_per_group={params.seq_per_group_regional}
            region="${{rgn//[_y]/}}"
            region="${{region//[-y]/ }}"
            echo "Filtering for a focal run on $region - $seq_per_group per division"
            regionarg="--exclude-where region!="
            echo "   This is passed to augur as $regionarg'$region'"
        fi

        augur filter \
            --sequences {input.sequences} \
            --metadata {input.metadata} \
            --include {input.include} \
            $regionarg"$region" \
            --group-by {params.group_by} \
            --sequences-per-group $seq_per_group \
            --output {output.sequences} \

        """

rule make_priorities:
    message:
        """
        determine priority for inclusion in as phylogenetic context by
        genetic similiarity to sequences in focal set.
        """
    input:
        alignment = rules.mask.output.alignment,
        metadata = rules.download.output.metadata,
        focal_alignment = rules.subsample_focus.output.sequences
    output:
        priorities = "results/subsampling_priorities{region}.tsv"
    shell:
        """
        #Figure out what region being wanted
        rgn="{wildcards.region}"

        if [ -z "$rgn" ]; then
            echo "Global run - no priorities needed"
            echo -n >{output.priorities}
        else
            region="${{rgn//[_y]/}}"
            region="${{region//[-y]/ }}"
            echo "Creating priorities for focal build on $region"
            python3 scripts/priorities.py --alignment {input.alignment} \
                            --metadata {input.metadata} \
                            --focal-alignment {input.focal_alignment} \
                            --output {output.priorities}
        fi
        """

rule add_priorities_metadata:
    input:
        metadata = rules.download.output.metadata,
        priorities = expand("results/subsampling_priorities{region}.tsv", region=REGIONS),
        config = "config/auspice_config.json"
    output:
        metadata = "results/metadata_priorities.tsv",
        config = "results/auspice_config_priors.json"
    shell:
        """
        python3 scripts/add_priorities_to_meta.py \
            --metadata {input.metadata} \
            --config {input.config} \
            --priorities {input.priorities} \
            --output-meta {output.metadata} \
            --output-config {output.config}
        """

rule export_priorities:
    message: "Exporting extra JSON with the priorities calculations"
    input:
        tree = rules.refine.output.tree,
        metadata = rules.add_priorities_metadata.output.metadata,
        branch_lengths = rules.refine.output.node_data,
        nt_muts = rules.ancestral.output.node_data,
        aa_muts = rules.translate.output.node_data,
        traits = rules.traits.output.node_data,
        auspice_config = rules.add_priorities_metadata.output.config,
        colors = rules.colors.output.colors,
        lat_longs = files.lat_longs,
        description = files.description,
        clades = "results/clades{region}.json",
        recency = rules.recency.output
    output:
        auspice_json = "auspice/ncov{region}-with-priorities.json"
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} {input.traits} {input.clades} {input.recency} \
            --auspice-config {input.auspice_config} \
            --colors {input.colors} \
            --lat-longs {input.lat_longs} \
            --description {input.description} \
            --output {output.auspice_json}
        """