Merge branch 'master' of https://github.com/biocore/gneiss into radia…

…lplot

gneiss/plot/_dendrogram.py

@@ -5,29 +5,7 @@
 #
 # The full license is in the file COPYING.txt, distributed with this software.
 # ----------------------------------------------------------------------------
-
-"""Drawing trees.
-
-Draws horizontal trees where the vertical spacing between taxa is
-constant.  Since dy is fixed dendrograms can be either:
- - square:     dx = distance
- - not square: dx = max(0, sqrt(distance**2 - dy**2))
-
-Also draws basic unrooted trees.
-
-For drawing trees use either:
- - UnrootedDendrogram
-
-Note: This is directly ported from pycogent.
-"""
-
-# Future:
-#  - font styles
-#  - orientation switch
-# Layout gets more complicated for rooted tree styles if dy is allowed to vary,
-# and constant-y is suitable for placing alongside a sequence alignment anyway.
-import abc
-from collections import namedtuple
+from skbio import TreeNode
 import pandas as pd
 import numpy
 from skbio import TreeNode
@@ -39,36 +17,34 @@ def _sign(x):
 
 
 class Dendrogram(TreeNode):
-    # One of these for each tree edge.  Extra attributes:
-    #    depth - distance from root to bottom of edge
-    #    height - max distance from a decendant leaf to top of edge
-    #    width - number of decendant leaves
-    # note these are named tree-wise, not geometricaly, so think
-    # of a vertical tree (for this part anyway)
-    #
-    #   x1, y1, x2, y2 - coordinates
-    # these are horizontal / vertical as you would expect
-    #
-    # The algorithm is split into 4 passes over the tree for easier
-    # code reuse - vertical drawing, new tree styles, new graphics
-    # libraries etc.
-
+    """ Stores data to be plotted as a dendrogram.
+
+    A `Dendrogram` object is represents a tree in addition to the
+    key information required to create a tree layout prior to
+    visualization.  No layouts are specified within this class,
+    since this serves as a super class for different tree layouts.
+
+    Parameters
+    ----------
+    use_lengths: bool
+        Specifies if the branch lengths should be included in the
+        resulting visualization (default True).
+
+    Attributes
+    ----------
+    length
+    """
     aspect_distorts_lengths = True
 
     def __init__(self, use_lengths=True, **kwargs):
         """ Constructs a Dendrogram object for visualization.
 
-        Parameters
-        ----------
-        use_lengths: bool
-            Specifies if the branch lengths should be included in the
-            resulting visualization (default True).
         """
         super().__init__(**kwargs)
         self.use_lengths_default = use_lengths
 
     def _cache_ntips(self):
-        for n in self.postorder(include_self=True):
+        for n in self.postorder():
             if n.is_tip():
                 n.leafcount = 1
             else:
@@ -113,19 +89,20 @@ def coords(self, height, width):
         -------
         pd.DataFrame
             index : str
-                name of node
+                Name of node.
             x : float
-                x-coordinate of point
+                x-coordinate of node.
             y : float
-                y-coordinate of point
-            left : str
-                name of left child node
-            right : str
-                name of right child node
+                y-coordinate of node.
+            child(i) : str
+                Name of ith child node in that specific node.
+                in the tree.
+            is_tip : str
+                Specifies if the node is a tip in the treee.
         """
         self.rescale(width, height)
         result = {}
-        for node in self.postorder(include_self=True):
+        for node in self.postorder():
             children = {'child%d' % i: n.name
                         for i, n in enumerate(node.children)}
             coords = {'x': node.x2, 'y': node.y2}
@@ -134,18 +111,31 @@ def coords(self, height, width):
         result = pd.DataFrame(result).T
 
         # reorder so that x and y are first
-        cols = list(result)
-        cols.insert(0, cols.pop(cols.index('y')))
-        cols.insert(0, cols.pop(cols.index('x')))
-        result = result.ix[:, cols]
-        return result
+        cols = ['x', 'y'] + sorted(list(set(result.columns) - set(['x', 'y'])))
+        return result.loc[:, cols]
 
     @abc.abstractmethod
     def rescale(self, width, height):
         pass
 
 
 class UnrootedDendrogram(Dendrogram):
+    """ Stores data to be plotted as an unrooted dendrogram.
+
+    A `UnrootedDendrogram` object is represents a tree in addition to the
+    key information required to create a radial tree layout prior to
+    visualization.
+
+    Parameters
+    ----------
+    use_lengths: bool
+        Specifies if the branch lengths should be included in the
+        resulting visualization (default True).
+
+    Attributes
+    ----------
+    length
+    """
     aspect_distorts_lengths = True
 
     def __init__(self, **kwargs):
@@ -172,14 +162,18 @@ def from_tree(cls, tree, use_lengths=True):
         -------
         UnrootedDendrogram
         """
