diff --git a/augur/export_v2.py b/augur/export_v2.py index af563bf50..cfb31a53a 100644 --- a/augur/export_v2.py +++ b/augur/export_v2.py @@ -117,6 +117,24 @@ def order_nodes(node): return od +def read_trees(filenames): + trees = [Phylo.read(fname, 'newick') for fname in filenames] + # augur export requires unique node names (both terminal and external) as these + # are used to associate metadata/node-data with nodes. Any duplication is fatal. + # The exception to this is unlabelled node names, which auspice will handle but + # won't be associated with any metadata within export. + node_names = [clade.name for tree in trees for clade in tree.root.find_clades()] + if None in node_names: + raise AugurError(f"Tree contains unnamed nodes. If these are internal nodes you may wish to run "+ + "`augur refine --tree --output-tree ` to name them.") + if len(set(node_names))!=len(node_names): + from collections import Counter + dups = [name for name, count in Counter(node_names).items() if count>1] + raise AugurError(f"{len(dups)} node names occur multiple times in the tree: " + + ", ".join([f"'{v}'" for v in dups[0:5]]) + ("..." if len(dups)>5 else "")) + return (trees, node_names) + + def node_div(T, node_attrs): """ Scans the provided tree & metadata to see if divergence is defined, and if so returns @@ -732,7 +750,12 @@ def _recursively_set_data(node): node['branch_attrs'] = branch_attrs[node['name']] for child in node.get("children", []): _recursively_set_data(child) - _recursively_set_data(data_json["tree"]) + + if isinstance(data_json["tree"], list): + for subtree in data_json['tree']: + _recursively_set_data(subtree) + else: + _recursively_set_data(data_json["tree"]) def set_node_attrs_on_tree(data_json, node_attrs, additional_metadata_columns): @@ -821,7 +844,11 @@ def _recursively_set_data(node): for child in node.get("children", []): _recursively_set_data(child) - _recursively_set_data(data_json["tree"]) + if isinstance(data_json["tree"], list): + for subtree in data_json['tree']: + _recursively_set_data(subtree) + else: + _recursively_set_data(data_json["tree"]) def node_data_prop_is_normal_trait(name): # those traits / keys / attrs which are not "special" and can be exported @@ -875,7 +902,7 @@ def register_parser(parent_subparsers): required = parser.add_argument_group( title="REQUIRED" ) - required.add_argument('--tree','-t', metavar="newick", required=True, help="Phylogenetic tree, usually output from `augur refine`") + required.add_argument('--tree','-t', metavar="newick", nargs='+', required=True, help="Phylogenetic tree(s), usually output from `augur refine`") required.add_argument('--output', metavar="JSON", required=True, help="Output file (typically for visualisation in auspice)") config = parser.add_argument_group( @@ -1059,12 +1086,12 @@ def create_branch_labels(branch_attrs, node_data, branch_data): continue branch_attrs[node_name]["labels"][label_key] = label_value -def parse_node_data_and_metadata(T, node_data, metadata): +def parse_node_data_and_metadata(node_names, node_data, metadata): node_data_names = set() metadata_names = set() # assign everything to node_attrs, exclusions considered later - node_attrs = {clade.name: {} for clade in T.root.find_clades()} + node_attrs = {name: {} for name in node_names} # first pass: metadata for metadata_id, node in metadata.items(): @@ -1088,7 +1115,7 @@ def parse_node_data_and_metadata(T, node_data, metadata): # third pass: create `branch_attrs`. The data comes from # (a) some keys within `node_data['nodes']` (for legacy reasons) # (b) the `node_data['branches']` dictionary, which currently only defines labels - branch_attrs = {clade.name: defaultdict(dict) for clade in T.root.find_clades()} + branch_attrs = {name: defaultdict(dict) for name in node_names} create_branch_mutations(branch_attrs, node_data) create_branch_labels(branch_attrs, node_data['nodes'], node_data.get('branches', {})) @@ -1173,9 +1200,9 @@ def run(args): metadata_file = {} # parse input files - T = Phylo.read(args.tree, 'newick') + (trees, node_names) = read_trees(args.tree) node_data, node_attrs, node_data_names, metadata_names, branch_attrs = \ - parse_node_data_and_metadata(T, node_data_file, metadata_file) + parse_node_data_and_metadata(node_names, node_data_file, metadata_file) config = get_config(args) additional_metadata_columns = get_additional_metadata_columns(config, args.metadata_columns, metadata_names) @@ -1205,7 +1232,8 @@ def run(args): set_filters(data_json, config) # set tree structure - data_json["tree"] = convert_tree_to_json_structure(T.root, node_attrs, node_div(T, node_attrs)) + trees_json = [convert_tree_to_json_structure(tree.root, node_attrs, node_div(tree, node_attrs)) for tree in trees] + data_json["tree"] = trees_json[0] if len(trees_json)==1 else trees_json set_node_attrs_on_tree(data_json, node_attrs, additional_metadata_columns) set_branch_attrs_on_tree(data_json, branch_attrs) diff --git a/augur/refine.py b/augur/refine.py index 95717861f..348898377 100644 --- a/augur/refine.py +++ b/augur/refine.py @@ -143,6 +143,7 @@ def register_parser(parent_subparsers): default='mutations-per-site', help='Units in which sequence divergences is exported.') parser.add_argument('--seed', type=int, help='seed for random number generation') parser.add_argument('--verbosity', type=int, default=1, help='treetime verbosity, between 0 and 6 (higher values more output)') + parser.add_argument('--content-aware-internal-node-names', action="store_true") parser.set_defaults(covariance=True) return parser @@ -342,6 +343,21 @@ def are_sequence_states_different(nuc1, nuc2): print("ERROR: divergence unit",args.divergence_units,"not supported!", file=sys.stderr) return 1 + if args.content_aware_internal_node_names: + terminals = [n.name for n in T.find_clades() if n.is_terminal()] + terminals.sort() + internals = {n.name for n in T.find_clades() if not n.is_terminal()} + assert all(n.startswith('NODE_') for n in internals) + # compute hash similar to git short commit hash + import hashlib + id = hashlib.sha256("".join(terminals).encode('utf-8')).hexdigest()[0:7] + def rename(name): + if name not in internals: return name + return f"NODE_{id}_{name.split('_')[1]}" + for n in T.find_clades(): + n.name = rename(n.name) + node_data['nodes'] = {rename(name):data for name,data in node_data['nodes'].items()} + # Export refined tree and node data tree_success = Phylo.write(T, tree_fname, 'newick', format_branch_length='%1.8f', branch_length_only=True) print("updated tree written to",tree_fname, file=sys.stdout) diff --git a/augur/validate_export.py b/augur/validate_export.py index 5cd3f4c4c..80da43bfe 100644 --- a/augur/validate_export.py +++ b/augur/validate_export.py @@ -7,7 +7,7 @@ import sys from collections import defaultdict -def ensure_no_duplicate_names(root, ValidateError): +def ensure_no_duplicate_names(tree, ValidateError): """ Check that all node names are identical, which is required for auspice (v2) JSONs. """ @@ -18,10 +18,14 @@ def recurse(node): names.add(node["name"]) if "children" in node: [recurse(child) for child in node["children"]] - recurse(root) + if isinstance(tree, list): + for subtree in tree: + recurse(subtree) + else: + recurse(tree) -def collectTreeAttrsV2(root, warn): +def collectTreeAttrsV2(tree, warn): """ Collect all keys specified on `node["node_attrs"]` throughout the tree and the values associated with them. Note that this will only look at @@ -47,7 +51,12 @@ def recurse(node): [recurse(child) for child in node["children"]] else: num_terminal += 1 - recurse(root) + + if isinstance(tree, list): + for subtree in tree: + recurse(subtree) + else: + recurse(tree) for data in seen.values(): if data["count"] == num_nodes: @@ -56,7 +65,7 @@ def recurse(node): return(seen, num_terminal) -def collectMutationGenes(root): +def collectMutationGenes(tree): """ Returns a set of all genes specified in the tree in the "aa_muts" objects """ @@ -67,17 +76,28 @@ def recurse(node): genes.update(mutations.keys()) if "children" in node: [recurse(child) for child in node["children"]] - recurse(root) + + if isinstance(tree, list): + for subtree in tree: + recurse(subtree) + else: + recurse(tree) + genes -= {"nuc"} return genes -def collectBranchLabels(root): +def collectBranchLabels(tree): labels = set() def recurse(node): labels.update(node.get("branch_attrs", {}).get("labels", {}).keys()) if "children" in node: [recurse(child) for child in node["children"]] - recurse(root) + + if isinstance(tree, list): + for subtree in tree: + recurse(subtree) + else: + recurse(tree) return labels def verifyMainJSONIsInternallyConsistent(data, ValidateError): diff --git a/tests/functional/export_v2/cram/duplicate-node-names.t b/tests/functional/export_v2/cram/duplicate-node-names.t new file mode 100644 index 000000000..2a8c8ca77 --- /dev/null +++ b/tests/functional/export_v2/cram/duplicate-node-names.t @@ -0,0 +1,10 @@ + $ source "$TESTDIR"/_setup.sh + + $ echo "(tipA:1,(tipA:1,tipA:1)internalBC:2,(tipD:3,tipE:4,tipF:1)internalDEF:5)ROOT:0;" > tree1.nwk + + $ ${AUGUR} export v2 \ + > --tree tree1.nwk \ + > --output minimal.json + ERROR: 1 node names occur multiple times in the tree: 'tipA' + [2] + diff --git a/tests/functional/export_v2/cram/minify-output.t b/tests/functional/export_v2/cram/minify-output.t index 68a573369..a9035a9a0 100644 --- a/tests/functional/export_v2/cram/minify-output.t +++ b/tests/functional/export_v2/cram/minify-output.t @@ -16,13 +16,16 @@ names to create larger file sizes with less recursion. > } A small tree is not automatically minified. -The unminified output is 13KB which is considered acceptable. +The unminified output is 15KB which is considered acceptable. + + $ echo "$(generate_newick 10);" > small_tree_raw.nwk + +Use augur refine to add internal node names + $ ${AUGUR} refine --tree small_tree_raw.nwk --output-tree small_tree.nwk &>/dev/null - $ echo "$(generate_newick 10);" > small_tree.nwk $ ${AUGUR} export v2 \ > --tree small_tree.nwk \ - > --skip-validation \ > --output output.json &>/dev/null $ head -c 20 output.json @@ -30,13 +33,12 @@ The unminified output is 13KB which is considered acceptable. "version": "v2", (no-eol) $ ls -l output.json | awk '{print $5}' - 13813 + 15039 It can be forcefully minified with an argument. $ ${AUGUR} export v2 \ > --tree small_tree.nwk \ - > --skip-validation \ > --minify-json \ > --output output.json &>/dev/null @@ -48,21 +50,22 @@ even if it may seem "falsey". $ AUGUR_MINIFY_JSON=0 ${AUGUR} export v2 \ > --tree small_tree.nwk \ - > --skip-validation \ > --output output.json &>/dev/null $ head -c 20 output.json {"version": "v2", "m (no-eol) -A large tree, when forcefully not minified, has an output size of 6MB which is +A large tree, when forcefully not minified, has an output size of 8MB which is considered large. - $ echo "$(generate_newick 500);" > big_tree.nwk + $ echo "$(generate_newick 500);" > big_tree_raw.nwk + +Use augur refine to add internal node names + $ ${AUGUR} refine --tree big_tree_raw.nwk --output-tree big_tree.nwk &>/dev/null $ ${AUGUR} export v2 \ > --tree big_tree.nwk \ - > --skip-validation \ > --no-minify-json \ > --output output.json &>/dev/null @@ -71,7 +74,7 @@ considered large. "version": "v2", (no-eol) $ ls -l output.json | awk '{print $5}' - 6568454 + 8591420 This means it is automatically minified. @@ -84,4 +87,4 @@ This means it is automatically minified. {"version": "v2", "m (no-eol) $ ls -l output.json | awk '{print $5}' - 561436 + 576414 diff --git a/tests/functional/export_v2/cram/missing-internal-nodes.t b/tests/functional/export_v2/cram/missing-internal-nodes.t new file mode 100644 index 000000000..3732a9c72 --- /dev/null +++ b/tests/functional/export_v2/cram/missing-internal-nodes.t @@ -0,0 +1,18 @@ +Setup + + $ source "$TESTDIR"/_setup.sh + +Test a tree with unlabelled internal nodes. +In Augur v24.3.0 and earlier these would work _sometimes_ depending +on what code paths were hit, but in most real-life cases would raise +an uncaught error. We now check for these when we parse the tree. + + $ echo "(tipA:1,(tipB:1,tipC:1):2,(tipD:3,tipE:4,tipF:1):5):0;" > tree1.nwk + + $ ${AUGUR} export v2 \ + > --tree tree1.nwk \ + > --metadata "$TESTDIR/../data/dataset1_metadata_with_name.tsv" \ + > --node-data "$TESTDIR/../data/div_node-data.json" \ + > --output test1.json + ERROR: Tree contains unnamed nodes.+ (re) + [2] diff --git a/tests/functional/export_v2/cram/multi-tree.t b/tests/functional/export_v2/cram/multi-tree.t new file mode 100644 index 000000000..a7a45b948 --- /dev/null +++ b/tests/functional/export_v2/cram/multi-tree.t @@ -0,0 +1,53 @@ + +Setup + + $ source "$TESTDIR"/_setup.sh + +Create a small second tree (which has different names/labels than 'tree.nwk') + $ cat > tree2.nwk <<~~ + > (tipG:1,(tipH:1,tipI:1)internalHI:2)SECOND_ROOT:0; + > ~~ + +Minimal export -- no node data, no metadata etc etc + $ ${AUGUR} export v2 \ + > --tree "$TESTDIR/../data/tree.nwk" tree2.nwk \ + > --output minimal.json &> /dev/null + +More realistic export - with node_data for all nodes and metadata for some of them + + $ cat > metadata.tsv <<~~ + > strain something + > tipA foo + > tipB foo + > tipC foo + > tipG bar + > tipH bar + > ~~ + + + $ cat > node-data.json <<~~ + > { + > "nodes": { + > "ROOT": {"mutation_length": 0}, + > "tipA": {"mutation_length": 1}, + > "internalBC": {"mutation_length": 2}, + > "tipB": {"mutation_length": 1}, + > "tipC": {"mutation_length": 1}, + > "internalDEF": {"mutation_length": 5}, + > "tipD": {"mutation_length": 3}, + > "tipE": {"mutation_length": 4}, + > "tipF": {"mutation_length": 1}, + > "SECOND_ROOT": {"mutation_length": 0}, + > "tipG": {"mutation_length": 1}, + > "internalHI": {"mutation_length": 2}, + > "tipH": {"mutation_length": 1}, + > "tipI": {"mutation_length": 1} + > } + > } + > ~~ + + $ ${AUGUR} export v2 \ + > --tree "$TESTDIR/../data/tree.nwk" tree2.nwk \ + > --metadata metadata.tsv --color-by-metadata something \ + > --node-data node-data.json \ + > --output output.json &> /dev/null