script/assign_rbd_levels: Update clade parsing

The JSON output from `augur clades` was updated to separate `nodes`
and `branches` in nextstrain/augur#728 so now
the `assign_rbd_levels` script needs to parse the `branches` in order
to find the basal node.

scripts/assign_rbd_levels.py

@@ -7,12 +7,12 @@
 def find_matching_nodes(clades_fname, basal_clade_label, tree_fname):
     basal_node_name = None
     with open(clades_fname) as fh:
-        for name, node_data in json.load(fh)['nodes'].items():
-            if node_data.get('clade_annotation', '') == basal_clade_label:
+        for name, node_data in json.load(fh)['branches'].items():
+            if node_data.get('labels', {}).get('clade', '') == basal_clade_label:
                 basal_node_name = name
                 break
     if not basal_node_name:
-        print(f"WARNING: no node found with a clade_annotation of {basal_clade_label}. This script will proceed, but no levels will be exported.")
+        print(f"WARNING: no branch found with a clade of {basal_clade_label}. This script will proceed, but no levels will be exported.")
         return set()
     print(f"Node representing {basal_clade_label}: {basal_node_name}")
     T = Phylo.read(tree_fname, 'newick')
@@ -92,4 +92,4 @@ def classify_into_levels(spike_seq, rbd_mutations):
         node_data['rbd_level_details'][name] = ", ".join([f"S:{x[0][1]}{x[1]} ({x[2]})" for x in zip(rbd_mutations, codons, calls)])
 
     with open(args.output_node_data, 'w') as fh:
-        json.dump(node_data, fh, indent=2)
+        json.dump(node_data, fh, indent=2)