class MultiPlot: separate layout from contents

Examples/ff/SasPyramid4.py

@@ -19,5 +19,5 @@ def simulate_at(omega):
 
 namedResults = [simulate_at(omega) for omega in [0, 30]]
 
-bp.make_plot(namedResults, det, "ff_Pyramid4")
+bp.make_plot_row(namedResults, det, "ff_Pyramid4")
 # bp.save_results(namedResults, "ff_Pyramid4")

Examples/ff/bornplot2.py

@@ -8,6 +8,7 @@ from matplotlib import rc
 rc('font', **{'family': 'sans-serif', 'sans-serif': ['Helvetica']})
 rc('text', usetex=True)
 rc('image', cmap='inferno')
+mpl.rcParams['image.interpolation'] = 'none'
 
 import bornagain as ba
 from bornagain import deg, angstrom
@@ -34,77 +35,96 @@ def save_results(namedResults, name):
         numpy.savetxt(fname, namedResults[i].result.array())
 
 
-def make_plot(namedResults, det, name, nrow=1):
+class MultiPlot:
+    def __init__(self, n, ncol):
+        self.n = n
+        self.ncol = ncol
+        self.nrow = 1 + (self.n - 1) // self.ncol
+
+        # Parameters as fraction of subfig size.
+        yskip = 0.2
+        bottomskip = yskip
+        topskip = yskip/2
+        xskip = 0.18
+        leftskip = xskip
+        rightskip = 0.28 + ncol*0.03
+        xtot = self.ncol*1.0 + (self.ncol - 1)*xskip + leftskip + rightskip
+        ytot = self.nrow*1.0 + (self.nrow - 1)*yskip + bottomskip + topskip
+
+        # We need parameters as fraction of total fig size.
+        self.xskip = xskip / xtot
+        self.leftskip = leftskip / xtot
+        self.rightskip = rightskip / xtot
+        self.yskip = yskip / ytot
+        self.bottomskip = bottomskip / ytot
+        self.topskip = topskip / ytot
+
+        # Set total figure dimensions.
+        ftot = 5
+        self.fontsize = 18 + 36.0/(ncol + 2)
+        # Create the figure 'fig' and its subplots axes ('tmp'->'axes').
+        self.fig, tmp = plt.subplots(self.nrow, self.ncol, figsize=(ftot*xtot, ftot*ytot))
+        if n > 1:
+            self.axes = tmp.flat
+        else:
+            self.axes = [tmp]
+
+        # Adjust whitespace around and between subfigures.
+        plt.subplots_adjust(wspace=self.xskip,
+                            hspace=self.yskip,
+                            left=self.leftskip,
+                            right=1 - self.rightskip,
+                            bottom=self.bottomskip,
+                            top=1 - self.topskip)
+
+    def plot_colorbar(self, im):
+        # Plot the color scale.
+        cbar_ax = self.fig.add_axes([
+            1 - self.rightskip + 0.4*self.xskip, self.bottomskip, 0.25*self.xskip,
+            1 - self.bottomskip - self.topskip
+        ])
+        cb = self.fig.colorbar(im, cax=cbar_ax)
+        cb.set_label(r'$\left|F(q)\right|^2/V^{\,2}$', fontsize=self.fontsize)
+
+
+def make_plot_row(namedResults, det, name):
+    make_plot(namedResults, det, name, len(namedResults))
+
+
+def make_plot(namedResults, det, name, ncol):
     """
     Make a plot consisting of one detector image for each Result in results,
     plus one common color scale.
 
     :param results: List of simulation results
     :param det: Detector
-    :param name: Filename for plot during save
-    :param nrow: Number of rows for different plots
+    :param name: Filename for multiplot during save
+    :param ncol: Number of columns in multiplot
     """
-    mpl.rcParams['image.interpolation'] = 'none'
-    n = len(namedResults)
-    ncol = 1 + (n - 1) // nrow
-    # Parameters as fraction of subfig size.
-    yskip = 0.2  # +ncol*0.02
-    bottomskip = yskip
-    topskip = yskip/2
-    xskip = 0.18
-    leftskip = xskip
-    rightskip = 0.28 + ncol*0.03
-    xtot = ncol*1.0 + (ncol - 1)*xskip + leftskip + rightskip
-    ytot = nrow*1.0 + (nrow - 1)*yskip + bottomskip + topskip
-    # We need parameters as fraction of total fig size.
-    xskip /= xtot
-    leftskip /= xtot
-    rightskip /= xtot
-    yskip /= ytot
-    bottomskip /= ytot
-    topskip /= ytot
-    # Set total figure dimensions.
-    ftot = 5
-    fontsize = 18 + 36.0/(ncol + 2)
-    # Create the figure 'fig' and its subplots axes ('tmp'->'axes').
-    fig, tmp = plt.subplots(nrow, ncol, figsize=(ftot*xtot, ftot*ytot))
-    if n > 1:
-        axes = tmp.flat
-    else:
-        axes = [tmp]
+    multiPlot = MultiPlot(len(namedResults), ncol)
+
     # Always the same color scale, to facilitate comparisons between figures.
     norm = mpl.colors.LogNorm(1e-8, 1)
     # Plot the subfigures.
     for i in range(len(namedResults)):
         item = namedResults[i]
-        ax = axes[i]
+        ax = multiPlot.axes[i]
         im = ax.imshow(item.result.array(),
                        norm=norm,
                        extent=rectangle(det),
                        aspect=1)
-        ax.set_xlabel(r'$\phi_{\rm f} (^{\circ})$', fontsize=fontsize)
+        ax.set_xlabel(r'$\phi_{\rm f} (^{\circ})$', fontsize=multiPlot.fontsize)
         if i % ncol == 0:
             ax.set_ylabel(r'$\alpha_{\rm f} (^{\circ})$',
-                          fontsize=fontsize)
+                          fontsize=multiPlot.fontsize)
         if item.title != "":
-            ax.set_title(item.title, fontsize=fontsize)
+            ax.set_title(item.title, fontsize=multiPlot.fontsize)
         ax.tick_params(axis='both',
                        which='major',
-                       labelsize=fontsize*21/24)
-    # Adjust whitespace around and between subfigures.
-    plt.subplots_adjust(wspace=xskip,
-                        hspace=yskip,
-                        left=leftskip,
-                        right=1 - rightskip,
-                        bottom=bottomskip,
-                        top=1 - topskip)
-    # Plot the color scale.
-    cbar_ax = fig.add_axes([
-        1 - rightskip + 0.4*xskip, bottomskip, 0.25*xskip,
-        1 - bottomskip - topskip
-    ])
-    cb = fig.colorbar(im, cax=cbar_ax)
-    cb.set_label(r'$\left|F(q)\right|^2/V^{\,2}$', fontsize=fontsize)
+                       labelsize=multiPlot.fontsize*21/24)
+
+    multiPlot.plot_colorbar(im)
+
     # Show or export
     plt.savefig(name + ".pdf", format="pdf", bbox_inches='tight')
     # plt.show()