-        for n in tree.postorder(include_self=True):
+        for n in tree.postorder():
             n.__class__ = UnrootedDendrogram
 
         tree.update_geometry(use_lengths)
         return tree
 
     def rescale(self, width, height):
-        """ Find best scaling factor for fitting the tree in the canvas
+        """ Find best scaling factor for fitting the tree in the dimensions
+        specified by width and height.
+
+        This method will find the best orientation and scaling possible
+        to fit the tree within the dimensions specified by width and height.
 
         Parameters
         ----------
@@ -195,24 +189,25 @@ def rescale(self, width, height):
 
         Notes
         -----
-        This function has a little bit of recursion.  This will
-        need to be refactored to remove the recursion.
         """
-
-        angle = 2*numpy.pi / self.leafcount
+        angle = (2 * numpy.pi) / self._n_tips
         # this loop is a horrible brute force hack
         # there are better (but complex) ways to find
         # the best rotation of the tree to fit the display.
         best_scale = 0
         for i in range(60):
             direction = i / 60.0 * numpy.pi
-            points = self.update_coordinates(1.0, 0, 0, direction, angle)
-            xs = [x for (x, y) in points]
-            ys = [y for (x, y) in points]
+            # TODO:
+            # This function has a little bit of recursion.  This will
+            # need to be refactored to remove the recursion.
 
+            points = self.update_coordinates(1.0, 0, 0, direction, angle)
+            xs, ys = zip(*points)
             # double check that the tree fits within the margins
             scale = min(float(width) / (max(xs) - min(xs)),
                         float(height) / (max(ys) - min(ys)))
+            # TODO: This margin seems a bit arbituary.
+            # will need to investigate.
             scale *= 0.95  # extra margin for labels
             if scale > best_scale:
                 best_scale = scale
@@ -226,6 +221,12 @@ def rescale(self, width, height):
     def update_coordinates(self, s, x1, y1, a, da):
         """ Update x, y coordinates of tree nodes in canvas.
 
+        `update_coordinates` will recursively updating the
+        plotting parameters for all of the nodes within the tree.
+        This can be applied when the tree becomes modified (i.e. pruning
+        or collapsing) and the resulting coordinates need to be modified
+        to reflect the changes to the tree structure.
+
         Parameters
         ----------
         s : float
@@ -235,33 +236,36 @@ def update_coordinates(self, s, x1, y1, a, da):
         y1 : float
             y midpoint
         a : float
-            direction (degrees)
-        da : float
             angle (degrees)
+        da : float
+            angle resolution (degrees)
 
         Returns
         -------
         points : list of tuple
-            2D coordinates of all of the notes
+            2D coordinates of all of the nodes.
 
         Notes
         -----
         This function has a little bit of recursion.  This will
         need to be refactored to remove the recursion.
         """
         # Constant angle algorithm.  Should add maximum daylight step.
-        x2 = x1+self.length*s*numpy.sin(a)
-        y2 = y1+self.length*s*numpy.cos(a)
+        x2 = x1 + self.length * s * numpy.sin(a)
+        y2 = y1 + self.length * s * numpy.cos(a)
         (self.x1, self.y1, self.x2, self.y2, self.angle) = (x1, y1, x2, y2, a)
         # TODO: Add functionality that allows for collapsing of nodes
         a = a - self.leafcount * da / 2
         if self.is_tip():
             points = [(x2, y2)]
         else:
             points = []
-            # recurse down the tree
+            # TODO:
+            # This function has a little bit of recursion.  This will
+            # need to be refactored to remove the recursion.
             for child in self.children:
-                ca = child.leafcount * da
+                # calculate the arc that covers the subtree.
+                ca = child._n_tips * da
                 points += child.update_coordinates(s, x2, y2, a+ca/2, da)
                 a += ca
         return points

gneiss/plot/tests/test_blobtree.py

@@ -43,5 +43,6 @@ def test_no_colors(self):
         self.assertTrue(os.path.exists(self.fname))
         self.assertTrue(os.path.getsize(self.fname) > 0)
 
+
 if __name__ == "__main__":
     unittest.main()

gneiss/plot/tests/test_dendrogram.py

@@ -109,5 +109,6 @@ def test_update_coordinates(self):
         res = pd.DataFrame(t.update_coordinates(1, 0, 0, 2, 1))
         pdt.assert_frame_equal(res, exp, check_less_precise=True)
 
+
 if __name__ == "__main__":
     unittest.main()

gneiss/plot/tests/test_heatmap.py

@@ -24,5 +24,6 @@ def test_not_fail(self):
         self.assertTrue(os.path.exists(self.fname))
         self.assertTrue(os.path.getsize(self.fname) > 0)
 
+
 if __name__ == "__main__":
     unittest.main()