Snakefile_base

import datetime
from datetime import date
import pandas as pd
from treetime.utils import numeric_date
from scripts.flu_regions import region_names

path_to_fauna = '../fauna'
min_length = 900
frequency_regions = region_names


localrules: download_titers, download_sequences

def vpm(v):
    vpm = {'6m':360, '2y':90, '3y':60, '6y':30, '12y':15, '60y':5}
    return vpm[v.resolution] if v.resolution in vpm else 5

def reference_strain(v):
    references = {'h3n2':"A/Beijing/32/1992",
                  'h1n1pdm':"A/California/07/2009",
                  'vic':"B/HongKong/02/1993",
                  'yam':"B/Singapore/11/1994"
                  }
    return references[v.lineage]

genes_to_translate = {'ha':['SigPep', 'HA1', 'HA2'], 'na':['NA'],
                      'pb1':['PB1'], 'pb2':['PB2'], 'pa':['PA'],
                      'np':['NP'], 'ma':['M'], 'ns':['NEP']}

def gene_names(w):
    return genes_to_translate[w.segment]

def translations(w):
    genes = gene_names(w)
    return ["results/aa-seq_%s_%s_%s_%s_%s_%s_%s.fasta"%(w.center, w.lineage, w.segment, w.resolution, w.passage, w.assay, g)
            for g in genes]

def pivot_interval(w):
    """Returns the number of months between pivots by build resolution.
    """
    pivot_intervals_by_resolution = {'6m': 1, '2y': 1, '3y': 2, '6y': 3, '12y': 6, '60y': 6}
    return pivot_intervals_by_resolution[w.resolution]

def min_date(wildcards):
    now = numeric_date(date.today())
    if wildcards.resolution[-1] == "y":
        years_back = int(wildcards.resolution[:-1])
    elif wildcards.resolution[-1] == "m":
        years_back = int(wildcards.resolution[:-1]) / 12.
    else:
        years_back = 3
    return now - years_back

def max_date(w):
    # Estimate frequencies a given number of months back in the past to account
    # for lag in data availability.
    if "frequency_max_date_month_offset" in config:
        date_offset = pd.DateOffset(months=config["frequency_max_date_month_offset"])
        return numeric_date(pd.to_datetime(date.today()) - date_offset)
    else:
        return numeric_date(date.today())

def clock_rate(w):
    # these rates are from 12y runs on 2019-10-18
    rate = {
     ('h1n1pdm', 'ha'): 0.00329,
 	 ('h1n1pdm', 'na'): 0.00326,
 	 ('h1n1pdm', 'np'): 0.00221,
 	 ('h1n1pdm', 'pa'): 0.00217,
	 ('h1n1pdm', 'pb1'): 0.00205,
 	 ('h1n1pdm', 'pb2'): 0.00277,
 	 ('h3n2', 'ha'): 0.00382,
 	 ('h3n2', 'na'): 0.00267,
	 ('h3n2', 'np'): 0.00157,
 	 ('h3n2', 'pa'): 0.00178,
 	 ('h3n2', 'pb1'): 0.00139,
 	 ('h3n2', 'pb2'): 0.00218,
 	 ('vic', 'ha'): 0.00145,
 	 ('vic', 'na'): 0.00133,
 	 ('vic', 'np'): 0.00132,
 	 ('vic', 'pa'): 0.00178,
 	 ('vic', 'pb1'): 0.00114,
 	 ('vic', 'pb2'): 0.00106,
 	 ('yam', 'ha'): 0.00176,
 	 ('yam', 'na'): 0.00177,
 	 ('yam', 'np'): 0.00133,
 	 ('yam', 'pa'): 0.00112,
 	 ('yam', 'pb1'): 0.00092,
 	 ('yam', 'pb2'): 0.00113}
    return rate.get((w.lineage, w.segment), 0.001)


def clock_std_dev(w):
    return 0.2*clock_rate(w)

#
# Define clades functions
#
def _get_clades_file_for_wildcards(wildcards):
    if wildcards.segment == "ha":
        return "config/clades_%s_ha.tsv"%(wildcards.lineage)
    else:
        return "results/clades_%s_%s_ha_%s_%s_%s.json"%(wildcards.center, wildcards.lineage,
                                                        wildcards.resolution, wildcards.passage, wildcards.assay)


