george-githinji · November 23, 2021 03:18
diff --git a/README.md b/README.md
diff --git a/auspice_tree_to_table.py b/auspice_tree_to_table.py
 import argparse
 from augur.utils import json_to_tree
 import json
 import pandas as pd
 import sys


 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("tree", help="auspice tree JSON")
    parser.add_argument("output", help="tab-delimited file of attributes per node of the given tree")
    parser.add_argument("--include-internal-nodes", action="store_true", help="include data from internal nodes in output")
    parser.add_argument("--attributes", nargs="+", help="names of attributes to export from the given tree")

    args = parser.parse_args()

    # Load tree from JSON.
    with open(args.tree, "r", encoding="utf-8") as fh:
        tree_json = json.load(fh)

    tree = json_to_tree(tree_json)

    # Collect attributes per node from the tree to export.
    records = []

    if args.attributes:
        attributes = args.attributes
    else:
        attributes = sorted(
            set(tree.root.node_attrs.keys()) |
            set(tree.root.branch_attrs.keys())
        )

    for node in tree.find_clades():
        if node.is_terminal() or args.include_internal_nodes:
            record = {
                "name": node.name
            }

            for attribute in attributes:
                if attribute in node.node_attrs:
                    value = node.node_attrs[attribute]
                elif attribute in node.branch_attrs:
                    value = node.branch_attrs[attribute]
                else:
                    print(f"Could not find attribute '{attribute}' for node '{node.name}'.", file=sys.stderr)
                    value = None

                if value is not None:
                    if isinstance(value, dict) and "value" in value:
                        value = value["value"]

                record[attribute] = value

            records.append(record)

    # Convert records to a data frame and save as a tab-delimited file.
    df = pd.DataFrame(records)
    df.to_csv(args.output, sep="\t", header=True, index=False, columns=["name"] + list(attributes), float_format="%.2f")
	import argparse
	from augur.utils import json_to_tree
	import json
	import pandas as pd
	import sys


	if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("tree", help="auspice tree JSON")
	parser.add_argument("output", help="tab-delimited file of attributes per node of the given tree")
	parser.add_argument("--include-internal-nodes", action="store_true", help="include data from internal nodes in output")
	parser.add_argument("--attributes", nargs="+", help="names of attributes to export from the given tree")

	args = parser.parse_args()

	# Load tree from JSON.
	with open(args.tree, "r", encoding="utf-8") as fh:
	tree_json = json.load(fh)

	tree = json_to_tree(tree_json)

	# Collect attributes per node from the tree to export.
	records = []

	if args.attributes:
	attributes = args.attributes
	else:
	attributes = sorted(
	set(tree.root.node_attrs.keys()) \|
	set(tree.root.branch_attrs.keys())
	)

	for node in tree.find_clades():
	if node.is_terminal() or args.include_internal_nodes:
	record = {
	"name": node.name
	}

	for attribute in attributes:
	if attribute in node.node_attrs:
	value = node.node_attrs[attribute]
	elif attribute in node.branch_attrs:
	value = node.branch_attrs[attribute]
	else:
	print(f"Could not find attribute '{attribute}' for node '{node.name}'.", file=sys.stderr)
	value = None

	if value is not None:
	if isinstance(value, dict) and "value" in value:
	value = value["value"]

	record[attribute] = value

	records.append(record)

	# Convert records to a data frame and save as a tab-delimited file.
	df = pd.DataFrame(records)
	df.to_csv(args.output, sep="\t", header=True, index=False, columns=["name"] + list(attributes), float_format="%.2f")
No results found