Merge 9683507 into 4e76bcd

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Version 0.3.2
 * Added `balance_boxplot` and `balance_barplot` to make interpretation balance partitions easier.
+* Added `assign_ids` command to allow for ids to be added manually.
 
 ## Version 0.3.0
 * Added q2 support for linear regression and linear mixed effects models [#98](https://github.com/biocore/gneiss/pull/98)

gneiss/cluster/_cluster.py

@@ -5,6 +5,7 @@
 #
 # The full license is in the file COPYING.txt, distributed with this software.
 # ----------------------------------------------------------------------------
+import uuid
 import pandas as pd
 import numpy as np
 import skbio
@@ -16,6 +17,7 @@
 from gneiss.plugin_setup import plugin
 from gneiss.cluster._pba import correlation_linkage, gradient_linkage
 from gneiss.sort import gradient_sort, mean_niche_estimator
+from gneiss.util import rename_internal_nodes
 
 
 def correlation_clustering(table: pd.DataFrame) -> skbio.TreeNode:
@@ -121,3 +123,30 @@ def gradient_clustering(table: pd.DataFrame,
                  'This method is primarily used to sort the table to reveal '
                  'the underlying block-like structures.')
 )
+
+
+def assign_ids(tree: skbio.TreeNode) -> skbio.TreeNode:
+
+    t = tree.copy()
+    t.bifurcate()
+    ids = ['%sL-%s' % (i, uuid.uuid4())
+           for i, n in enumerate(t.levelorder(include_self=True))
+           if not n.is_tip()]
+    t = rename_internal_nodes(t, names=ids)
+    return t
+
+
+plugin.methods.register_function(
+    function=assign_ids,
+    inputs={'tree': Phylogeny[Rooted]},
+    outputs=[('tree', Phylogeny[Rooted])],
+    name='Assigns ids on internal nodes in the tree.',
+    input_descriptions={
+        'tree': ('The input tree with potential missing ids.')},
+    parameters={},
+    output_descriptions={
+        'tree': ('A tree with uniquely identifying ids.')},
+    description=('Assigns UUIDs to uniquely identify internal nodes '
+                 'in the tree.  Also corrects for polytomies to create '
+                 'strictly bifurcating trees.')
+)

gneiss/cluster/tests/test_cluster.py

@@ -39,7 +39,6 @@ def test_gradient_artifact(self):
         res = gradient_clustering(in_table, in_metadata.get_category('x'))
         res_clust = res.clustering._view(TreeNode)
         exp_str = '((o1:0.5,o2:0.5)y1:0.5,(o3:0.5,o4:0.5)y2:0.5)y0;\n'
-        print(str(res_clust))
         self.assertEqual(exp_str, str(res_clust))
 
     def test_gradient_artifact_weighted(self):
@@ -59,6 +58,28 @@ def test_gradient_artifact_weighted(self):
 
         self.assertNotEqual(str(res_clust_uw), str(res_clust_w))
 
+    def test_assign_ids(self):
+        from qiime2.plugins.gneiss.methods import assign_ids
+        tree_f = get_data_path("tree.qza")
+        tree = qiime2.Artifact.load(tree_f)
+        out_tree = assign_ids(tree)
+        res_t = out_tree.tree._view(TreeNode)
+        for n in res_t.levelorder(include_self=True):
+            self.assertTrue(n.name is not None)
+
+    def test_assign_ids_polytomy(self):
+        from qiime2.plugins.gneiss.methods import assign_ids
+        tree_f = get_data_path("polytomy.qza")
+        tree = qiime2.Artifact.load(tree_f)
+        out_tree = assign_ids(tree)
+        res_t = out_tree.tree._view(TreeNode)
+        res_nontips = []
+        for n in res_t.levelorder(include_self=True):
+            self.assertTrue(n.name is not None)
+            if not n.is_tip():
+                res_nontips.append(n.name)
+        self.assertEqual(len(res_nontips), 4)
+
 
 if __name__ == '__main__':
     unittest.main()