#
# Define titer data sets to be used.
#
def _get_tdb_databases(wildcards):
    if wildcards.center in ['cdc', 'crick', 'niid', 'vidrl']:
        return wildcards.center + "_tdb tdb"
    else:
        return "cdc_tdb crick_tdb niid_tdb vidrl_tdb tdb"


def exclude_where(wildcards):
    if wildcards.passage == 'cell':
        return "country=? region=? passage=egg"
    else:
        return "country=? region=?"

#
# Define LBI parameters and functions.
#
LBI_params = {
    '6m': {"tau": 0.3, "time_window": 0.5},
    '2y': {"tau": 0.3, "time_window": 0.5},
    '3y': {"tau": 0.4, "time_window": 0.6},
    '6y': {"tau": 0.25, "time_window": 0.75},
    '12y': {"tau": 0.25, "time_window": 0.75},
    '60y': {"tau": 0.25, "time_window": 0.75}
}

def _get_lbi_tau_for_wildcards(wildcards):
    return LBI_params[wildcards.resolution]["tau"]

def _get_lbi_window_for_wildcards(wildcards):
    return LBI_params[wildcards.resolution]["time_window"]

#
# Configure distance maps (for amino acid and other distances).
#

# Load distance map configuration for lineages and segments.
distance_map_config = pd.read_table("config/distance_maps.tsv")

def _get_build_distance_map_config(wildcards):
    config = distance_map_config[(distance_map_config["lineage"] == wildcards.lineage) &
                          (distance_map_config["segment"] == wildcards.segment)]
    if config.shape[0] > 0:
        return config
    else:
        return None

def _get_distance_comparisons_by_lineage_and_segment(wildcards):
    config = _get_build_distance_map_config(wildcards)
    return " ".join(config.loc[:, "compare_to"].values)

def _get_distance_attributes_by_lineage_and_segment(wildcards):
    config = _get_build_distance_map_config(wildcards)
    return " ".join(config.loc[:, "attribute"].values)

def _get_distance_maps_by_lineage_and_segment(wildcards):
    config = _get_build_distance_map_config(wildcards)
    return [
        "config/distance_maps/{wildcards.lineage}/{wildcards.segment}/{distance_map}.json".format(wildcards=wildcards, distance_map=distance_map)
        for distance_map in config.loc[:, "distance_map"].values
    ]

def _get_glyc_alignment(w):
    return "results/aa-seq_{c}_{l}_{seg}_{res}_{p}_{a}_{gene}.fasta".format(
            c=w.center, l=w.lineage, seg=w.segment, res=w.resolution, p=w.passage, a=w.assay,
            gene='HA1' if w.segment=='ha' else 'NA')

#
# Define node data table functions.
#

def float_to_datestring(time):
    """Convert a floating point date from TreeTime `numeric_date` to a date string
    """
    # Extract the year and remainder from the floating point date.
    year = int(time)
    remainder = time - year

    # Calculate the day of the year (out of 365 + 0.25 for leap years).
    tm_yday = int(remainder * 365.25)
    if tm_yday == 0:
        tm_yday = 1

    # Construct a date object from the year and day of the year.
    date = datetime.datetime.strptime("%s-%s" % (year, tm_yday), "%Y-%j")

    # Build the date string with zero-padded months and days.
    date_string = "%s-%.2i-%.2i" % (date.year, date.month, date.day)

    return date_string

def _get_start_timepoint(wildcards):
    return float_to_datestring(min_date(wildcards))

def _get_end_timepoint(wildcards):
    return float_to_datestring(max_date(wildcards))

def _get_annotations_for_node_data(wildcards):
    annotations = ["%s=%s" % (key, value) for key, value in wildcards.items()]
    annotations.append("timepoint=%s" % _get_end_timepoint(wildcards))
    return " ".join(annotations)

def _get_excluded_fields_arg(wildcards):
    if config.get("excluded_node_data_fields"):
        return "--excluded-fields %s" % " ".join(config["excluded_node_data_fields"])
    else:
        return ""

#
# Define forecasting helper functions.
#

def _get_delta_months_to_forecast(wildcards):
    return " ".join([str(month) for month in config["fitness_model"]["delta_months"]])

#
# Define shared rules
#

