WIP : Rename clades

gneiss/tests/test_util.py

@@ -12,7 +12,7 @@
 import pandas.util.testing as pdt
 from skbio import TreeNode
 from gneiss.util import match, match_tips, design_formula
-from gneiss.util import (rename_internal_nodes,
+from gneiss.util import (rename_internal_nodes, rename_clades,
                          _type_cast_to_float, block_diagonal, band_diagonal,
                          split_balance, check_internal_nodes,
                          _xarray_match_tips)
@@ -505,6 +505,27 @@ def test_formula(self):
 
 class TestUtil(unittest.TestCase):
 
+    def test_rename_clades(self):
+        tree = TreeNode.read([u"(((a,b), c)Firmicutes,d)r;"])
+        exp_tree = TreeNode.read([u"(((a,b)clade0, c)Firmicutes,d)r;"])
+        res_tree = rename_clades(tree)
+        self.assertEqual(str(exp_tree), str(res_tree))
+
+    def test_rename_clades_immutable(self):
+        tree = TreeNode.read([u"(((a,b)Clostridia, c),d)r;"])
+        rename_clades(tree)
+        self.assertEqual(str(tree), "(((a,b)Clostridia,c),d)r;\n")
+
+    def test_rename_clades_mutable(self):
+        tree = TreeNode.read([u"(((a,b)Clostridia, c),d)r;"])
+        rename_clades(tree, inplace=True)
+        self.assertEqual(str(tree), "(((a,b)Clostridia,c)clade0,d)r;\n")
+
+    def test_rename_clades_dup(self):
+        tree = TreeNode.read([u"(((a,b)Clostridia, c)Clostridia,d)r;"])
+        rename_clades(tree, inplace=True)
+        self.assertEqual(str(tree), "(((a,b)Clostridia1,c)Clostridia,d)r;\n")
+
     def test_rename_internal_nodes(self):
         tree = TreeNode.read([u"(((a,b), c),d)r;"])
         exp_tree = TreeNode.read([u"(((a,b)y2, c)y1,d)y0;"])

gneiss/util.py

@@ -29,6 +29,7 @@
 import warnings
 import numpy as np
 from skbio.stats.composition import closure
+from collections import Counter
 import pandas as pd
 from patsy import dmatrix
 import biom
@@ -357,6 +358,51 @@ def rename_internal_nodes(tree, names=None, inplace=False):
     return _tree
 
 
+def rename_clades(tree, inplace=False):
+    """ Names the internal according to level ordering.
+
+    The tree will be traversed in level order (i.e. top-down, left to right)
+    and missing labels will be replaced with `clade_i`, where `i` is the
+    level ordering position. Duplicate internal node names will also be
+    renamed according to level ordering position.
+
+
+    Parameters
+    ----------
+    tree : skbio.TreeNode
+        Tree object where the leafs correspond to the features.
+    inplace : bool, optional
+        Specifies if the operation should be done on the original tree or not.
+
+    Returns
+    -------
+    skbio.TreeNode
+       Tree with renamed internal nodes.
+
+    Raises
+    ------
+    ValueError:
+        Raised if `tree` and `name` have incompatible sizes.
+    """
+    if inplace:
+        _tree = tree
+    else:
+        _tree = tree.copy()
+
+    i = 0
+    obs_clades = Counter()
+    for n in _tree.levelorder(include_self=True):
+        if n.name is None:
+            n.name = f"clade{i}"
+            i += 1
+        elif n.name in obs_clades.keys():
+            n.name = f"{n.name}{obs_clades[n.name]}"
+            obs_clades[n.name] += 1
+        else:
+            obs_clades[n.name] += 1
+    return _tree
+
+
 def _type_cast_to_float(df):
     """ Attempt to cast all of the values in dataframe to float.