Snakefile

import re
import json
import requests
from pathlib import Path

configfile: "config/config.yaml"


def read_fasta_IDs(path):
    ids = []
    with open(path) as f:
        for line in f:
            if line.startswith(">"):
                ids.append(line[1:].strip().split(" ")[0])
    return ids


def download_file(url, path):
    response = requests.get(url)
    with open(path, "wb") as fw:
        fw.write(response.content)


def copy_file(from_path, to_path, chunk_size=1024):
    with open(from_path, "rb") as f:
        with open(to_path, "wb") as fw:
            fw.write(f.read())


INPUT_DIR = Path(config["INPUT_DIR"])
OUTPUT_DIR = Path(config["OUTPUT_DIR"])
DATASETS = [path.stem for path in INPUT_DIR.glob("*.txt")]


rule all:
    input:
        expand(OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/clusters.csv", dataset=DATASETS, prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/summary.json", dataset=DATASETS, prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/fitness/{days}d/estimated_fitness.csv", dataset=DATASETS, days=config["FITNESS_PADDING_DAYS"], prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/timeline/timeline.svg", dataset=DATASETS, prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/age/vs_background.svg", dataset=DATASETS, prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"summary/[{prop}]_[{size}]/timeline/timeline.svg", prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"summary/[{prop}]_[{size}]/age/transmission_age.svg", prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report.svg", days=config["FITNESS_PADDING_DAYS"], prop=config["MIN_PROP"], size=config["MIN_SIZE"]),
        expand(OUTPUT_DIR/"summary/[{prop}]_[{size}]/missing/age.svg", prop=config["MIN_PROP"], size=config["MIN_SIZE"])


rule get_problematic_vcf:
    threads: 1
    shadow: "shallow"
    output:
        vcf_path = "data/problematic_sites_sarsCov2.vcf"
    resources:
        mem_mb = 2000,
        runtime = 15
    retries: 3
    run:
        if config["PROBLEMATIC_VCF"].startswith("https://"):
            download_file(config["PROBLEMATIC_VCF"], output.vcf_path)
        else:
            copy_file(config["PROBLEMATIC_VCF"], output.vcf_path)


rule get_reference_tree:
    threads: 1
    shadow: "shallow"
    output:
        tree_path = "data/public.all.masked.pb.gz"
    params:
        reference_tree = config["REFERENCE_TREE"]
    resources:
        mem_mb = 2000,
        runtime = 15
    retries: 3
    run:
        if params.reference_tree.startswith("https://"):
            download_file(params.reference_tree, output.tree_path)
        else:
            copy_file(params.reference_tree, output.tree_path)


rule extract_records:
    threads: 1
    shadow: "shallow"
    conda: "envs/bio.yaml"
    input:
        ids_file = INPUT_DIR/"{dataset}.txt"
    params:
        full_fasta = config["ALIGNED_FULL_FASTA"]
    output:
        fasta = OUTPUT_DIR/"{dataset}/sequences.aligned.fasta",
        extracted_ids = OUTPUT_DIR/"{dataset}/extracted_ids.csv",
        targets = OUTPUT_DIR/"{dataset}/targets.txt"
    resources:
        mem_mb = 4000,
        runtime = 60
    script:
        "scripts/extract_records.py"


rule phylogenetic_placement:
    threads: 8
    shadow: "shallow"
    conda: "envs/usher.yaml"
    input:
        new_samples = OUTPUT_DIR/"{dataset}/sequences.aligned.fasta",
        masking_vcf = "data/problematic_sites_sarsCov2.vcf",
        tree_protobuf = "data/public.all.masked.pb.gz"
    params:
        reference_fasta = config["REFERENCE_FASTA"],
        reference_id = read_fasta_IDs(config["REFERENCE_FASTA"])[0],
        max_equally_parsimonious_placements = 1000000
    output:
        tree = OUTPUT_DIR/"{dataset}/tree.nwk",
        clades = OUTPUT_DIR/"{dataset}/clades.txt",
        tree_pb = temp(OUTPUT_DIR/"{dataset}/tree.annotated.pb"),
        usher_dir = temp(directory(OUTPUT_DIR/"{dataset}/usher_results"))
    resources:
        mem_mb = 16000,
        runtime = 12 * 60
    shell:
        """
        cat {params.reference_fasta} {input.new_samples} > input.fasta
        faToVcf -maskSites='{input.masking_vcf}' -ref='{params.reference_id}' input.fasta input.vcf
        usher -T {threads} -i {input.tree_protobuf} -o {output.tree_pb} -v input.vcf -D -u -e {params.max_equally_parsimonious_placements} -d usher_results
        mv usher_results/uncondensed-final-tree.nh {output.tree}
        mv usher_results/clades.txt {output.clades}
        mv usher_results {output.usher_dir}
        """