def _get_auspice_config(wildcards):
    if wildcards.lineage == "h3n2" and wildcards.segment == "ha" and wildcards.resolution == "2y":
        return "config/auspice_config_h3n2_fitness.json"
    else:
        return "config/auspice_config_{lineage}.json".format(lineage=wildcards.lineage)

rule files:
    params:
        outliers = "config/outliers_{lineage}.txt",
        exclude_sites = "config/exclude-sites_{lineage}.txt",
        references = "config/references_{lineage}.txt",
        reference = "config/reference_{lineage}_{segment}.gb",
        colors = "config/colors.tsv",
        auspice_config = _get_auspice_config,
        vaccine_json = "config/vaccines_{lineage}.json",
        description = "config/description.md"

files = rules.files.params


rule download_sequences:
    message: "Downloading sequences from fauna"
    output:
        sequences = "data/{lineage}_{segment}.fasta"
    params:
        fasta_fields = "strain virus accession collection_date virus_inclusion_date region country division location passage_category originating_lab submitting_lab age gender"
    shell:
        """
        python3 {path_to_fauna}/vdb/download.py \
            --database vdb \
            --virus flu \
            --fasta_fields {params.fasta_fields} \
            --resolve_method split_passage \
            --select locus:{wildcards.segment} lineage:seasonal_{wildcards.lineage} \
            --path data \
            --fstem {wildcards.lineage}_{wildcards.segment}
        """

rule download_titers:
    message: "Downloading titers from fauna: {wildcards.lineage}, {wildcards.assay}, {wildcards.center}"
    output:
        titers = "data/{center}_{lineage}_{passage}_{assay}_titers.tsv"
    params:
        dbs = _get_tdb_databases,
        assays = lambda wildcards: "fra,hint" if wildcards.assay == "fra" else wildcards.assay
    shell:
        """
        python3 {path_to_fauna}/tdb/download.py \
            --database {params.dbs} \
            --virus flu \
            --subtype {wildcards.lineage} \
            --select assay_type:{params.assays} serum_passage_category:{wildcards.passage} \
            --path data \
            --fstem {wildcards.center}_{wildcards.lineage}_{wildcards.passage}_{wildcards.assay}
        """

rule parse:
    message: "Parsing fasta into sequences and metadata"
    input:
        sequences = rules.download_sequences.output.sequences
    output:
        sequences = "results/sequences_{lineage}_{segment}.fasta",
        metadata = "results/metadata_{lineage}_{segment}.tsv"
    params:
        fasta_fields =  "strain virus accession date date_submitted region country division location passage originating_lab submitting_lab age gender",
        prettify_fields = "region country division location originating_lab submitting_lab"
    shell:
        """
        augur parse \
            --sequences {input.sequences} \
            --output-sequences {output.sequences} \
            --output-metadata {output.metadata} \
            --fields {params.fasta_fields} \
            --prettify-fields {params.prettify_fields}
        """

rule filter:
    message:
        """
        Filtering {wildcards.lineage} {wildcards.segment} {wildcards.passage} sequences:
          - less than {params.min_length} bases
          - outliers
          - samples with missing region and country metadata
          - samples that are egg-passaged if cell build
        """
    input:
        metadata = rules.parse.output.metadata,
        sequences = rules.parse.output.sequences,
        exclude = files.outliers
    output:
        sequences = 'results/filtered_{lineage}_{segment}_{passage}.fasta'
    params:
        min_length = min_length,
        exclude_where = exclude_where
    shell:
        """
        augur filter \
            --sequences {input.sequences} \
            --metadata {input.metadata} \
            --min-length {params.min_length} \
            --non-nucleotide \
            --exclude {input.exclude} \
            --exclude-where {params.exclude_where} \
            --output {output}
        """

rule select_strains:
    input:
        sequences = expand("results/filtered_{{lineage}}_{segment}_{{passage}}.fasta", segment=segments),
        metadata = expand("results/metadata_{{lineage}}_{segment}.tsv", segment=segments),
        titers = rules.download_titers.output.titers,
        include = files.references
    output:
        strains = "results/strains_{center}_{lineage}_{resolution}_{passage}_{assay}.txt",
    params:
        viruses_per_month = vpm
    shell:
        """
        python3 scripts/select_strains.py \
            --sequences {input.sequences} \
            --metadata {input.metadata} \
            --segments {segments} \
            --include {input.include} \
            --lineage {wildcards.lineage} \
            --resolution {wildcards.resolution} \
            --viruses-per-month {params.viruses_per_month} \
            --titers {input.titers} \
            --output {output.strains}
        """

