Support fitting titer models per reference virus

Allow titer collections to request a titer model run per reference virus
in each collection. The workflow aggregates the resulting model results
into a single measurements panel JSON which allows users to display
inferred titer measurements per virus in the tree in both the
measurements panel and in the tree coloring by measurements for a
reference selected from the panel. This logic attempts to recreate a
nextflu feature which allows users to click on a titer reference virus
to color the tree by the measurements against that virus and then choose
to color by the titer model fit to that virus.

Related to #214

Author: huddlej

profiles/full-trees.yaml

@@ -116,6 +116,7 @@ builds:
           prefix: cell_hi_
           title: "Cell-passaged HI titers from ferret sera"
           genes: ["HA1"]
+          run_reference_models: true
         - name: egg_hi
           data: "data/h1n1pdm/who_ferret_egg_hi_titers.tsv"
           prefix: egg_hi_

scripts/generate_collection_config_json.py

@@ -71,23 +71,26 @@ def calc_row(clade):
     }
 
     # Read collection.
-    collection_df = pd.read_csv(args.collection, sep="\t", usecols=args.groupings + ["reference_date", "clade_reference"])
+    collection_df = pd.read_csv(args.collection, sep="\t", usecols=args.groupings)
 
     # Map y-axis positions in the phylogeny to reference strains.
     collection_df["y_axis_position_in_phylogeny"] = collection_df["reference_strain"].map(y_axis_positions_per_tip_name)
 
