update augur clades, add tests

Commit should be squashed after review

Author: jameshadfield

augur/clades.py

@@ -133,7 +133,7 @@ def assign_clades(clade_designations, all_muts, tree, ref=None):
                     node.sequences[gene][pos] = d
 
 
-    # second pass to assign 'clade_annotation' to basal nodes within each clade
+    # store names of basal nodes of each clade in `basal_clade_nodes` and `clade_membership` dicts.
     # if multiple nodes match, assign annotation to largest
     # otherwise occasional unwanted cousin nodes get assigned the annotation
     for clade_name, clade_alleles in clade_designations.items():
@@ -147,7 +147,7 @@ def assign_clades(clade_designations, all_muts, tree, ref=None):
             basal_clade_nodes[target_node.name] = clade_name
             clade_membership[target_node.name] = clade_name # basal nodes are members of the clade
 
-    # third pass to propagate 'clade_membership'
+    # propagate 'clade_membership' to children nodes
     # don't propagate if encountering 'clade_annotation'
     for node in tree.find_clades(order = 'preorder'):
         for child in node:
@@ -157,29 +157,6 @@ def assign_clades(clade_designations, all_muts, tree, ref=None):
     return (basal_clade_nodes, clade_membership)
 
 
-def create_node_data_structure(basal_clade_nodes, clade_membership, args):
-    node_data = {}
-
-    if (not args.label_name and not args.trait_name):
-        print("WARNING: running `augur clades` without specifying --label-name and/or")
-        print("         --trait-name is deprecated. To preserve backwards compatibility")
-        print("         we will use 'clade' and 'clade_membership', respectively.")
-        print("         (Note that 'clade' is now exported as a 'branch_label')")
-
-        label_name = "clade"
-        trait_name = "clade_membership"
-    else:
-        label_name = args.label_name
-        trait_name = args.trait_name
-
-    if trait_name:
-        node_data['nodes'] = {node: {trait_name: clade} for node,clade in clade_membership.items()}
-    if label_name:
-        node_data['branch_labels'] = {node: {label_name: clade} for node,clade in basal_clade_nodes.items()}
-
-    return node_data
-
-
 def get_reference_sequence_from_root_node(all_muts, root_name):
     # attach sequences to root
     ref = {}
@@ -202,8 +179,7 @@ def register_arguments(parser):
     parser.add_argument('--mutations', nargs='+', help='JSON(s) containing ancestral and tip nucleotide and/or amino-acid mutations ')
     parser.add_argument('--reference', nargs='+', help='fasta files containing reference and tip nucleotide and/or amino-acid sequences ')
     parser.add_argument('--clades', type=str, help='TSV file containing clade definitions by amino-acid')
-    parser.add_argument('--trait-name', type=str, help='name to use to store clade membership (set for every node belonging to the clade)')
-    parser.add_argument('--label-name', type=str, help='name to use for branch labels (set on basal branches for each clade)')
+    parser.add_argument('--attribute-name', type=str, default="clade", help="name to use for clade membership & branch labels", required=False)
     parser.add_argument('--output-node-data', type=str, help='name of JSON file to save clade assignments to')
 
 
@@ -230,8 +206,12 @@ def run(args):
 
     (basal_clade_nodes, clade_membership) = assign_clades(clade_designations, all_muts, tree, ref)
 
-    node_data = create_node_data_structure(basal_clade_nodes, clade_membership, args)
+    # create node_data for export as a JSON
+    node_data = {
+      'nodes': {node: {args.attribute_name: clade} for node,clade in clade_membership.items()},
+      'branch_labels': {node: {args.attribute_name: clade} for node,clade in basal_clade_nodes.items()}
+    }
 
     out_name = get_json_name(args)
     write_json(node_data, out_name)
-    print("clades written to", out_name, file=sys.stdout)
+    print(f"clades written to {out_name} using attribute name {args.attribute_name}", file=sys.stdout)

tests/functional/clades.t

