Add missing lineages instead of recombinant lineages

scripts/modify_covizu_tree.py

@@ -8,27 +8,22 @@ def main():
     warnings.filterwarnings('ignore')
 
     # command line interface
-    parser = argparse.ArgumentParser(description='Add recombinant lineages to covizu tree')
+    parser = argparse.ArgumentParser(description='Add missing lineages to covizu tree')
     parser.add_argument('--covizu-tree', required=True, help='CoVizu time-scaled tree file')
     parser.add_argument('--output', required=True, help='Modified tree file')
     args = parser.parse_args()
 
-    # get recombinant lineages from github repo: pango-designation
+    # add lineages to tree
+    t = Tree(args.covizu_tree, format=1)
+    leaves = [leaf.name for leaf in t.get_leaves()]
     url = 'https://raw.githubusercontent.com/cov-lineages/pango-designation/master/lineage_notes.txt'
     resp = requests.get(url)
     content = resp.content.decode()
-    recomb = []
     for i in content.split('\n')[1:]:
         l = i.split('\t')[0]
-        if l.startswith('X'):
-            recomb.append(l)
-
-    # add recombinant lineages to tree
-    t = Tree(args.covizu_tree, format=1)
-    leaves = [leaf.name for leaf in t.get_leaves()]
-    for r in recomb:
-        if r not in leaves:
-            t.add_child(name=r, dist=1)
+        if (not l.startswith('*')) and (l != ''):
+            if l not in leaves:
+                t.add_child(name=l, dist=1)
 
     # write modified tree
     t.write(outfile=args.output)