Merge pull request PaNOSC-ViNYL#100 from PaNOSC-ViNYL/return_fig_ax

Allows make_plot and make_sub_plot to return fig and ax object

Author: mads-bertelsen

mcstasscript/_version.py

@@ -1 +1 @@
-__version__ = "0.0.77"
+__version__ = "0.0.78"

mcstasscript/helper/plot_helper.py

@@ -6,6 +6,7 @@
 
 from matplotlib.ticker import MaxNLocator
 from matplotlib.colors import BoundaryNorm
+import matplotlib.ticker as mticker
 
 from mcstasscript.data.data import McStasData
 from mcstasscript.data.data import McStasDataEvent
@@ -42,6 +43,18 @@ def _fmt(x, pos):
         return r'${}\cdot 10^{{{}}}$'.format(a, b)
 
 
+def _fmt(x, pos):
+    a, b = f"{x:.2e}".split("e")
+    b = int(b)
+
+    if abs(float(a) - 1) < 0.01:
+        return rf"10^{{{b}}}"
+    else:
+        return rf"{a}\cdot 10^{{{b}}}"
+
+formatter = mticker.FuncFormatter(lambda x, pos: rf"${_fmt(x, pos)}$")
+
+
 def _find_min_max_I(data):
     """
     Returns minimum and maximum intensity to plot given dataset
@@ -209,8 +222,13 @@ def _plot_fig_ax(data, fig, ax, **kwargs):
             if "colorbar_axes" in kwargs:
                 cax = kwargs["colorbar_axes"]
 
-            colorbar = fig.colorbar(im, ax=ax, cax=cax,
-                                    format=matplotlib.ticker.FuncFormatter(_fmt))
+            colorbar = fig.colorbar(im, ax=ax, cax=cax)
+
+            # If your colorbar is logarithmic:
+            #colorbar.locator = mticker.LogLocator(base=10)
+            #colorbar.formatter = mticker.LogFormatterMathtext(base=10)
+
+            #colorbar.update_ticks()
 
             if data.metadata.zlabel is not None:
                 colorbar.set_label(data.metadata.zlabel)

mcstasscript/interface/plotter.py

@@ -31,12 +31,18 @@ def make_plot(data_list, **kwargs):
         figsize = (13, 7)
 
     for data in data_list:
+        #fig, ax0 = plt.subplots(figsize=figsize, constrained_layout=True)
         fig, ax0 = plt.subplots(figsize=figsize, tight_layout=True)
         _plot_fig_ax(data, fig, ax0, **kwargs)
 
+        if "return_fig_ax" in kwargs and len(data_list) == 1:
+            if kwargs["return_fig_ax"]:
+                return fig, ax0
+
     if "filename" in kwargs:
-        fig.tight_layout()
+        #fig.tight_layout()
         fig.savefig(kwargs["filename"])
+        plt.close()
     else:
         plt.show()
 
@@ -96,6 +102,9 @@ def make_sub_plot(data_list, **kwargs):
 
     fig.tight_layout()
 
+    if "return_fig_ax" in kwargs and kwargs["return_fig_ax"]:
+        return fig, ax
+
     if "filename" in kwargs:
         fig.tight_layout()
         fig.savefig(kwargs["filename"])