@@ -0,0 +1,34 @@
+Integration tests for augur clades.
+
+  $ pushd "$TESTDIR" > /dev/null
+  $ export AUGUR="../../bin/augur"
+
+Run augur clades without --attribute-name. We expect the name to be "clade"
+
+  $ ${AUGUR} clades \
+  >  --tree clades/tree.nwk \
+  >  --clades clades/clades.tsv \
+  >  --mutations clades/nt_muts.json \
+  >  --output-node-data "$TMP/default.json" > /dev/null
+
+  $ diff -u "clades/expected-output-default.json" "$TMP/default.json"
+(This diff will fail once we have a major version bump... is there a better JSON diff tool to use?)
+
+Run augur clades with a custom --attribute-name
+
+  $ ${AUGUR} clades \
+  >  --tree clades/tree.nwk \
+  >  --clades clades/clades.tsv \
+  >  --mutations clades/nt_muts.json \
+  >  --attribute-name custom \
+  >  --output-node-data "$TMP/custom-attr.json" > /dev/null
+
+  $ diff -u "clades/expected-output-custom-attr.json" "$TMP/custom-attr.json"
+(This diff will fail once we have a major version bump... is there a better JSON diff tool to use?)
+
+# Ensure the only change between runs of `augur clades` is the attr name used
+  $ cat "$TMP/default.json" | sed "s/clade/custom/" > "$TMP/default-now-custom.json"
+  $ diff -u "$TMP/default-now-custom.json" "$TMP/custom-attr.json"
+
+Cleanup
+  $ rm -f "$TMP/default.json" "$TMP/custom-attr.json" "$TMP/default-now-custom.json"

tests/functional/clades/clades.tsv

@@ -0,0 +1,12 @@
+clade	gene	site	alt
+
+# the 1b mutation only once, on the branch leading to tips B and C
+# thus we expect the clade label to be on node `internalAB`
+cladeCB	nuc	1	B
+# the 2c mutation appears twice -- on branch `internalAB` and `internalDEF`
+# as the latter has 3 descendants, it is chosen over the former
+cladeDEF	nuc	2	C
+# mutation 3e appears only on a terminal node (tipE)
+# but we still expect both a branch label and a node_attr
+# this means that tipE should be annotated "cladeE" and _not_ "cladeDEF"
+cladeE	nuc	3	E

tests/functional/clades/expected-output-custom-attr.json

@@ -0,0 +1,40 @@
+{
+  "branch_labels": {
+    "internalBC": {
+      "custom": "cladeCB"
+    },
+    "internalDEF": {
+      "custom": "cladeDEF"
+    },
+    "tipE": {
+      "custom": "cladeE"
+    }
+  },
+  "generated_by": {
+    "program": "augur",
+    "version": "12.0.0"
+  },
+  "nodes": {
+    "internalBC": {
+      "custom": "cladeCB"
+    },
+    "internalDEF": {
+      "custom": "cladeDEF"
+    },
+    "tipB": {
+      "custom": "cladeCB"
+    },
+    "tipC": {
+      "custom": "cladeCB"
+    },
+    "tipD": {
+      "custom": "cladeDEF"
+    },
+    "tipE": {
+      "custom": "cladeE"
+    },
+    "tipF": {
+      "custom": "cladeDEF"
+    }
+  }
+}

tests/functional/clades/expected-output-default.json

@@ -0,0 +1,40 @@
+{
+  "branch_labels": {
+    "internalBC": {
+      "clade": "cladeCB"
+    },
+    "internalDEF": {
+      "clade": "cladeDEF"
+    },
+    "tipE": {
+      "clade": "cladeE"
+    }
+  },
+  "generated_by": {
+    "program": "augur",
+    "version": "12.0.0"
+  },
+  "nodes": {
+    "internalBC": {
+      "clade": "cladeCB"
+    },
+    "internalDEF": {
+      "clade": "cladeDEF"
+    },
+    "tipB": {
+      "clade": "cladeCB"
+    },
+    "tipC": {
+      "clade": "cladeCB"
+    },
+    "tipD": {
+      "clade": "cladeDEF"
+    },
+    "tipE": {
+      "clade": "cladeE"
+    },
+    "tipF": {
+      "clade": "cladeDEF"
+    }
+  }
+}

tests/functional/clades/nt_muts.json

@@ -0,0 +1,19 @@
+{
+  "nodes": {
+    "tipA": {"muts": [], "aa_muts": {}},
+    "tipB": {"muts": [], "aa_muts": {}},
+    "tipC": {"muts": [], "aa_muts": {}},
+    "tipD": {"muts": [], "aa_muts": {}},
+    "tipE": {"muts": ["A3E"], "aa_muts": {}},
+    "tipF": {"muts": [], "aa_muts": {}},
+    "internalBC": {
+      "muts": ["A1B", "A2C"],
+      "aa_muts": {}
+    },
+    "internalDEF": {
+      "muts": ["A2C"],
+      "aa_muts": {}
+    },
+    "ROOT":{"muts": [], "aa_muts": {}}
+  }
+}

tests/functional/clades/tree.nwk

@@ -0,0 +1 @@
+(tipA:1,(tipB:1,tipC:1)internalBC:2,(tipD:3,tipE:4,tipF:1)internalDEF:5)ROOT:1;