rule build_phylo4:
    threads: 1
    shadow: "shallow"
    conda: "envs/trees.yaml"
    input:
        tree = OUTPUT_DIR/"{dataset}/tree.nwk",
        extracted_ids = OUTPUT_DIR/"{dataset}/extracted_ids.csv"
    output:
        tree_p4 = OUTPUT_DIR/"{dataset}/tree.p4.csv"
    resources:
        mem_mb = 8000,
        runtime = 24 * 60
    script:
        "scripts/build_phylo4.R"


rule calculate_clusters:
    threads: 1
    shadow: "shallow"
    conda: "envs/tcfinder.yaml"
    input:
        tree_p4 = OUTPUT_DIR/"{dataset}/tree.p4.csv",
        targets = OUTPUT_DIR/"{dataset}/targets.txt"
    output:
        clusters = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/clusters.csv"
    resources:
        mem_mb = lambda wildcards, attempt: 8000 * attempt,
        runtime = 12 * 60
    retries: 3
    shell: "tcfinder -i {input.tree_p4} -t {input.targets} -o {output.clusters} -p {wildcards.prop} -s {wildcards.size}"


rule summarize_results:
    threads: 1
    shadow: "shallow"
    conda: "envs/tables.yaml"
    input:
        clusters = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/clusters.csv",
        extracted_ids = OUTPUT_DIR/"{dataset}/extracted_ids.csv",
        ids_file = INPUT_DIR/"{dataset}.txt"
    output:
        table = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/summary.json"
    resources:
        mem_mb = 4000,
        runtime = 15
    script:
        "scripts/summarize_results.py"


rule build_haplotype_metadata:
    threads: 1
    conda: "envs/trees.yaml"
    input:
        clusters = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/clusters.csv",
        extracted_ids = OUTPUT_DIR/"{dataset}/extracted_ids.csv",
        full_metadata = config["METADATA"]
    params:
        metadata_originalid_column = config["METADATA_COLS"]["FULL_METADATA"]["ORIGINAL_ID"],
        metadata_id_column = config["METADATA_COLS"]["FULL_METADATA"]["UNIQUE_ID"],
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_location_column = config["METADATA_COLS"]["FULL_METADATA"]["LOCATION"],
        metadata_gender_column = config["METADATA_COLS"]["FULL_METADATA"]["GENDER"],
        metadata_lineage_column = config["METADATA_COLS"]["FULL_METADATA"]["LINEAGE"],
        metadata_age_column = config["METADATA_COLS"]["FULL_METADATA"]["AGE"]
    output:
        metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/metadata.csv"
    resources:
        mem_mb = 16000,
        runtime = 2 * 60
    script: "scripts/build_haplotype_metadata.R"


rule estimate_fitness:
    threads: 1
    input:
        full_metadata = config["METADATA"],
        haplotype_metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/metadata.csv"
    params:
        metadata_id_column = config["METADATA_COLS"]["FULL_METADATA"]["UNIQUE_ID"],
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_location_column = config["METADATA_COLS"]["FULL_METADATA"]["LOCATION"]
    output:
        fitness = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/fitness/{days_padding}d/estimated_fitness.csv"
    script: "scripts/estimate_fitness.py"


rule report_estimated_fitness:
    threads: 1
    conda: "envs/report.yaml"
    input:
        fitnesses = expand(OUTPUT_DIR/"{dataset}/[{{prop}}]_[{{size}}]/fitness/{{days}}d/estimated_fitness.csv", dataset=DATASETS)
    params:
        width_per_haplotype_mm = 50, height_mm = 240,
        separate_location = True
    output:
        report = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report.svg",
        report_simple = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report.simple.svg",
        report_data = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report.csv",
        report_adding = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report_adding.svg",
        report_slicing = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report_slicing.svg",
        summary = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/summary.csv",
        report_contrast = OUTPUT_DIR/"summary/[{prop}]_[{size}]/fitness/{days}d/report_contrast.csv"
    script: "scripts/report_estimated_fitness.R"


