rework plotting to allow for plot_axes

contact_map/contact_count.py

@@ -16,13 +16,14 @@
 # pandas 0.25 not available on py27; can drop this when we drop py27
 _PD_VERSION = tuple(int(x) for x in pd.__version__.split('.')[:2])
 
-def _colorbar(with_colorbar, cmap_f, norm, min_val):
+
+def _colorbar(with_colorbar, cmap_f, norm, min_val, ax=None):
     if with_colorbar is False:
         return None
     elif with_colorbar is True:
         cbmin = np.floor(min_val)  # [-1.0..0.0] => -1; [0.0..1.0] => 0
         cbmax = 1.0
-        cb = ranged_colorbar(cmap_f, norm, cbmin, cbmax)
+        cb = ranged_colorbar(cmap_f, norm, cbmin, cbmax, ax=ax)
     # leave open other inputs to be parsed later (like tuples)
     return cb
 
@@ -198,18 +199,39 @@ def plot(self, cmap='seismic', vmin=-1.0, vmax=1.0, with_colorbar=True,
         """
         if not HAS_MATPLOTLIB:  # pragma: no cover
             raise RuntimeError("Error importing matplotlib")
+        fig, ax = plt.subplots(**kwargs)
+
+        # Check the number of pixels of the figure
+        self._check_number_of_pixels(fig)
+        self.plot_axes(ax=ax, cmap=cmap, vmin=vmin, vmax=vmax)
+
+        return (fig, ax)
+
+    def plot_axes(self, ax, cmap='seismic', vmin=-1.0, vmax=1.0,
+                  with_colorbar=True):
+        """
+        Plot contact matrix on a matplotlib.axes
+
+        Parameters
+        ----------
+        ax : matplotlib.axes
+           axes to plot the contact matrix on
+        cmap : str
+            color map name, default 'seismic'
+        vmin : float
+            minimum value for color map interpolation; default -1.0
+        vmax : float
+            maximum value for color map interpolation; default 1.0
+        with_colorbar : bool
+            If a colorbar is added to the axes
+        """
+
         norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
         cmap_f = plt.get_cmap(cmap)
-
-        fig, ax = plt.subplots(**kwargs)
         ax.axis([0, self.n_x, 0, self.n_y])
         ax.set_facecolor(cmap_f(norm(0.0)))
 
         min_val = 0.0
-
-        # Check the number of pixels of the figure
-        self._check_number_of_pixels(fig)
-
         for (pair, value) in self.counter.items():
             if value < min_val:
                 min_val = value
@@ -227,9 +249,7 @@ def plot(self, cmap='seismic', vmin=-1.0, vmax=1.0, with_colorbar=True,
             ax.add_patch(patch_0)
             ax.add_patch(patch_1)
 
-        _colorbar(with_colorbar, cmap_f, norm, min_val)
-
-        return (fig, ax)
+        _colorbar(with_colorbar, cmap_f, norm, min_val, ax=ax)
 
     def most_common(self, obj=None):
         """

contact_map/plot_utils.py

@@ -7,7 +7,8 @@
 except ImportError:  # pragma: no cover
     pass
 
-def ranged_colorbar(cmap, norm, cbmin, cbmax, name="Partial Map"):
+
+def ranged_colorbar(cmap, norm, cbmin, cbmax, ax=None):
     """Create a colorbar with given endpoints.
 
     Parameters
@@ -20,8 +21,8 @@ def ranged_colorbar(cmap, norm, cbmin, cbmax, name="Partial Map"):
         minimum value for the colorbar
     cbmax : float
         maximum value for the colorbar
-    name : str
-        name for the submap to be created
+    ax : matplotlib.axes
+        the axes to take space from to plot the colorbar
 
     Returns
     -------
@@ -33,6 +34,12 @@ def ranged_colorbar(cmap, norm, cbmin, cbmax, name="Partial Map"):
         cmap_f = plt.get_cmap(cmap)
     else:
         cmap_f = cmap
+
+    if ax is None:
+        fig = plt
+    else:
+        fig = ax.figure
+
     cbmin_normed = float(cbmin - norm.vmin) / (norm.vmax - norm.vmin)
     cbmax_normed = float(cbmax - norm.vmin) / (norm.vmax - norm.vmin)
     n_colors = int(round((cbmax_normed - cbmin_normed) * cmap_f.N))
@@ -42,7 +49,5 @@ def ranged_colorbar(cmap, norm, cbmin, cbmax, name="Partial Map"):
     new_norm = matplotlib.colors.Normalize(vmin=cbmin, vmax=cbmax)
     sm = plt.cm.ScalarMappable(cmap=new_cmap, norm=new_norm)
     sm._A = []
-    cb = plt.colorbar(sm, fraction=0.046, pad=0.04)
+    cb = fig.colorbar(sm, ax=ax, fraction=0.046, pad=0.04)
     return cb
-
-