Merge pull request mwaskom#551 from mwaskom/default_linewidths

mwaskom · mwaskom · commit d0956f3fb6db · 2015-05-10T11:59:11.000-07:00
Some changes to help heatmap and clustermap plot large matrices
diff --git a/doc/releases/v0.6.0.txt b/doc/releases/v0.6.0.txt
@@ -63,10 +63,14 @@ Other additions
 
 - Added a ``savefig`` method to :class:`JointGrid` that defaults to a tight bounding box to make it easier to save figures using this class.
 
+- Changed the default ``linewidths`` in :func:`heatmap` and :`clustermap` to 0 so that larger matrices plot correctly. This parameter still exists and can be used to get the old effect of lines demarcating each cell in the heatmap (the old default ``linewidths`` was 0.5).
+
+- You can now pass an integer to the ``xticklabels`` and ``yticklabels`` parameter of :func:`heatmap` (and, by extension, :func:`clustermap`). This will make the plot use the ticklabels inferred from the data, but only plot every ``n`` label, where ``n`` is the number you pass. This can help when visualizing larger matrices with some sensible ordering to the rows or columns of the dataframe.
+
 Bug fixes
 ~~~~~~~~~
 
-- Fixed a bug in :func:`clustermap` where the mask was not being reorganized using the dendrograms.
+- Fixed bugs in :func:`clustermap` where the mask and specified ticklabels were not being reorganized using the dendrograms.
 
 - Fixed a bug in :class:`FacetGrid` and :class:`PairGrid` that lead to incorrect legend labels when levels of the ``hue`` variable appeared in ``hue_order`` but not in the data.
 
