feat(itol): Polyphyletic phyla will be the same colour.

Author: aaronmussig

gtdb_itol_decorate/__main__.py

@@ -1,9 +1,10 @@
+import sys
 from pathlib import Path
 
 import typer
 
 from gtdb_itol_decorate.gtdb import load_taxonomy_file, get_taxon_to_phylum
-from gtdb_itol_decorate.itol import get_phylum_to_lca, get_phylum_colours, write_color_datastrip, \
+from gtdb_itol_decorate.itol import get_phylum_colours, write_color_datastrip, \
     get_internal_nodes_with_labels, write_internal_node_labels, write_tree_colours, write_collapse_file, \
     write_popup_file
 from gtdb_itol_decorate.newick import load_newick_to_tree, validate_dendropy_namespace, \
@@ -28,14 +29,18 @@ def main(tree_path: Path, tax_path: Path, out_dir: Path):
 
     log('Reverse mapping taxon to phylum')
     d_taxon_to_phylum = get_taxon_to_phylum(d_tax)
+    phyla = frozenset(d_taxon_to_phylum.values())
+    if len(phyla) == 0:
+        log('No phyla could be found, please send me your tree!')
+        sys.exit(1)
 
     log('Annotating internal nodes with descendant taxa')
     set_node_desc_taxa(tree)
-    set_taxon_label_for_internal_nodes(tree, d_tax)
+    d_taxon_to_lca = set_taxon_label_for_internal_nodes(tree, d_tax)
 
     log('Getting the last common ancestor of each phylum.')
-    d_phylum_to_lca = get_phylum_to_lca(tree)
-    d_phylum_palette = get_phylum_colours(d_phylum_to_lca)
+    d_phylum_to_lca = {k: v for k, v in d_taxon_to_lca.items() if k.startswith('p__')}
+    d_phylum_palette = get_phylum_colours(phyla)
     write_color_datastrip(d_phylum_to_lca, d_phylum_palette, out_dir / 'itol_dataset_strip_phylum.txt')
 
     log('Making tree compatible with iTOL (stripping labels)')

gtdb_itol_decorate/decorate.py

@@ -1,16 +1,13 @@
 from collections import defaultdict
 
 
-
 def decorate_tree(tree, d_gid_to_tax):
-
     d_node_to_leaf_taxa = defaultdict(set)
     for node in tree.postorder_node_iter():
         # desc_taxa = get_node_desc_taxa(node)
 
         print(node)
 
-
     # Create a mapping from the taxon to each gid
     d_taxon_to_gids = defaultdict(set)
     for gid, taxon in d_gid_to_tax.items():

gtdb_itol_decorate/gtdb.py

@@ -1,4 +1,3 @@
-from collections import defaultdict
 from typing import Dict
 
 from gtdblib.taxon.rank import TaxonRank
@@ -10,8 +9,8 @@
 
 def load_taxonomy_file(path: str, limit_to_gids: set):
     gtdb_ranks = [TaxonRank.DOMAIN, TaxonRank.PHYLUM, TaxonRank.CLASS,
-                   TaxonRank.ORDER, TaxonRank.FAMILY, TaxonRank.GENUS,
-                   TaxonRank.SPECIES]
+                  TaxonRank.ORDER, TaxonRank.FAMILY, TaxonRank.GENUS,
+                  TaxonRank.SPECIES]
     out = dict()
     with open(path) as f:
         for line in f.readlines():
@@ -34,7 +33,6 @@ def load_taxonomy_file(path: str, limit_to_gids: set):
     return out
 
 
-
 def get_taxon_to_phylum(d_tax: Dict[str, Taxonomy]):
     out = dict()
     for gid, taxonomy in d_tax.items():
@@ -43,3 +41,9 @@ def get_taxon_to_phylum(d_tax: Dict[str, Taxonomy]):
             out[taxon] = taxonomy.p.value
     return out
 
+
+def remove_polyphyletic_suffix(taxon):
+    if not (taxon.startswith('g__') or taxon.startswith('s__')):
+        if taxon[-2] == '_':
+            return taxon[:-2]
+    return taxon

gtdb_itol_decorate/itol.py

@@ -1,8 +1,9 @@
-from collections import deque, defaultdict
+from collections import defaultdict
 from pathlib import Path