rule extract:
    input:
        sequences = rules.filter.output.sequences,
        strains = rules.select_strains.output.strains
    output:
        sequences = 'results/extracted_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.fasta'
    shell:
        """
        python3 scripts/extract_sequences.py \
            --sequences {input.sequences} \
            --samples {input.strains} \
            --output {output}
        """

rule annotate_recency_of_submissions:
    input:
        metadata = rules.parse.output.metadata,
    params:
        submission_date_field=config["submission_date_field"],
        date_bins=config["recency"]["date_bins"],
        date_bin_labels=config["recency"]["date_bin_labels"],
        upper_bin_label=config["recency"]["upper_bin_label"],
    output:
        node_data = "results/recency_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        python3 scripts/construct-recency-from-submission-date.py \
            --metadata {input.metadata} \
            --submission-date-field {params.submission_date_field} \
            --date-bins {params.date_bins} \
            --date-bin-labels {params.date_bin_labels:q} \
            --upper-bin-label {params.upper_bin_label} \
            --output {output.node_data}
        """

rule align:
    message:
        """
        Aligning sequences to {input.reference}
        """
    input:
        sequences = rules.extract.output.sequences,
        reference = files.reference
    output:
        alignment = "results/aligned_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.fasta"
    shell:
        """
        python3 scripts/codon_align.py \
            --sequences {input.sequences} \
            --reference {input.reference} \
            --output {output.alignment}
        """

rule tree:
    message: "Building tree"
    input:
        alignment = rules.align.output.alignment,
        exclude_sites = files.exclude_sites
    output:
        tree = "results/tree-raw_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.nwk"
    shell:
        """
        augur tree \
            --alignment {input.alignment} \
            --output {output.tree} \
            --nthreads 1 \
            --exclude-sites {input.exclude_sites}
        """

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 = rules.parse.output.metadata
    output:
        tree = "results/tree_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.nwk",
        node_data = "results/branch-lengths_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        coalescent = "const",
        date_inference = "marginal",
        clock_filter_iqd = 4,
        clock_rate = clock_rate,
        clock_std_dev = clock_std_dev
    shell:
        """
        augur refine \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --metadata {input.metadata} \
            --output-tree {output.tree} \
            --output-node-data {output.node_data} \
            --timetree \
            --no-covariance \
            --clock-rate {params.clock_rate} \
            --clock-std-dev {params.clock_std_dev} \
            --coalescent {params.coalescent} \
            --date-confidence \
            --date-inference {params.date_inference} \
            --clock-filter-iqd {params.clock_filter_iqd}
        """

rule ancestral:
    message: "Reconstructing ancestral sequences and mutations"
    input:
        tree = rules.refine.output.tree,
        alignment = rules.align.output
    output:
        node_data = "results/nt-muts_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        inference = "joint"
    shell:
        """
        augur ancestral \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --output-node-data {output.node_data} \
            --inference {params.inference}
        """

rule translate:
    message: "Translating amino acid sequences"
    input:
        tree = rules.refine.output.tree,
        node_data = rules.ancestral.output.node_data,
        reference = files.reference
    output:
        node_data = "results/aa-muts_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        augur translate \
            --tree {input.tree} \
            --ancestral-sequences {input.node_data} \
            --reference-sequence {input.reference} \
            --output {output.node_data} \
        """

rule reconstruct_translations:
    message: "Reconstructing translations required for titer models and frequencies"
    input:
        tree = rules.refine.output.tree,
        node_data = "results/aa-muts_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    output:
        aa_alignment = "results/aa-seq_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_{gene}.fasta"
    shell:
        """
        augur reconstruct-sequences \
            --tree {input.tree} \
            --mutations {input.node_data} \
            --gene {wildcards.gene} \
            --output {output.aa_alignment} \
            --internal-nodes
        """


