add_labels.py

importargparse
importjson
fromBioimportPhylo
fromcollectionsimportdefaultdict

defattach_labels(d, labeled_nodes):
if"children"ind:
forcind["children"]:
ifc["name"] inlabeled_nodes:
if"labels"notinc["branch_attrs"]:
c["branch_attrs"]["labels"] = {}
c['branch_attrs']['labels']['mlabel'] =labeled_nodes[c["name"]][0]
print(c['branch_attrs']['labels'])
attach_labels(c, labeled_nodes)


if__name__=='__main__':
parser=argparse.ArgumentParser(
description="Remove extraneous colorings",
formatter_class=argparse.ArgumentDefaultsHelpFormatter
 )

parser.add_argument('--input', type=str, metavar="JSON", required=True, help="input Auspice JSON")
parser.add_argument('--tree', type=str, required=True, help="tree file")
parser.add_argument('--clades', type=str, required=True, help="clades")
parser.add_argument('--mutations', type=str, required=True, help="mutations")
parser.add_argument('--output', type=str, metavar="JSON", required=True, help="output Auspice JSON")
args=parser.parse_args()

T=Phylo.read(args.tree, 'newick')

withopen(args.mutations, "r") asf:
mutation_json=json.load(f)['nodes']

withopen(args.clades, "r") asf:
clades_json=json.load(f)['nodes']

withopen(args.input, "r") asf:
input_json=json.load(f)

nodes= {}
forninT.find_clades(order='postorder'):
ifn.is_terminal():
n.tip_count=1
else:
n.tip_count=sum([c.tip_countforcinn])
nodes[n.name] = {'tip_count':n.tip_count}

labels=defaultdict(list)
fornodeinnodes:
forminmutation_json[node]['muts']:
ifm[0] in'ACGT'andm[-1] in'ACGT':
clade=clades_json[node]['clade_membership']
tmp_label= (clade, m)
labels[tmp_label].append((node, nodes[node]['tip_count']))

labeled_nodes=defaultdict(list)
forlabelinlabels:
node=sorted(labels[label], key=lambdax:-x[1])[0]
labeled_nodes[node[0]].append('/'.join(label))

attach_labels(input_json["tree"], labeled_nodes)

withopen(args.output, 'w') asf:
json.dump(input_json, f, indent=2)