Merge pull request #407 from jeromekelleher/svg_tree_height

Fix issues with max_tree_height with rank scale.

Author: jeromekelleher

python/tests/test_drawing.py

@@ -1163,6 +1163,47 @@ def test_one_mutation_label_colour(self):
         self.verify_basic_svg(svg)
         self.assertEqual(svg.count('stroke="{}"'.format(colour)), 1)
 
+    def test_bad_tree_height_scale(self):
+        t = self.get_binary_tree()
+        for bad_scale in ["te", "asdf", "", [], b'23']:
+            with self.assertRaises(ValueError):
+                t.draw_svg(tree_height_scale=bad_scale)
+
+    def test_bad_max_tree_height(self):
+        t = self.get_binary_tree()
+        for bad_height in ["te", "asdf", "", [], b'23']:
+            with self.assertRaises(ValueError):
+                t.draw_svg(max_tree_height=bad_height)
+
+    def test_height_scale_time_and_max_tree_height(self):
+        ts = msprime.simulate(5, recombination_rate=2, random_seed=2)
+        t = ts.first()
+        # The default should be the same as tree.
+        svg1 = t.draw_svg(max_tree_height="tree")
+        self.verify_basic_svg(svg1)
+        svg2 = t.draw_svg()
+        self.assertEqual(svg1, svg2)
+        svg3 = t.draw_svg(max_tree_height="ts")
+        self.assertNotEqual(svg1, svg3)
+        svg4 = t.draw_svg(max_tree_height=max(ts.tables.nodes.time))
+        self.assertEqual(svg3, svg4)
+
+    def test_height_scale_rank_and_max_tree_height(self):
+        # Make sure the rank height scale and max_tree_height interact properly.
+        ts = msprime.simulate(5, recombination_rate=2, random_seed=2)
+        t = ts.first()
+        # The default should be the same as tree.
+        svg1 = t.draw_svg(max_tree_height="tree", tree_height_scale="rank")
+        self.verify_basic_svg(svg1)
+        svg2 = t.draw_svg(tree_height_scale="rank")
+        self.assertEqual(svg1, svg2)
+        svg3 = t.draw_svg("tmp.svg", max_tree_height="ts", tree_height_scale="rank")
+        self.assertNotEqual(svg1, svg3)
+        self.verify_basic_svg(svg3)
+        # Numeric max tree height not supported for rank scale.
+        with self.assertRaises(ValueError):
+            t.draw_svg(max_tree_height=2, tree_height_scale="rank")
+
     #
     # TODO: update the tests below here to check the new SVG based interface.
     #

python/tskit/drawing.py