rule convert_translations_to_json:
    input:
        tree = rules.refine.output.tree,
        translations = translations
    output:
        translations = "results/aa-seq_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        gene_names = gene_names
    shell:
        """
        python3 flu-forecasting/scripts/convert_translations_to_json.py \
            --tree {input.tree} \
            --alignment {input.translations} \
            --gene-names {params.gene_names} \
            --output {output.translations}
        """


rule traits:
    message:
        """
        Inferring ancestral traits for {params.columns!s}
        """
    input:
        tree = rules.refine.output.tree,
        metadata = rules.parse.output.metadata
    output:
        node_data = "results/traits_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    params:
        columns = "region"
    shell:
        """
        augur traits \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --output {output.node_data} \
            --columns {params.columns} \
            --confidence
        """

rule titers_sub:
    input:
        titers = rules.download_titers.output.titers,
        aa_muts = rules.translate.output,
        alignments = translations,
        tree = rules.refine.output.tree
    params:
        genes = gene_names
    output:
        titers_model = "results/titers-sub-model_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        augur titers sub \
            --titers {input.titers} \
            --alignment {input.alignments} \
            --gene-names {params.genes} \
            --tree {input.tree} \
            --allow-empty-model \
            --output {output.titers_model}
        """

rule titers_tree:
    input:
        titers = rules.download_titers.output.titers,
        tree = rules.refine.output.tree
    output:
        titers_model = "results/titers-tree-model_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        augur titers tree \
            --titers {input.titers} \
            --tree {input.tree} \
            --allow-empty-model \
            --output {output.titers_model}
        """

rule tip_frequencies:
    input:
        tree = rules.refine.output.tree,
        metadata = rules.parse.output.metadata,
        weights = "config/frequency_weights_by_region.json"
    params:
        narrow_bandwidth = 2 / 12.0,
        wide_bandwidth = 3 / 12.0,
        proportion_wide = 0.0,
        weight_attribute = "region",
        min_date = min_date,
        max_date = max_date,
        pivot_interval = pivot_interval
    output:
        tip_freq = "auspice/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_tip-frequencies.json"
    shell:
        """
        augur frequencies \
            --method kde \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --narrow-bandwidth {params.narrow_bandwidth} \
            --wide-bandwidth {params.wide_bandwidth} \
            --proportion-wide {params.proportion_wide} \
            --weights {input.weights} \
            --weights-attribute {params.weight_attribute} \
            --pivot-interval {params.pivot_interval} \
            --min-date {params.min_date} \
            --max-date {params.max_date} \
            --output {output}
        """

