Hide branches without clade info

This hides potentially misleading internal branches connecting outbreak
clades.

Author: victorlin

phylogenetic/rules/export.smk

@@ -25,6 +25,24 @@ This part of the workflow usually includes the following steps:
 See Augur's usage docs for these commands for more details.
 """
 
+rule set_hidden_attribute:
+    input:
+        clades = "results/{build}/clades.json",
+        clade_defs = lambda w: config["build_params"][w.build]["files"].get("clades"),
+    output:
+        hidden = "results/{build}/hidden.json",
+    benchmark:
+        "benchmarks/{build}/set_hidden_attribute.txt"
+    log:
+        "logs/{build}/set_hidden_attribute.txt"
+    shell:
+        r"""
+        exec &> >(tee {log:q})
+
+        ./scripts/set-hidden-attribute.py {input.clades:q} {output.hidden:q}
+        """
+
+
 def export_node_data_inputs(w):
     files = [
         f"results/{w.build}/branch_lengths.json",
@@ -35,6 +53,7 @@ def export_node_data_inputs(w):
     if config["build_params"][w.build]["files"].get("clades"):
         files += [
             f"results/{w.build}/clades.json",
+            f"results/{w.build}/hidden.json",
         ]
     return files
 

phylogenetic/scripts/set-hidden-attribute.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+
+import argparse
+import json
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Mark certain nodes with a hidden flag."
+    )
+    parser.add_argument("node_attrs_in", help="Path to a node attributes file.")
+    parser.add_argument("node_attrs_out", help="Path for the output node attributes file.")
+
+    args = parser.parse_args()
+    process_json(args.node_attrs_in, args.node_attrs_out)
+
+
+def process_json(node_attrs_in, node_attrs_out):
+    with open(node_attrs_in, "r") as f:
+        data = json.load(f)
+
+    nodes_in = data.get("nodes", {})
+    nodes_out = {}
+
+    for node_id, node_info in nodes_in.items():
+        # Hide if the node has no clade info.
+        if node_info.get("clade_membership") == "unassigned":
+            nodes_out[node_id] = {"hidden": "timetree"}
+            continue
+
+    with open(node_attrs_out, "w") as f:
+        json.dump({"nodes": nodes_out}, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()