Merge pull request nextstrain#48 from nextstrain/phylogeo

Phylogeo

Snakefile

@@ -17,6 +17,7 @@ rule files:
         exclude = "config/exclude.txt",
         reference = "config/reference.gb",
         outgroup = "config/outgroup.fasta",
+        weights = "config/weights.tsv",
         auspice_config = "config/auspice_config.json",
         auspice_config_gisaid = "config/auspice_config_gisaid.json",
         auspice_config_zh = "config/auspice_config_zh.json",
@@ -235,6 +236,33 @@ rule translate:
             --output-node-data {output.node_data} \
         """
 
+rule traits:
+    message:
+        """
+        Inferring ancestral traits for {params.columns!s}
+          - increase uncertainty of reconstruction by {params.sampling_bias_correction} to partially account for sampling bias
+        """
+    input:
+        tree = rules.refine.output.tree,
+        metadata = rules.parse.output.metadata,
+        weights = files.weights
+    output:
+        node_data = "results/traits.json",
+    params:
+        columns = "division",
+        sampling_bias_correction = 2
+    shell:
+        """
+        augur traits \
+            --tree {input.tree} \
+            --metadata {input.metadata} \
+            --weights {input.weights} \
+            --output {output.node_data} \
+            --columns {params.columns} \
+            --confidence \
+            --sampling-bias-correction {params.sampling_bias_correction} \
+        """
+
 rule export:
     message: "Exporting data files for for auspice"
     input:
@@ -243,6 +271,7 @@ rule export:
         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 = files.auspice_config,
         colors = files.colors,
         lat_longs = files.lat_longs,
@@ -254,7 +283,7 @@ rule export:
         augur export v2 \
             --tree {input.tree} \
             --metadata {input.metadata} \
-            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} \
+            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} {input.traits} \
             --auspice-config {input.auspice_config} \
             --colors {input.colors} \
             --lat-longs {input.lat_longs} \

config/weights.tsv

@@ -0,0 +1,30 @@
+division	Hubei	100.0
+division	Jiangxi	1.0
+division	Chongqing	1.0
+division	Sichuan	1.0
+division	Yunnan	1.0
+division	Jiangsu	1.0
+division	Zhejiang	1.0
+division	Guangdong	1.0
+division	Taiwan	1.0
+division	Nonthaburi	1.0
+division	Aichi	1.0
+division	Tokyo	1.0
+division	Kyoto	1.0
+division	Gyeonggi	1.0
+division	Singapore	1.0
+division	Thanh Hoa	1.0
+division	Queensland	1.0
+division	New South Wales	1.0
+division	Victoria	1.0
+division	Ile De France	1.0
+division	Flanders	1.0
+division	Bavaria	1.0
+division	England	1.0
+division	Lapland	1.0
+division	Washington	1.0
+division	California	1.0
+division	Arizona	1.0
+division	Illinois	1.0
+division	Wisconsin	1.0
+division	Massachusetts	1.0