rule report_haplotype_timeline:
    threads: 1
    conda: "envs/report.yaml"
    input:
        haplotype_metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/metadata.csv"
    params:
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_location_column = config["METADATA_COLS"]["FULL_METADATA"]["LOCATION"],
        separate_location = config["METADATA_COLS"]["SEPARATE_LOCATION"],
        n_top_countries = 5,
        width_mm = 320, height_mm = 200
    output:
        report = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/timeline/timeline.svg",
        report_data = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/timeline/timeline.csv"
    script: "scripts/report_timeline.R"


rule report_summarized_timeline:
    threads: 1
    conda: "envs/report.yaml"
    input:
        haplotype_metadata_tables = expand(OUTPUT_DIR/"{dataset}/[{{prop}}]_[{{size}}]/metadata.csv", dataset=DATASETS)
    params:
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_location_column = config["METADATA_COLS"]["FULL_METADATA"]["LOCATION"],
        separate_location = config["METADATA_COLS"]["SEPARATE_LOCATION"],
        n_top_countries = 11,
        width_mm = 320, height_per_haplotype_mm = 50
    output:
        report = OUTPUT_DIR/"summary/[{prop}]_[{size}]/timeline/timeline.svg",
        report_data = OUTPUT_DIR/"summary/[{prop}]_[{size}]/timeline/timeline.csv"
    script: "scripts/report_summarized_timeline.R"


rule build_age_corrected_haplotype_metadata:
    threads: 1
    conda: "envs/report.yaml"
    input:
        haplotype_metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/metadata.csv"
    params:
        metadata_age_column = config["METADATA_COLS"]["FULL_METADATA"]["AGE"],
        cluster_size_threshold = lambda wildcards: wildcards.size,
        min_age = 0,
        max_age = 120  # Maria Branyas was 117 at the time
    output:
        age_corrected_metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/age/metadata.csv"
    script: "scripts/build_age_corrected_metadata.R"


rule report_age_differences:
    threads: 1
    conda: "envs/report.yaml"
    input:
        age_corrected_tables = expand(OUTPUT_DIR/"{dataset}/[{{prop}}]_[{{size}}]/age/metadata.csv", dataset=DATASETS)
    params:
        metadata_age_column = config["METADATA_COLS"]["FULL_METADATA"]["AGE"],
        cluster_size_threshold = lambda wildcards: wildcards.size,
        width_per_haplotype_mm = 50, height_mm = 300
    output:
        report_haplotype = OUTPUT_DIR/"summary/[{prop}]_[{size}]/age/haplotype_age.svg",
        report_haplotype_contrast = OUTPUT_DIR/"summary/[{prop}]_[{size}]/age/haplotype_age.contrast.csv",
        report_haplotype_simple = OUTPUT_DIR/"summary/[{prop}]_[{size}]/age/haplotype_age.simple.svg",
        report_transmission = OUTPUT_DIR/"summary/[{prop}]_[{size}]/age/transmission_age.svg"
    script: "scripts/report_age_differences.R"


rule report_age_vs_background:
    threads: 1
    conda: "envs/report.yaml"
    input:
        age_corrected_metadata = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/age/metadata.csv",
        full_metadata = config["METADATA"]
    params:
        metadata_age_column = config["METADATA_COLS"]["FULL_METADATA"]["AGE"],
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_id_column = config["METADATA_COLS"]["FULL_METADATA"]["UNIQUE_ID"],
        min_age = rules.build_age_corrected_haplotype_metadata.params.min_age,
        max_age = rules.build_age_corrected_haplotype_metadata.params.max_age,
        width_mm = 320, height_mm = 200
    output:
        report = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/age/vs_background.svg",
        report_data = OUTPUT_DIR/"{dataset}/[{prop}]_[{size}]/age/vs_background.csv"
    script: "scripts/report_age_vs_background.R"


rule report_missing_data:
    threads: 1
    conda: "envs/report.yaml"
    input:
        age_corrected_tables = expand(OUTPUT_DIR/"{dataset}/[{{prop}}]_[{{size}}]/age/metadata.csv", dataset=DATASETS)
    params:
        metadata_age_column = config["METADATA_COLS"]["FULL_METADATA"]["AGE"],
        metadata_date_column = config["METADATA_COLS"]["FULL_METADATA"]["DATE"],
        metadata_location_column = config["METADATA_COLS"]["FULL_METADATA"]["LOCATION"],
        width_mm = 320, height_mm = 200
    output:
        report_age = OUTPUT_DIR/"summary/[{prop}]_[{size}]/missing/age.svg",
        report_location = OUTPUT_DIR/"summary/[{prop}]_[{size}]/missing/location.svg",
        report_date = OUTPUT_DIR/"summary/[{prop}]_[{size}]/missing/date.svg"
    script: "scripts/report_missing_data.R"