-from typing import Dict, Set
+from typing import Dict, FrozenSet
 
 import dendropy
+import seaborn as sns
 from gtdblib.file.itol.collapse import iTolCollapseFile
 from gtdblib.file.itol.dataset_color_strip import iTolDatasetColorStripFile
 from gtdblib.file.itol.label import iTolLabelFile
@@ -11,37 +12,24 @@
 from gtdblib.taxonomy.taxonomy import Taxonomy
 from gtdblib.util.color import TABLEAU_20, rgb_to_hex
 
-from gtdb_itol_decorate.newick import parse_label
-
-import seaborn as sns
-
-from gtdb_itol_decorate.util import canonical_gid
-
-
-def get_lca_str(node: dendropy.Node):
-    if node.is_leaf():
-        return node.taxon.label
-    if len(node.child_nodes()) < 2:
-        return node.child_nodes()[0].leaf_nodes()[0].taxon.label
-    left = node.child_nodes()[0].leaf_nodes()[0]
-    right = node.child_nodes()[1].leaf_nodes()[0]
-    return f'{left.taxon.label}|{right.taxon.label}'
-
+from gtdb_itol_decorate.gtdb import remove_polyphyletic_suffix
+from gtdb_itol_decorate.newick import get_lca_str
 
 
-def get_phylum_colours(d_phylum_to_lca):
-    """Generate a list of colours for each phylum. Order is inferred through a
-    depth first search, to ensure no clade colours are side by side.
+def get_phylum_colours(phyla: FrozenSet[str]):
+    """Generate a list of colours for each phylum.
     """
     colours = TABLEAU_20
 
+    phyla_mono = {remove_polyphyletic_suffix(x) for x in phyla}
     d_phylum_to_colour = dict()
-    for phylum in d_phylum_to_lca.keys():
+    for phylum in phyla_mono:
         d_phylum_to_colour[phylum] = colours[len(d_phylum_to_colour) % len(colours)]
 
     # Generate a colour palette for each phylum (increasing brightness)
     out = dict()
-    for phylum, colour in d_phylum_to_colour.items():
+    for phylum in phyla:
+        colour = d_phylum_to_colour[remove_polyphyletic_suffix(phylum)]
         cur_pal = sns.light_palette(colour, 6, reverse=True)
         out[phylum] = [rgb_to_hex(*[round(y * 255) for y in x]) for x in cur_pal]
     return out
@@ -62,7 +50,6 @@ def get_phylum_to_lca(tree: dendropy.Tree):
     return out
 
 
-
 def write_color_datastrip(d_phylum_to_lca, d_phylum_palette, path):
     file = iTolDatasetColorStripFile(path, 'Phylum Labels', '#000000', strip_width=100, show_internal=True)
 
@@ -83,7 +70,8 @@ def get_internal_nodes_with_labels(tree: dendropy.Tree):
             out[';'.join(node.tax_label)].append(get_lca_str(node))
     return out
 
-def write_internal_node_labels( d_label_to_lca,path: Path):
+
+def write_internal_node_labels(d_label_to_lca, path: Path):
     file = iTolLabelFile(path)
 
     # 1. Add internal node labels (trivial)
@@ -93,7 +81,6 @@ def write_internal_node_labels( d_label_to_lca,path: Path):
     file.write()
 
 
-
 def write_tree_colours(tree, d_taxon_to_phylum, path, d_color_palette):
     file = iTolTreeColorsFile(path)
 
@@ -122,7 +109,6 @@ def write_tree_colours(tree, d_taxon_to_phylum, path, d_color_palette):
     file.write()
 
 
-
 def write_collapse_file(d_int_label_to_lca, path, taxon_prefix):
     file = iTolCollapseFile(path)
     for taxon, lst_lr_nodes in d_int_label_to_lca.items():
@@ -132,6 +118,7 @@ def write_collapse_file(d_int_label_to_lca, path, taxon_prefix):
                     file.insert(lr_node)
     file.write()
 
+
 def write_popup_file(tree: dendropy.Tree, d_tax: Dict[str, Taxonomy], path: Path):
     file = iTolPopupFile(path)
 
@@ -153,4 +140,4 @@ def write_popup_file(tree: dendropy.Tree, d_tax: Dict[str, Taxonomy], path: Path
         file.insert(gid, row.s.value, ''.join(lines))
 
     file.write()
-    return
+    return

gtdb_itol_decorate/newick.py

@@ -1,6 +1,7 @@
-import dendropy
 from collections import Counter, deque, defaultdict
 
+import dendropy
+
 from gtdb_itol_decorate.util import canonical_gid, log, is_float
 
 
@@ -25,6 +26,7 @@ def load_newick_to_tree(path: str) -> dendropy.Tree:
                                          rooting='force-rooted',
                                          preserve_underscores=True)
 
+
 def validate_dendropy_namespace(taxa):
     taxa_count = Counter(taxa)
     duplicates = {k: v for k, v in taxa_count.items() if v > 1}
@@ -33,12 +35,23 @@ def validate_dendropy_namespace(taxa):
     return
 
 
+def get_lca_str(node: dendropy.Node):
+    if node.is_leaf():
+        return node.taxon.label
+    if len(node.child_nodes()) < 2:
+        return node.child_nodes()[0].leaf_nodes()[0].taxon.label
+    left = node.child_nodes()[0].leaf_nodes()[0]
+    right = node.child_nodes()[1].leaf_nodes()[0]
+    return f'{left.taxon.label}|{right.taxon.label}'
+
+
 def get_canonical_mapping(gids):
     out = dict()
     for gid in gids:
         out[canonical_gid(gid)] = gid
     return out
 
+
 def validate_sets(newick_gids, tax_gids):
     if newick_gids != tax_gids:
         log('The following genomes are in the newick file but not in the taxonomy file:')
@@ -85,6 +98,7 @@ def parse_label(label):
 
     return support, taxon, auxiliary_info
 
+
 def strip_tree_labels(tree: dendropy.Tree):
     for node in tree.preorder_node_iter():
         if node.is_leaf():
@@ -94,6 +108,7 @@ def strip_tree_labels(tree: dendropy.Tree):
             node.label = str(support)
     return
 
+
 def get_node_depth(tree: dendropy.Tree):
     """Return the depth of every node in a tree."""
     depth_to_nodes = defaultdict(list)
@@ -126,7 +141,7 @@ def set_node_desc_taxa(tree: dendropy.Tree):
 
 def set_taxon_label_for_internal_nodes(tree: dendropy.Tree, d_tax):
     # Find the highest internal node that forms a monophyletic group
-
+    taxon_to_lca = defaultdict(list)
     ranks = ('d', 'p', 'c', 'o', 'f', 'g', 's')
 
     queue = deque([(tree.seed_node, 'd')])
@@ -139,15 +154,17 @@ def set_taxon_label_for_internal_nodes(tree: dendropy.Tree, d_tax):
         # Check if this forms a monophyletic group at the target taxon
         desc_taxa = {getattr(d_tax[x], target_rank).value for x in node.desc_taxa}
         if len(desc_taxa) == 1:
-            node.tax_label.append(list(desc_taxa)[0])
+            cur_taxon = desc_taxa.pop()
+            node.tax_label.append(cur_taxon)
+            taxon_to_lca[cur_taxon].append(get_lca_str(node))
 
             # Re-queue this node at the front, in-case it can be extended
             if target_rank != 's':
-                next_rank = ranks[ranks.index(target_rank)+1]
+                next_rank = ranks[ranks.index(target_rank) + 1]
                 queue.appendleft((node, next_rank))
 
         # Otherwise, keep going down
         else:
             for child in node.child_nodes():
                 queue.append((child, target_rank))
-    return
+    return taxon_to_lca

gtdb_itol_decorate/util.py

@@ -8,6 +8,7 @@ def log(msg: str, title=False):
     if title:
         print('-' * 80)
 
+
 def canonical_gid(gid):
     """Get canonical form of NCBI genome accession.
 
@@ -30,7 +31,6 @@ def canonical_gid(gid):
     return gid
 
 
-
 def is_float(s):
     """Check if a string can be converted to a float.
 
@@ -50,4 +50,4 @@ def is_float(s):
     except ValueError:
         return False
 
-    return True
+    return True