-    # Find minimum y-axis position for reference strains within each clade. This
-    # position represents the earliest instance of the clade in the tree.
-    min_y_axis_position_by_reference_clade = collection_df.groupby("subclade_reference")["y_axis_position_in_phylogeny"].min().reset_index().rename(
-        columns={"y_axis_position_in_phylogeny": "min_y_axis_position_in_phylogeny"}
-    )
-
-    # Annotate min y-axis position per clade to collection.
-    collection_df = collection_df.merge(
-        min_y_axis_position_by_reference_clade,
-        on="subclade_reference",
-        how="left",
-    )
+    if "subclade_reference" in collection_df.columns:
+        # Find minimum y-axis position for reference strains within each clade. This
+        # position represents the earliest instance of the clade in the tree.
+        min_y_axis_position_by_reference_clade = collection_df.groupby("subclade_reference")["y_axis_position_in_phylogeny"].min().reset_index().rename(
+            columns={"y_axis_position_in_phylogeny": "min_y_axis_position_in_phylogeny"}
+        )
+
+        # Annotate min y-axis position per clade to collection.
+        collection_df = collection_df.merge(
+            min_y_axis_position_by_reference_clade,
+            on="subclade_reference",
+            how="left",
+        )
+    else:
+        collection_df["min_y_axis_position_in_phylogeny"] = collection_df["y_axis_position_in_phylogeny"]
 
     # Sort collection by y-axis position.
     sorted_df = collection_df.sort_values(

scripts/get_antigenic_distances_for_reference_model.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+import argparse
+from augur.utils import read_node_data
+import pandas as pd
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--titer-model", required=True, help="node data JSON from titer model with inferred titers annotated in 'nodes' key by field ending with 'cTiterSub'")
+    parser.add_argument("--titers", required=True, help="TSV of titers used to fit the given model")
+    parser.add_argument("--annotations", nargs="+", help="additional annotations to add to the output table in the format of 'key=value' pairs")
+    parser.add_argument("--output", required=True, help="table of antigenic distances in log2 titers between reference and test strains")
+
+    args = parser.parse_args()
+
+    # Load raw titers to get the reference name.
+    raw_titers = pd.read_csv(
+        args.titers,
+        sep="\t",
+        nrows=2,
+    )
+    reference = raw_titers["serum_strain"].values[0]
+
+    # Load titer model data.
+    titer_data = read_node_data(args.titer_model)["nodes"]
+
+    # Convert titer data to a data frame.
+    titer_records = []
+    for test_strain, test_strain_values in titer_data.items():
+        for key, value in test_strain_values.items():
+            if key.endswith("cTiterSub"):
+                titer_records.append({
+                    "reference_strain": reference,
+                    "test_strain": test_strain,
+                    "log2_titer": value,
+                })
+
+    titer_table = pd.DataFrame(titer_records)
+
+    # Add any additional annotations requested by the user in the format of
+    # "key=value" pairs where each key becomes a new column with the given
+    # value.
+    if args.annotations:
+        for annotation in args.annotations:
+            key, value = annotation.split("=")
+            titer_table[key] = value
+
+    # Save the annotated table.
+    titer_table.to_csv(
+        args.output,
+        sep="\t",
+        index=False,
+        float_format="%.4f"
+    )

scripts/get_titers_per_reference.py

@@ -0,0 +1,50 @@
+"""Split titers into separate files per reference virus.
+"""
+import argparse
+import pandas as pd
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--titers", required=True, help="TSV of titers to split by reference")
+    parser.add_argument("--output-references", required=True, help="text file listing the references with titer outputs")
+    parser.add_argument("--output-titers-directory", required=True, help="directory where split titers TSV are placed per reference")
+
+    args = parser.parse_args()
+
+    titers = pd.read_csv(
+        args.titers,
+        sep="\t",
+    )
+
+    # Find references with autologous and heterologous measurements.
+    distinct_pairs = titers.loc[:, ["virus_strain", "serum_strain"]].drop_duplicates()
+    print(f"Found {distinct_pairs.shape[0]} distinct pairs")
+
+    has_autologous_measurement = (distinct_pairs["virus_strain"] == distinct_pairs["serum_strain"])
+    autologous_references = set(distinct_pairs.loc[has_autologous_measurement, "serum_strain"].drop_duplicates().values)
+    print(f"Found {len(autologous_references)} autologous references")
+
+    has_heterologous_measurement = (distinct_pairs["virus_strain"] != distinct_pairs["serum_strain"])
+    heterologous_references = set(distinct_pairs.loc[has_heterologous_measurement, "serum_strain"].drop_duplicates().values)
+    print(f"Found {len(heterologous_references)} heterologous references")
+
+    selected_references = autologous_references & heterologous_references
+    print(f"Found {len(selected_references)} references")
+
+    selected_titers = titers[titers["serum_strain"].isin(selected_references)].copy()
+    selected_titers["reference_path"] = selected_titers["serum_strain"].apply(
+        lambda strain: strain.replace("/", "_")
+    )
+    selected_reference_paths = selected_titers["reference_path"].drop_duplicates().values
+
+    for reference, reference_titers in selected_titers.groupby("reference_path"):
+        reference_titers.to_csv(
+            f"{args.output_titers_directory}/{reference}.tsv",
+            sep="\t",
+            index=False,
+        )
+
+    with open(args.output_references, "w", encoding="utf-8") as oh:
+        for reference in selected_reference_paths:
+            print(reference, file=oh)

workflow/snakemake_rules/titer_models.smk

@@ -232,11 +232,170 @@ rule export_measurements:
             --output-json {output.measurements} 2>&1 | tee {log}
         """
 
+checkpoint get_titers_per_reference:
+    input:
+        titers="builds/{build_name}/titers/{titer_collection}.tsv",
+    output:
+        references="builds/{build_name}/titer_references/{titer_collection}.txt",
+        reference_titers_directory=directory("builds/{build_name}/reference_titers/{titer_collection}/"),
+    conda: "../envs/nextstrain.yaml"
+    shell:
+        r"""
+        mkdir -p {output.reference_titers_directory};
+
+        python scripts/get_titers_per_reference.py \
+            --titers {input.titers} \
+            --output-references {output.references} \
+            --output-titers-directory {output.reference_titers_directory}
+        """
+
+rule reference_model_titers_sub:
+    input:
+        titers = build_dir +"/{build_name}/reference_titers/{titer_collection}/{reference}.tsv",
+        tree = rules.refine.output.tree,
+        translations_done = build_dir + "/{build_name}/{segment}/translations.done"
+    params:
+        genes = get_titer_collection_genes,
+        translations = lambda wildcards: [f"{build_dir}/{wildcards.build_name}/{wildcards.segment}/translations/{gene}_withInternalNodes.fasta" for gene in get_titer_collection_genes(wildcards)],
+        attribute_prefix_argument = get_titer_collection_attribute_prefix_argument,
+    output:
+        titers_model = build_dir + "/{build_name}/{segment}/reference-titers-sub-model/{titer_collection}/{reference}.json",
+    conda: "../envs/nextstrain.yaml"
+    benchmark:
+        "benchmarks/titers_sub_{build_name}_{segment}_{titer_collection}_{reference}.txt",
+    log:
+        "logs/titers_sub_{build_name}_{segment}_{titer_collection}_{reference}.txt",
+    resources:
+        mem_mb=8000,
+    shell:
+        """
+        augur titers sub \
+            --titers {input.titers} \
+            --alignment {params.translations} \
+            --gene-names {params.genes} \
+            --tree {input.tree} \
+            --allow-empty-model \
+            {params.attribute_prefix_argument} \
+            --output {output.titers_model} 2>&1 | tee {log}
+        """
+
+rule reference_model_antigenic_distances_between_strains:
+    input:
+        titer_model="builds/{build_name}/{segment}/reference-titers-sub-model/{titer_collection}/{reference}.json",
+        titers="builds/{build_name}/reference_titers/{titer_collection}/{reference}.tsv",
+    output:
+        distances="builds/{build_name}/{segment}/reference_model_antigenic_distances_between_strains/{titer_collection}/{reference}.tsv",
+    benchmark:
+        "benchmarks/reference_model_antigenic_distances_between_strains_{build_name}_{segment}_{titer_collection}_{reference}.txt"
+    log:
+        "logs/reference_model_antigenic_distances_between_strains_{build_name}_{segment}_{titer_collection}_{reference}.txt"
+    conda: "../envs/nextstrain.yaml"
+    shell:
+        """
+        python3 scripts/get_antigenic_distances_for_reference_model.py \
+            --titer-model {input.titer_model} \
+            --titers {input.titers} \
+            --output {output.distances} &> {log}
+        """
+
+def aggregate_reference_model_distances_input(wildcards):
+    with checkpoints.get_titers_per_reference.get(**wildcards).output["references"].open() as fh:
+        distances = [
+            f"builds/{wildcards.build_name}/{wildcards.segment}/reference_model_antigenic_distances_between_strains/{wildcards.titer_collection}/{reference.strip()}.tsv"
+            for reference in fh
+        ]
+
+    return distances
+
+rule aggregate_reference_model_distances:
+    input:
+        distances=aggregate_reference_model_distances_input,
+    output:
+        distances="builds/{build_name}/{segment}/reference_model_antigenic_distances_between_strains/{titer_collection}.tsv",
+    conda: "../envs/nextstrain.yaml"
+    shell:
+        r"""
+        tsv-append -H {input.distances} > {output.distances}
+        """
+
+rule generate_reference_model_collection_config_json:
+    input:
+        distances="builds/{build_name}/{segment}/reference_model_antigenic_distances_between_strains/{titer_collection}.tsv",
+        tree="builds/{build_name}/{segment}/tree.nwk",
+    output:
+        config_json="builds/{build_name}/{segment}/reference_model_measurements_collection_config/{titer_collection}.json",
+    conda: "../envs/nextstrain.yaml"
+    params:
+        groupings=[
+            "reference_strain",
+        ],
+        fields=[
+            "strain",
+            "reference_strain",
+            "value",
+        ],
+    log:
+        "logs/generate_reference_model_collection_config_json_{build_name}_{segment}_{titer_collection}.txt"
+    shell:
+        """
+        python3 scripts/generate_collection_config_json.py \
+            --tree {input.tree} \
+            --collection {input.distances} \
+            --groupings {params.groupings:q} \
+            --fields {params.fields:q} \
+            --output {output.config_json} &> {log}
+        """
+
+rule export_reference_model_measurements:
+    input:
+        distances="builds/{build_name}/{segment}/reference_model_antigenic_distances_between_strains/{titer_collection}.tsv",
+        configuration="builds/{build_name}/{segment}/reference_model_measurements_collection_config/{titer_collection}.json",
+    output:
+        measurements="builds/{build_name}/{segment}/reference_model_measurements/{titer_collection}.json",
+    conda: "../envs/nextstrain.yaml"
+    benchmark:
+        "benchmarks/export_reference_model_measurements_{build_name}_{segment}_{titer_collection}.txt"
+    log:
+        "logs/export_reference_model_measurements_{build_name}_{segment}_{titer_collection}.txt"
+    params:
+        strain_column="test_strain",
+        value_column="log2_titer",
+        title=lambda wildcards: get_titer_collection_title(wildcards) + " (inferred)",
+        x_axis_label="inferred log2 titer",
+        thresholds=[0.0, 2.0],
+        filters=[
+            "reference_strain",
+        ],
+        include_columns=[
+            "reference_strain",
+        ],
+    shell:
+        """
+        augur measurements export \
+            --collection {input.distances} \
+            --collection-config {input.configuration} \
+            --include-columns {params.include_columns:q} \
+            --strain-column {params.strain_column} \
+            --value-column {params.value_column} \
+            --key {wildcards.titer_collection}_inferred \
+            --title {params.title:q} \
+            --x-axis-label {params.x_axis_label:q} \
+            --thresholds {params.thresholds} \
+            --filters {params.filters} \
+            --show-threshold \
+            --hide-overall-mean \
+            --minify-json \
+            --output-json {output.measurements} 2>&1 | tee {log}
+        """
+
 def get_titer_collections(wildcards):
     files = []
     for collection in config["builds"][wildcards.build_name]["titer_collections"]:
         files.append(f"builds/{wildcards.build_name}/{wildcards.segment}/measurements/{collection['name']}.json")
 
+        if collection.get("run_reference_models"):
+            files.append(f"builds/{wildcards.build_name}/{wildcards.segment}/reference_model_measurements/{collection['name']}.json")
+
     return files
 
 rule concat_measurements: