Changeset 158 for trunk/src/plotting.py

Timestamp:

04/18/08 14:18:44 (9 months ago)

Author:

apdavison

Message:

Added SimpleMultiplot class to the plotting module.

Files:

• trunk/src/plotting.py (modified) (2 diffs)

trunk/src/plotting.py

@@ -2 +2 @@
-import numpy #, pylab
-import sys
-
+
+
+
-"""
-Plotting helpers
@@ -50 +52 @@
-
-
+
+def set_frame(ax,boollist,linewidth=2):
+    assert len(boollist) == 4
+    bottom = Line2D([0, 1], [0, 0], transform=ax.transAxes, linewidth=linewidth, color='k')
+    left   = Line2D([0, 0], [0, 1], transform=ax.transAxes, linewidth=linewidth, color='k')
+    top    = Line2D([0, 1], [1, 1], transform=ax.transAxes, linewidth=linewidth, color='k')
+    right  = Line2D([1, 0], [1, 1], transform=ax.transAxes, linewidth=linewidth, color='k')
+    # anti-aliased?
+    if boollist != [True,True,True,True]:
+        ax.set_frame_on(False)
+        for side,draw in zip([left,bottom,right,top],boollist):
+            if draw:
+                ax.add_line(side)
+
+class SimpleMultiplot(object):
+    """
+    A figure consisting of multiple panels, all with the same datatype and
+    the same x-range.
+    """
+    
+    def __init__(self, nrows, ncolumns, title="", xlabel=None, ylabel=None,
+                 scaling=('linear','linear')):
+        self.fig = Figure()
+        self.canvas = FigureCanvas(self.fig)
+        self.axes = []
+        self.all_panels = self.axes
+        self.nrows = nrows
+        self.ncolumns = ncolumns
+        self.n = nrows*ncolumns
+        self._curr_panel = 0
+        self.title = title
+        topmargin = 0.06
+        rightmargin = 0.02
+        bottommargin = 0.1
+        leftmargin=0.1
+        panelsep = 0.05
+        panelheight = (1 - topmargin - bottommargin - (nrows-1)*panelsep)/nrows
+        panelwidth = (1 - leftmargin - rightmargin - (ncolumns-1)*panelsep)/ncolumns
+        assert panelheight > 0
+        
+        bottomlist = [bottommargin + i*panelsep + i*panelheight for i in range(nrows)]
+        leftlist = [leftmargin + j*panelsep + j*panelwidth for j in range(ncolumns)]
+        bottomlist.reverse()
+        for j in range(ncolumns):
+            for i in range(nrows):
+                ax = self.fig.add_axes([leftlist[j],bottomlist[i],panelwidth,panelheight])
+                set_frame(ax,[True,True,False,False])
+                ax.xaxis.tick_bottom()
+                ax.yaxis.tick_left()
+                self.axes.append(ax)
+        if xlabel:
+            self.axes[self.nrows-1].set_xlabel(xlabel)
+        if ylabel:
+            self.fig.text(0.5*leftmargin,0.5,ylabel,
+                          rotation='vertical',
+                          horizontalalignment='center',
+                          verticalalignment='center')
+        if scaling == ("linear","linear"):
+            self.plot_function = "plot"
+        elif scaling == ("log", "log"):
+            self.plot_function = "loglog"
+        elif scaling == ("log", "linear"):
+            self.plot_function = "semilogx"
+        elif scaling == ("linear", "log"):
+            self.plot_function = "semilogy"
+        else:
+            raise Exception("Invalid value for scaling parameter")
+    
+    def finalise(self):
+        """Adjustments to be made after all panels have been plotted."""
+        # Turn off tick labels for all x-axes except the bottom one
+        self.fig.text(0.5, 0.99, self.title, horizontalalignment='center',
+                      verticalalignment='top')
+        for ax in self.axes[0:self.nrows-1]+self.axes[self.nrows:]:
+            ax.xaxis.set_ticklabels([])
+
+    def save(self,filename):
+        """Save/print the figure to file."""
+        self.finalise()
+        self.canvas.print_figure(filename)
+    
+    def next_panel(self):
+        ax = self.axes[self._curr_panel]
+        self._curr_panel += 1
+        if self._curr_panel >= self.n:
+            self._curr_panel = 0
+        ax.plot1 = getattr(ax, self.plot_function)
+        return ax
+        
+    def panel(self,i):
+        """Return panel i."""
+        ax = self.axes[i]
+        ax.plot1 = getattr(ax, self.plot_function)
+        return ax
+