enable passing of external Axes to plot on

Adds to the plot_image and plot_tabular functions an optional argument ax: plt.Axes. When given, this is the Axes that will be used to plot on and the internal plt.subplots call is skipped. This is useful for using these functions in custom multi-panel plots (we want to use this in a paper).

Author: egpbos

dianna/visualization/image.py

@@ -16,7 +16,8 @@ def plot_image(heatmap,
                heatmap_range=(None, None),  # (vmin, vmax)
                data_cmap=None,
                show_plot=True,
-               output_filename=None):
+               output_filename=None,
+               ax=None):
     """Plots a heatmap image.
 
     Optionally, the heatmap (typically a saliency map of an explainer) can be
@@ -38,13 +39,18 @@ def plot_image(heatmap,
         show_plot: Shows plot if true (for testing or writing plots to disk
                    instead).
         output_filename: Name of the file to save the plot to (optional).
+        ax: matplotlib.Axes object to plot on (optional).
 
     Returns:
         None
     """
     # default cmap depends on shape: grayscale or colour
 
-    fig, ax = plt.subplots()
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.get_figure()
+
     alpha = 1
     if original_data is not None:
         if len(original_data.shape) == 2 and data_cmap is None:

dianna/visualization/tabular.py

@@ -13,6 +13,7 @@ def plot_tabular(
     num_features: Optional[int] = None,
     show_plot: Optional[bool] = True,
     output_filename: Optional[str] = None,
+    ax: Optional[plt.Axes] = None,
 ) -> plt.Figure:
     """Plot feature importance with segments highlighted.
 
@@ -26,21 +27,23 @@ def plot_tabular(
             plots to disk instead).
         output_filename (str, optional): Name of the file to save
             the plot to (optional).
+        ax (matplotlib.Axes, optional): externally created canvas to plot on.
 
     Returns:
         plt.Figure
     """
     if not num_features:
         num_features = len(x)
-
-
     abs_values = [abs(i) for i in x]
     top_values = [x for _, x in sorted(zip(abs_values, x), reverse=True)][:num_features]
     top_features = [x for _, x in sorted(zip(abs_values, y), reverse=True)][
         :num_features
     ]
 
-    fig, ax = plt.subplots()
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.get_figure()
     colors = ["r" if x >= 0 else "b" for x in top_values]
     ax.barh(top_features, top_values, color=colors)
     ax.set_xlabel(x_label)