Merge pull request #171 from mortonjt/issue-143

Fixing highlight issue in heatmap
biocore · Apr 27, 2017 · 4e76bcd · 4e76bcd
2 parents 510cb96 + 91a5a1c
commit 4e76bcd
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 1 deletion.
diff --git a/gneiss/plot/_heatmap.py b/gneiss/plot/_heatmap.py
@@ -169,6 +169,7 @@ def _plot_highlights_dendrogram(ax_highlights, table, t, highlights):
     hcoords = []
     for i, n in enumerate(highlights.index):
         node = t.find(n)
+
         k, l, r = node._k, node._l, node._r
 
         ax_highlights.add_patch(

diff --git a/gneiss/plot/_plot.py b/gneiss/plot/_plot.py
@@ -451,7 +451,8 @@ def dendrogram_heatmap(output_dir: str, table: pd.DataFrame,
                        tree: TreeNode, metadata: MetadataCategory,
                        ndim=10, method='clr', color_map='viridis'):
 
-    nodes = [n.name for n in tree.levelorder()]
+    nodes = [n.name for n in tree.levelorder() if not n.is_tip()]
+
     nlen = min(ndim, len(nodes))
     highlights = pd.DataFrame([['#00FF00', '#FF0000']] * nlen,
                               index=nodes[:nlen])

diff --git a/gneiss/plot/tests/test_plot.py b/gneiss/plot/tests/test_plot.py
@@ -317,6 +317,36 @@ def test_visualization(self):
             self.assertIn('<h1>Dendrogram heatmap</h1>',
                           html)
 
+    def test_visualization_small(self):
+        # tests the scenario where ndim > number of tips
+        np.random.seed(0)
+        num_otus = 11  # otus
+        table = pd.DataFrame(np.random.random((num_otus, 5)),
+                             index=np.arange(num_otus).astype(np.str)).T
+
+        x = np.random.rand(num_otus)
+        dm = DistanceMatrix.from_iterable(x, lambda x, y: np.abs(x-y))
+        lm = ward(dm.condensed_form())
+        t = TreeNode.from_linkage_matrix(lm, np.arange(len(x)).astype(np.str))
+
+        for i, n in enumerate(t.postorder()):
+            if not n.is_tip():
+                n.name = "y%d" % i
+            n.length = np.random.rand()*3
+
+        md = MetadataCategory(
+            pd.Series(['a', 'a', 'a', 'b', 'b']))
+
+        dendrogram_heatmap(self.results, table, t, md)
+
+        index_fp = os.path.join(self.results, 'index.html')
+        self.assertTrue(os.path.exists(index_fp))
+
+        with open(index_fp, 'r') as fh:
+            html = fh.read()
+            self.assertIn('<h1>Dendrogram heatmap</h1>',
+                          html)
+
 
 if __name__ == "__main__":
     unittest.main()