@@ -24,11 +24,50 @@
 Module responsible for visualisations.
 """
 import collections
-from _tskit import NULL
+import numbers
 
 import svgwrite
 import numpy as np
 
+from _tskit import NULL
+
+LEFT = "left"
+RIGHT = "right"
+TOP = "top"
+BOTTOM = "bottom"
+
+
+def check_orientation(orientation):
+    if orientation is None:
+        orientation = TOP
+    else:
+        orientation = orientation.lower()
+        orientations = [LEFT, RIGHT, TOP, BOTTOM]
+        if orientation not in orientations:
+            raise ValueError(
+                "Unknown orientiation: choose from {}".format(orientations))
+    return orientation
+
+
+def check_max_tree_height(max_tree_height, allow_numeric=True):
+    if max_tree_height is None:
+        max_tree_height = "tree"
+    is_numeric = isinstance(max_tree_height, numbers.Real)
+    if max_tree_height not in ["tree", "ts"] and not allow_numeric:
+        raise ValueError("max_tree_height must be 'tree' or 'ts'")
+    if max_tree_height not in ["tree", "ts"] and (allow_numeric and not is_numeric):
+        raise ValueError(
+            "max_tree_height must be a numeric value or one of 'tree' or 'ts'")
+    return max_tree_height
+
+
+def check_tree_height_scale(tree_height_scale):
+    if tree_height_scale is None:
+        tree_height_scale = "time"
+    if tree_height_scale not in ["time", "log_time", "rank"]:
+        raise ValueError("tree_height_scale must be 'time', 'log_time' or 'rank'")
+    return tree_height_scale
+
 
 def check_format(format):
     if format is None:
@@ -264,28 +303,39 @@ def setup_drawing(self):
         self.mutation_right_labels = self.mutation_labels.add(dwg.g(text_anchor="end"))
 
     def assign_y_coordinates(self, tree_height_scale, max_tree_height):
+        tree_height_scale = check_tree_height_scale(tree_height_scale)
+        max_tree_height = check_max_tree_height(
+            max_tree_height, tree_height_scale != "rank")
         ts = self.tree.tree_sequence
         node_time = ts.tables.nodes.time
-        if tree_height_scale in [None, "time", "log_time"]:
-            if max_tree_height in [None, "tree"]:
+
+        if tree_height_scale == "rank":
+            assert tree_height_scale == "rank"
+            if max_tree_height == "tree":
+                # We only rank the times within the tree in this case.
+                t = np.zeros_like(node_time) + node_time[self.tree.left_root]
+                for u in self.tree.nodes():
+                    t[u] = node_time[u]
+                node_time = t
+            depth = {t: 2 * j for j, t in enumerate(np.unique(node_time))}
+            node_height = [depth[node_time[u]] for u in range(ts.num_nodes)]
+            max_tree_height = max(depth.values())
+        else:
+            assert tree_height_scale in ["time", "log_time"]
+            if max_tree_height == "tree":
                 max_tree_height = max(self.tree.time(root) for root in self.tree.roots)
             elif max_tree_height == "ts":
                 max_tree_height = ts.max_root_time
-        if tree_height_scale == "log_time":
-            # add 1 so that don't reach log(0) = -inf error.
-            # just shifts entire timeset by 1 year so shouldn't affect anything
-            node_height = np.log(ts.tables.nodes.time + 1)
-        elif tree_height_scale in [None, "time"]:
-            node_height = node_time
-        else:
-            if tree_height_scale != "rank":
-                raise ValueError(
-                    "Only 'time', 'log_time', "
-                    "and 'rank' are supported for tree_height_scale")
-            depth = {t: 2 * j for j, t in enumerate(np.unique(node_time))}
-            node_height = [depth[node_time[u]] for u in range(ts.num_nodes)]
-            if max_tree_height is None:
-                max_tree_height = max(depth.values())
+
+            if tree_height_scale == "log_time":
+                # add 1 so that don't reach log(0) = -inf error.
+                # just shifts entire timeset by 1 year so shouldn't affect anything
+                node_height = np.log(ts.tables.nodes.time + 1)
+            elif tree_height_scale == "time":
+                node_height = node_time
+
+        assert float(max_tree_height) == max_tree_height
+
         # In pathological cases, all the roots are at 0
         if max_tree_height == 0:
             max_tree_height = 1
@@ -466,32 +516,6 @@ def __str__(self):
         return "".join(self.canvas.reshape(self.width * self.height))
 
 
-LEFT = "left"
-RIGHT = "right"
-TOP = "top"
-BOTTOM = "bottom"
-
-
-def check_orientation(orientation):
-    if orientation is None:
-        orientation = TOP
-    else:
-        orientation = orientation.lower()
-        orientations = [LEFT, RIGHT, TOP, BOTTOM]
-        if orientation not in orientations:
-            raise ValueError(
-                "Unknown orientiation: choose from {}".format(orientations))
-    return orientation
-
-
-def check_max_tree_height(max_tree_height):
-    if max_tree_height is None:
-        max_tree_height = "tree"
-    if max_tree_height not in ["tree", "ts"]:
-        raise ValueError("max_tree_height must be 'tree' or 'ts'")
-    return max_tree_height
-
-
 def to_np_unicode(string):
     """
     Converts the specified string to a numpy unicode array.
@@ -574,7 +598,8 @@ def __init__(
             self, tree, node_labels=None, max_tree_height=None, use_ascii=False,
             orientation=None):
         self.tree = tree
-        self.max_tree_height = check_max_tree_height(max_tree_height)
+        self.max_tree_height = check_max_tree_height(
+            max_tree_height, allow_numeric=False)
         self.use_ascii = use_ascii
         self.orientation = check_orientation(orientation)
         self.horizontal_line_char = '━'