rule tree_frequencies:
    input:
        tree = rules.refine.output.tree,
        metadata = rules.parse.output.metadata,
    params:
        min_date = min_date,
        max_date = max_date,
        pivot_interval = pivot_interval,
        regions = ['global'] + frequency_regions,
        min_clade = 20
    output:
        frequencies = "results/tree-frequencies_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        augur frequencies \
            --method diffusion \
            --include-internal-nodes \
            --tree {input.tree} \
            --regions {params.regions:q} \
            --metadata {input.metadata} \
            --pivot-interval {params.pivot_interval} \
            --minimal-clade-size {params.min_clade} \
            --min-date {params.min_date} \
            --max-date {params.max_date} \
            --output {output}
        """

rule diffusion_frequencies:
    input:
        tree = rules.refine.output.tree,
        metadata = rules.parse.output.metadata,
    params:
        min_date = min_date,
        max_date = max_date,
        pivot_interval = config["fitness_model"]["pivot_interval"],
        regions = 'global'
    output:
        frequencies = "results/diffusion-frequencies_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json",
    shell:
        """
        augur frequencies \
            --method diffusion \
            --include-internal-nodes \
            --tree {input.tree} \
            --regions {params.regions:q} \
            --metadata {input.metadata} \
            --pivot-interval {params.pivot_interval} \
            --min-date {params.min_date} \
            --max-date {params.max_date} \
            --output {output}
        """

rule delta_frequency:
    input:
        tree = rules.refine.output.tree,
        frequencies = rules.diffusion_frequencies.output.frequencies
    output:
        delta_frequency = "results/delta_frequency_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        delta_pivots = config["fitness_model"]["delta_pivots"],
        method = "diffusion"
    shell:
        """
        python3 flu-forecasting/scripts/calculate_delta_frequency.py \
            --tree {input.tree} \
            --frequencies {input.frequencies} \
            --frequency-method {params.method} \
            --delta-pivots {params.delta_pivots} \
            --output {output.delta_frequency}
        """

rule clades:
    message: "Annotating clades"
    input:
        tree = "results/tree_{center}_{lineage}_ha_{resolution}_{passage}_{assay}.nwk",
        nt_muts = rules.ancestral.output,
        aa_muts = rules.translate.output,
        clades = _get_clades_file_for_wildcards
    output:
        clades = "results/clades_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    run:
        if wildcards.segment == 'ha':
            shell("""
                augur clades \
                    --tree {input.tree} \
                    --mutations {input.nt_muts} {input.aa_muts} \
                    --clades {input.clades} \
                    --output {output.clades}
            """)
        else:
            shell("""
                python3 scripts/import_tip_clades.py \
                    --tree {input.tree} \
                    --clades {input.clades} \
                    --output {output.clades}
            """)

rule distances:
    input:
        tree = rules.refine.output.tree,
        alignments = translations,
        distance_maps = _get_distance_maps_by_lineage_and_segment
    params:
        genes = gene_names,
        comparisons = _get_distance_comparisons_by_lineage_and_segment,
        attribute_names = _get_distance_attributes_by_lineage_and_segment
    output:
        distances = "results/distances_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    shell:
        """
        augur distance \
            --tree {input.tree} \
            --alignment {input.alignments} \
            --gene-names {params.genes} \
            --compare-to {params.comparisons} \
            --attribute-name {params.attribute_names} \
            --map {input.distance_maps} \
            --output {output}
        """

rule pairwise_titer_tree_distances:
    input:
        tree = rules.refine.output.tree,
        frequencies = rules.tip_frequencies.output.tip_freq,
        model = rules.titers_tree.output.titers_model,
        date_annotations = rules.refine.output.node_data
    output:
        distances = "results/pairwise-titer-tree-distances_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        attribute_names = "cTiter_pairwise",
        months_back_for_current_samples = config["fitness_model"]["months_back_for_current_samples"],
        years_back_to_compare = config["fitness_model"]["max_years_for_distances"]
    shell:
        """
        python3 flu-forecasting/scripts/pairwise_titer_tree_distances.py \
            --tree {input.tree} \
            --frequencies {input.frequencies} \
            --model {input.model} \
            --attribute-name {params.attribute_names} \
            --date-annotations {input.date_annotations} \
            --months-back-for-current-samples {params.months_back_for_current_samples} \
            --years-back-to-compare {params.years_back_to_compare} \
            --output {output}
        """

rule titer_tree_cross_immunities:
    input:
        frequencies = rules.tip_frequencies.output.tip_freq,
        distances = rules.pairwise_titer_tree_distances.output.distances,
        date_annotations = rules.refine.output.node_data
    output:
        cross_immunities = "results/titer-tree-cross-immunity_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    params:
        distance_attributes = "cTiter_pairwise",
        immunity_attributes = "cTiter_x",
        decay_factors = "14.0",
        years_to_wane = config["fitness_model"]["max_years_for_distances"]
    shell:
        """
        python3 flu-forecasting/src/cross_immunity.py \
            --frequencies {input.frequencies} \
            --distances {input.distances} \
            --date-annotations {input.date_annotations} \
            --distance-attributes {params.distance_attributes} \
            --immunity-attributes {params.immunity_attributes} \
            --decay-factors {params.decay_factors} \
            --output {output}
        """

rule glyc:
    input:
        tree = rules.refine.output.tree,
        alignment = _get_glyc_alignment
    output:
        glyc = "results/glyc_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    shell:
        """
        python3 scripts/glyc.py \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --output {output}
        """

rule lbi:
    message: "Calculating LBI"
    input:
        tree = rules.refine.output.tree,
        branch_lengths = rules.refine.output.node_data
    params:
        tau = _get_lbi_tau_for_wildcards,
        window = _get_lbi_window_for_wildcards,
        names = "lbi"
    output:
        lbi = "results/lbi_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    shell:
        """
        augur lbi \
            --tree {input.tree} \
            --branch-lengths {input.branch_lengths} \
            --output {output} \
            --attribute-names {params.names} \
            --tau {params.tau} \
            --window {params.window}
        """