diff --git a/seaborn/matrix.py b/seaborn/matrix.py
@@ -107,25 +107,38 @@ def __init__(self, data, vmin, vmax, cmap, center, robust, annot, fmt,
         plot_data = np.ma.masked_where(np.asarray(mask), plot_data)
 
         # Get good names for the rows and columns
-        if isinstance(xticklabels, bool) and xticklabels:
-            self.xticklabels = _index_to_ticklabels(data.columns)
+        xtickevery = 1
+        if isinstance(xticklabels, int) and xticklabels > 1:
+            xtickevery = xticklabels
+            xticklabels = _index_to_ticklabels(data.columns)
+        elif isinstance(xticklabels, bool) and xticklabels:
+            xticklabels = _index_to_ticklabels(data.columns)
         elif isinstance(xticklabels, bool) and not xticklabels:
-            self.xticklabels = ['' for _ in range(data.shape[1])]
-        else:
-            self.xticklabels = xticklabels
-
-        xlabel = _index_to_label(data.columns)
-
-        if isinstance(yticklabels, bool) and yticklabels:
-            self.yticklabels = _index_to_ticklabels(data.index)
+            xticklabels = ['' for _ in range(data.shape[1])]
+
+        ytickevery = 1
+        if isinstance(yticklabels, int) and yticklabels > 1:
+            ytickevery = yticklabels
+            yticklabels = _index_to_ticklabels(data.index)
+        elif isinstance(yticklabels, bool) and yticklabels:
+            yticklabels = _index_to_ticklabels(data.index)
         elif isinstance(yticklabels, bool) and not yticklabels:
-            self.yticklabels = ['' for _ in range(data.shape[0])]
+            yticklabels = ['' for _ in range(data.shape[0])]
         else:
-            self.yticklabels = yticklabels[::-1]
+            yticklabels = yticklabels[::-1]
 
-        ylabel = _index_to_label(data.index)
+        # Get the positions and used label for the ticks
+        nx, ny = data.T.shape
+        xstart, xend, xstep = 0, nx, xtickevery
+        self.xticks = np.arange(xstart, xend, xstep) + .5
+        self.xticklabels = xticklabels[xstart:xend:xstep]
+        ystart, yend, ystep = (ny - 1) % ytickevery, ny, ytickevery
+        self.yticks = np.arange(ystart, yend, ystep) + .5
+        self.yticklabels = yticklabels[ystart:yend:ystep]
 
         # Get good names for the axis labels
+        xlabel = _index_to_label(data.columns)
+        ylabel = _index_to_label(data.index)
         self.xlabel = xlabel if xlabel is not None else ""
         self.ylabel = ylabel if ylabel is not None else ""
 
@@ -204,8 +217,7 @@ def plot(self, ax, cax, kws):
         ax.set(xlim=(0, self.data.shape[1]), ylim=(0, self.data.shape[0]))
 
         # Add row and column labels
-        nx, ny = self.data.T.shape
-        ax.set(xticks=np.arange(nx) + .5, yticks=np.arange(ny) + .5)
+        ax.set(xticks=self.xticks, yticks=self.yticks)
         xtl = ax.set_xticklabels(self.xticklabels)
         ytl = ax.set_yticklabels(self.yticklabels, rotation="vertical")
 
@@ -233,7 +245,7 @@ def plot(self, ax, cax, kws):
 
 def heatmap(data, vmin=None, vmax=None, cmap=None, center=None, robust=False,
             annot=False, fmt=".2g", annot_kws=None,
-            linewidths=.5, linecolor="white",
+            linewidths=0, linecolor="white",
             cbar=True, cbar_kws=None, cbar_ax=None,
             square=False, ax=None, xticklabels=True, yticklabels=True,
             mask=None,
@@ -276,9 +288,9 @@ def heatmap(data, vmin=None, vmax=None, cmap=None, center=None, robust=False,
     annot_kws : dict of key, value mappings, optional
         Keyword arguments for ``ax.text`` when ``annot`` is True.
     linewidths : float, optional
-        Width of the lines that divide each cell.
+        Width of the lines that will divide each cell.
     linecolor : color, optional
-        Color of the lines that divide each cell.
+        Color of the lines that will divide each cell.
     cbar : boolean, optional
         Whether to draw a colorbar.
     cbar_kws : dict of key, value mappings, optional
@@ -292,14 +304,16 @@ def heatmap(data, vmin=None, vmax=None, cmap=None, center=None, robust=False,
     ax : matplotlib Axes, optional
         Axes in which to draw the plot, otherwise use the currently-active
         Axes.
-    xticklabels : list-like or bool, optional
+    xticklabels : list-like, int, or bool, optional
         If True, plot the column names of the dataframe. If False, don't plot
         the column names. If list-like, plot these alternate labels as the
-        xticklabels
-    yticklabels : list-like or bool, optional
+        xticklabels. If an integer, use the column names but plot only every
+        n label.
+    yticklabels : list-like, int, or bool, optional
         If True, plot the row names of the dataframe. If False, don't plot
         the row names. If list-like, plot these alternate labels as the
-        yticklabels
+        yticklabels. If an integer, use the index names but plot only every
+        n label.
     mask : boolean array or DataFrame, optional
         If passed, data will not be shown in cells where ``mask`` is True.
         Cells with missing values are automatically masked.
@@ -835,8 +849,22 @@ def plot_colors(self, xind, yind, **kws):
     def plot_matrix(self, colorbar_kws, xind, yind, **kws):
         self.data2d = self.data2d.iloc[yind, xind]
         self.mask = self.mask.iloc[yind, xind]
+
+        # Try to reorganize specified tick labels, if provided
+        xtl = kws.pop("xticklabels", True)
+        try:
+            xtl = np.asarray(xtl)[xind]
+        except (TypeError, IndexError):
+            pass
+        ytl = kws.pop("yticklabels", True)
+        try:
+            ytl = np.asarray(ytl)[yind]
+        except (TypeError, IndexError):
+            pass
+
         heatmap(self.data2d, ax=self.ax_heatmap, cbar_ax=self.cax,
-                cbar_kws=colorbar_kws, mask=self.mask, **kws)
+                cbar_kws=colorbar_kws, mask=self.mask,
+                xticklabels=xtl, yticklabels=ytl, **kws)
         self.ax_heatmap.yaxis.set_ticks_position('right')
         self.ax_heatmap.yaxis.set_label_position('right')
 
diff --git a/seaborn/tests/test_matrix.py b/seaborn/tests/test_matrix.py
@@ -200,6 +200,22 @@ def test_custom_ticklabels(self):
         nt.assert_equal(p.xticklabels, xticklabels)
         nt.assert_equal(p.yticklabels, yticklabels[::-1])
 
+    def test_custom_ticklabel_interval(self):
+
+        kws = self.default_kws.copy()
+        kws['xticklabels'] = 2
+        kws['yticklabels'] = 3
+        p = mat._HeatMapper(self.df_norm, **kws)
+
+        nx, ny = self.df_norm.T.shape
+        ystart = (ny - 1) % 3
+        npt.assert_array_equal(p.xticks, np.arange(0, nx, 2) + .5)
+        npt.assert_array_equal(p.yticks, np.arange(ystart, ny, 3) + .5)
+        npt.assert_array_equal(p.xticklabels,
+                               self.df_norm.columns[::2])
+        npt.assert_array_equal(p.yticklabels,
+                               self.df_norm.index[::-1][ystart:ny:3])
+
     def test_heatmap_annotation(self):
 
         ax = mat.heatmap(self.df_norm, annot=True, fmt=".1f",
@@ -879,3 +895,24 @@ def test_mask_reorganization(self):
                                    g.dendrogram_col.reordered_ind])
 
         plt.close("all")
+
+    def test_ticklabel_reorganization(self):
+
+        kws = self.default_kws.copy()
+        xtl = np.arange(self.df_norm.shape[1])
+        kws["xticklabels"] = list(xtl)
+        ytl = self.letters.ix[:self.df_norm.shape[0]]
+        kws["yticklabels"] = ytl
+
+        g = mat.clustermap(self.df_norm, **kws)
+
+        xtl_actual = [t.get_text() for t in g.ax_heatmap.get_xticklabels()]
+        ytl_actual = [t.get_text() for t in g.ax_heatmap.get_yticklabels()]
+
+        xtl_want = xtl[g.dendrogram_col.reordered_ind].astype("<U1")
+        ytl_want = ytl[g.dendrogram_row.reordered_ind].astype("<U1")[::-1]
+
+        npt.assert_array_equal(xtl_actual, xtl_want)
+        npt.assert_array_equal(ytl_actual, ytl_want)
+
+        plt.close("all")
