Add ax arg to plot method

brainglobe_heatmap/heatmaps.py

@@ -173,11 +173,13 @@ def render(self, **kwargs) -> Scene:
     def plot(
         self,
         show_legend: bool = False,
-        xlabel: str = "μm",
-        ylabel: str = "μm",
+        xlabel: str = "µm",
+        ylabel: str = "µm",
         hide_axes: bool = False,
         filename: Optional[str] = None,
         cbar_label: Optional[str] = None,
+        ax: Optional[plt.Axes] = None,
+        show_cbar: bool = True,
         **kwargs,
     ) -> plt.Figure:
         """
@@ -187,7 +189,11 @@ def plot(
             self.regions_meshes, self.scene.root
         )
 
-        f, ax = plt.subplots(figsize=(9, 9))
+        if ax is None:
+            f, ax = plt.subplots(figsize=(9, 9))
+        else:
+            f = plt.gcf()
+
         for r, coords in projected.items():
             name, segment = r.split("_segment_")
             ax.fill(
@@ -201,30 +207,33 @@ def plot(
                 alpha=0.3 if name == "root" else None,
             )
 
-        # make colorbar
-        divider = make_axes_locatable(ax)
-        cax = divider.append_axes("right", size="5%", pad=0.05)
-
-        # cmap = mpl.cm.cool
-        norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
-        if self.label_regions is True:
-            cbar = f.colorbar(
-                mpl.cm.ScalarMappable(
-                    norm=None,
-                    cmap=mpl.cm.get_cmap(self.cmap, len(self.values)),
-                ),
-                cax=cax,
-            )
-        else:
-            cbar = f.colorbar(
-                mpl.cm.ScalarMappable(norm=norm, cmap=self.cmap), cax=cax
-            )
+        if show_cbar:
+            # make colorbar
+            divider = make_axes_locatable(ax)
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+
+            # cmap = mpl.cm.cool
+            norm = mpl.colors.Normalize(vmin=self.vmin, vmax=self.vmax)
+            if self.label_regions is True:
+                cbar = f.colorbar(
+                    mpl.cm.ScalarMappable(
+                        norm=None,
+                        cmap=mpl.cm.get_cmap(self.cmap, len(self.values)),
+                    ),
+                    cax=cax,
+                )
+            else:
+                cbar = f.colorbar(
+                    mpl.cm.ScalarMappable(norm=norm, cmap=self.cmap), cax=cax
+                )
 
-        if cbar_label is not None:
-            cbar.set_label(cbar_label)
+            if cbar_label is not None:
+                cbar.set_label(cbar_label)
 
-        if self.label_regions is True:
-            cbar.ax.set_yticklabels([r.strip() for r in self.values.keys()])
+            if self.label_regions is True:
+                cbar.ax.set_yticklabels(
+                    [r.strip() for r in self.values.keys()]
+                )
 
         # style axes
         ax.invert_yaxis()
@@ -250,6 +259,9 @@ def plot(
 
         if show_legend:
             ax.legend()
-        plt.show()
 
-        return f
+        if ax is None:
+            plt.show()
+            return f
+
+        return ax