Allow newicks to be output without branch lengths

Author: hyanwong

python/CHANGELOG.rst

@@ -4,6 +4,9 @@
 
 **Features**
 
+- ``Tree.newick()`` now has extra option ``include_branch_lengths`` to allow branch
+  lengths to be omitted (:user:`hyanwong`, :pr:`931`).
+
 - Added ``TableCollection.equals``, a method that compares two table 
   collections with the options to ignore top-level metadata/schema or
   provenance tables. (:user:`mufernando`, :issue:`896`, :pr:`897`).
@@ -14,6 +17,7 @@
 **Breaking changes**
 
 - The argument to ``ts.dump`` and ``tskit.load`` has been renamed `file` from `path`.
+- All arguments to ``Tree.newick()`` except precision are now keyword-only.
 
 --------------------
 [0.3.2] - 2020-09-29

python/tests/test_highlevel.py

@@ -2054,6 +2054,28 @@ def test_newick_large_times(self):
                 newick_c = tree.newick(precision=precision)
                 assert newick_c == newick_py
 
+    def test_bifurcating_newick(self):
+        for n_tips in range(2, 6):
+            ts = msprime.simulate(n_tips, random_seed=1)  # msprime trees are binary
+            for tree in ts.trees():
+                base_newick = tree.newick(include_branch_lengths=False).strip(";")
+                for i in range(n_tips):
+                    # Each tip number (i+1) mentioned once
+                    assert base_newick.count(str(i + 1)) == 1
+                # Binary newick trees have 3 chars per extra tip: "(,)"
+                assert len(base_newick) == n_tips + 3 * (n_tips - 1)
+
+    def test_newick_topology_equiv(self):
+        replace_numeric = {ord(x): None for x in "1234567890:."}
+        for ts in get_example_tree_sequences():
+            for tree in ts.trees():
+                if tree.num_roots > 1:
+                    continue
+                plain_newick = tree.newick(node_labels={}, include_branch_lengths=False)
+                newick1 = tree.newick().translate(replace_numeric)
+                newick2 = tree.newick(node_labels={}).translate(replace_numeric)
+                assert newick1 == newick2 == plain_newick
+
     def test_newick_buffer_too_small_bug(self):
         nodes = io.StringIO(
             """\

python/tests/test_phylo_formats.py

@@ -89,10 +89,17 @@ class TestNewick(TreeExamples):
     external Newick parser.
     """
 
-    def verify_newick_topology(self, tree, root=None, node_labels=None):
+    def verify_newick_topology(
+        self, tree, root=None, node_labels=None, include_branch_lengths=True
+    ):
         if root is None:
             root = tree.root
-        ns = tree.newick(precision=16, root=root, node_labels=node_labels)
+        ns = tree.newick(
+            precision=16,
+            root=root,
+            node_labels=node_labels,
+            include_branch_lengths=include_branch_lengths,
+        )
         if node_labels is None:
             leaf_labels = {u: str(u + 1) for u in tree.leaves(root)}
         else:
@@ -104,7 +111,9 @@ def verify_newick_topology(self, tree, root=None, node_labels=None):
             name = leaf_labels[u]
             node = newick_tree.get_node(name)
             while u != root:
-                self.assertAlmostEqual(node.length, tree.branch_length(u))
+                # the newick library encodes absent branch lengths as 0.0
+                branch_length = tree.branch_length(u) if include_branch_lengths else 0.0
+                self.assertAlmostEqual(node.length, branch_length)
                 node = node.ancestor
                 u = tree.parent(u)
             assert node.ancestor is None
@@ -162,6 +171,10 @@ def test_single_node_label(self):
             None for _ in range(len(list(tree.nodes())) - 1)
         ]
 
+    def test_no_lengths(self):
+        t = msprime.simulate(5, random_seed=2).first()
+        self.verify_newick_topology(t, include_branch_lengths=False)
+
 
 class TestNexus(TreeExamples):
     """

python/tskit/trees.py

@@ -2148,10 +2148,10 @@ def nodes(self, root=None, order="preorder"):
                 yield from iterator(u)
 
     # TODO make this a bit less embarrassing by using an iterative method.
-    def __build_newick(self, node, precision, node_labels):
+    def __build_newick(self, *, node, precision, node_labels, include_branch_lengths):
         """
         Simple recursive version of the newick generator used when non-default
-        node labels are needed.
+        node labels are needed, or when branch lengths are omitted
         """
         label = node_labels.get(node, "")
         if self.is_leaf(node):
@@ -2160,12 +2160,26 @@ def __build_newick(self, node, precision, node_labels):
             s = "("
             for child in self.children(node):
                 branch_length = self.branch_length(child)
-                subtree = self.__build_newick(child, precision, node_labels)
-                s += subtree + ":{0:.{1}f},".format(branch_length, precision)
+                subtree = self.__build_newick(
+                    node=child,
+                    precision=precision,
+                    node_labels=node_labels,
+                    include_branch_lengths=include_branch_lengths,
+                )
+                if include_branch_lengths:
+                    subtree += ":{0:.{1}f}".format(branch_length, precision)
+                s += subtree + ","
             s = s[:-1] + f"){label}"
         return s
 
-    def newick(self, precision=14, root=None, node_labels=None):
+    def newick(
+        self,
+        precision=14,  # Should probably be keyword only, left positional for legacy use
+        *,
+        root=None,
+        node_labels=None,
+        include_branch_lengths=True,
+    ):
         """
         Returns a `newick encoding <https://en.wikipedia.org/wiki/Newick_format>`_
         of this tree. If the ``root`` argument is specified, return a representation
@@ -2187,6 +2201,8 @@ def newick(self, precision=14, root=None, node_labels=None):
         :param dict node_labels: If specified, show custom labels for the nodes
             that are present in the map. Any nodes not specified in the map will
             not have a node label.
+        :param include_branch_lengths: If True (default), output branch lengths in the
+            Newick file. If False, only output the topology, without branch lengths.
         :return: A newick representation of this tree.
         :rtype: str
         """
@@ -2198,6 +2214,9 @@ def newick(self, precision=14, root=None, node_labels=None):
                     "newick trees, one for each root."
                 )
             root = self.root
+        if not include_branch_lengths and node_labels is None:
+            # Force the python generator for simplicity, by specifying the default labels
+            node_labels = {i: str(i + 1) for i in self.leaves()}
         if node_labels is None:
             root_time = max(1, self.time(root))
             max_label_size = math.ceil(math.log10(self.tree_sequence.num_nodes))
@@ -2210,7 +2229,15 @@ def newick(self, precision=14, root=None, node_labels=None):
             )
             s = s.decode()
         else:
-            return self.__build_newick(root, precision, node_labels) + ";"
+            s = (
+                self.__build_newick(
+                    node=root,
+                    precision=precision,
+                    node_labels=node_labels,
+                    include_branch_lengths=include_branch_lengths,
+                )
+                + ";"
+            )
         return s
 
     def as_dict_of_dicts(self):
@@ -4489,7 +4516,7 @@ def to_nexus(self, precision=14):
         for tree in self.trees():
             start_interval = "{0:.{1}f}".format(tree.interval.left, precision)
             end_interval = "{0:.{1}f}".format(tree.interval.right, precision)
-            newick = tree.newick(precision, node_labels=node_labels)
+            newick = tree.newick(precision=precision, node_labels=node_labels)
             s += f"\tTREE tree{start_interval}_{end_interval} = {newick}\n"
 
         s += "END;\n"