From 0162bf4a39d5abf68821af67c31d7ba649b9de4e Mon Sep 17 00:00:00 2001
From: james hadfield <hadfield.james@gmail.com>
Date: Thu, 22 Aug 2024 14:57:33 +1200
Subject: [PATCH] Add sampling year coloring

For builds where we've chosen not to infer a timetree it's still very
useful to communicate the sampling date. Here we export the year which
forms a nice continuous coloring. (An ordinal coloring would be better
but the date range in this dataset means Auspice repeats colors which
isn't ideal. Temporal coloring will be better once it accepts YYYY-MM-DD
strings, but as of Auspice 2.55 it interprets integers as January 1st.)
---
 phylogenetic/defaults/auspice_config.json |  5 +++++
 phylogenetic/rules/annotate_phylogeny.smk | 15 +++++++++++++++
 phylogenetic/rules/export.smk             |  3 ++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/phylogenetic/defaults/auspice_config.json b/phylogenetic/defaults/auspice_config.json
index ea4eea1..49af8f7 100644
--- a/phylogenetic/defaults/auspice_config.json
+++ b/phylogenetic/defaults/auspice_config.json
@@ -27,6 +27,11 @@
       "title": "Country",
       "type": "categorical"
     },
+    {
+      "key": "year",
+      "title": "Sampling year",
+      "type": "continuous"
+    },
     {
       "key": "host",
       "title": "Host",
diff --git a/phylogenetic/rules/annotate_phylogeny.smk b/phylogenetic/rules/annotate_phylogeny.smk
index f032e10..0cb802b 100644
--- a/phylogenetic/rules/annotate_phylogeny.smk
+++ b/phylogenetic/rules/annotate_phylogeny.smk
@@ -51,3 +51,18 @@ rule translate:
             --output {output.node_data} \
             2>&1 | tee {log}
         """
+
+rule add_year_metadata:
+    input:
+        metadata = "data/metadata.tsv",
+    params:
+        strain_id = config["strain_id_field"],
+    output:
+        node_data = "results/year.json"
+    run:
+        from augur.io import read_metadata
+        import json
+        m = read_metadata(input.metadata, id_columns=[params.strain_id])
+        nodes = {name: {'year': date.split('-')[0]} for name,date in zip(m.index, m['date']) if date and not date.startswith('X')}
+        with open(output.node_data, 'w') as fh:
+            json.dump({"nodes": nodes}, fh)
diff --git a/phylogenetic/rules/export.smk b/phylogenetic/rules/export.smk
index 2937986..b85de5f 100644
--- a/phylogenetic/rules/export.smk
+++ b/phylogenetic/rules/export.smk
@@ -13,6 +13,7 @@ rule export:
         branch_lengths = "results/branch_lengths.json",
         nt_muts = "results/nt_muts.json",
         aa_muts = "results/aa_muts.json",
+        year = "results/year.json",
         colors = config["files"]["colors"],
         auspice_config = config["files"]["auspice_config"],
         description = config["files"]["description"]
@@ -30,7 +31,7 @@ rule export:
             --tree {input.tree} \
             --metadata {input.metadata} \
             --metadata-id-columns {params.strain_id} \
-            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} \
+            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} {input.year} \
             --colors {input.colors} \
             --auspice-config {input.auspice_config} \
             --include-root-sequence-inline \