Merge pull request #194 from jeetsukumaran/polytomize-root

Implement, test polytomize_root

src/dendropy/datamodel/treemodel/_tree.py

@@ -2363,6 +2363,21 @@ def _set_is_unrooted(self, val):
 
     is_unrooted = property(_get_is_unrooted, _set_is_unrooted)
 
+    def polytomize_root(self, set_as_unrooted_tree=True):
+        """Works recursively to set root as a degree-3 node.
+
+        If ``self`` has two children and at least one of them is an internal
+        node, then it will be converted to an out-degree three node (with the
+        edge length added as needed).
+
+        Similar to :meth:`Tree.collapse_basal_bifurcation()`, but works
+        recursively to set root as a degree-3 node.
+        """
+        if self.seed_node is not None:
+            self.seed_node._convert_node_to_root_polytomy()
+        if set_as_unrooted_tree:
+            self.is_rooted = False
+
     def collapse_basal_bifurcation(self, set_as_unrooted_tree=True):
         "Converts a degree-2 node at the root to a degree-3 node."
         seed_node = self.seed_node

tests/unittests/test_tree_operations_and_manipulations.py

@@ -938,6 +938,24 @@ def test_special_case2(self):
 
 class TreeRestructuring(dendropytest.ExtendedTestCase):
 
+    def test_polytomize_root(self):
+        t = dendropy.Tree.get(data="(B:0.2,((C:0.3,(D:0.4,Q:9.3)))E:0.5);", schema="newick")
+        t.polytomize_root(set_as_unrooted_tree=False)
+        self.assertEqual(t.as_string("newick").strip(), "(B:0.7,C:0.3,(D:0.4,Q:9.3));")
+
+        t = dendropy.Tree.get(data="(((A,((B,C)))));", schema="newick")
+        t.polytomize_root(set_as_unrooted_tree=False)
+        self.assertEqual(t.as_string("newick").strip(), "(A,B,C);")
+
+        newick_string = "(B:0.7,C:0.3,(D:0.4,Q:9.3):0.3):0.3;"
+        t = dendropy.Tree.get(data=newick_string, schema="newick")
+        t.polytomize_root(set_as_unrooted_tree=False)
+        self.assertEqual(t.as_string("newick").strip(), newick_string)
+
+        t = dendropy.Tree.get(data=newick_string, schema="newick")
+        t.polytomize_root(set_as_unrooted_tree=True)
+        self.assertEqual(t.is_rooted, False)
+
     def test_collapse_basal_bifurcation(self):
         self.assertFalse(self.fail_incomplete_tests())