Source Code for Module reflectometry.reduction.polplot

  1  """ 
  2  PolPlot2D is a panel for 4 cross-section 2D polarized reflectometry data. 
  3  """ 
  4   
  5  import wx,numpy,os 
  6  from math import log10, pow 
  7   
  8  import matplotlib as mpl 
  9  mpl.interactive(False) 
 10  #Use the WxAgg back end. The Wx one takes too long to render 
 11  mpl.use('WXAgg') 
 12   
 13  from copy import deepcopy 
 14  import matplotlib.cm 
 15  import matplotlib.colors 
 16  #from canvas import FigureCanvas as Canvas 
 17  from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as Canvas 
 18  from matplotlib.backend_bases import LocationEvent 
 19  from matplotlib.font_manager import FontProperties 
 20  from matplotlib.colors import Normalize,LogNorm 
 21  from cmapmenu import CMapMenu 
 22   
 23  import ticker 
 24   
 25   
 26 -def _rescale(lo,hi,step,pt=None,bal=None,scale='linear'): 
 27      """ 
 28      Rescale (lo,hi) by step, returning the new (lo,hi) 
 29      The scaling is centered on pt, with positive values of step 
 30      driving lo/hi away from pt and negative values pulling them in. 
 31      If bal is given instead of point, it is already in [0,1] coordinates. 
 32   
 33      This is a helper function for step-based zooming. 
 34      """ 
 35      # Convert values into the correct scale for a linear transformation 
 36      # TODO: use proper scale transformers 
 37      if scale=='log': 
 38          lo,hi = log10(lo),log10(hi) 
 39          if pt is not None: pt = log10(pt) 
 40   
 41      # Compute delta from axis range * %, or 1-% if persent is negative 
 42      if step > 0: 
 43          delta = float(hi-lo)*step/100 
 44      else: 
 45          delta = float(hi-lo)*step/(100-step) 
 46   
 47      # Add scale factor proportionally to the lo and hi values, preserving the 
 48      # point under the mouse 
 49      if bal is None: 
 50          bal = float(pt-lo)/(hi-lo) 
 51      lo = lo - bal*delta 
 52      hi = hi + (1-bal)*delta 
 53   
 54      # Convert transformed values back to the original scale 
 55      if scale=='log': 
 56          lo,hi = pow(10.,lo),pow(10.,hi) 
 57   
 58      return (lo,hi) 
 59   
 60 -def error_msg(msg, parent=None): 
 61      """ 
 62      Signal an error condition -- in a GUI, popup a error dialog 
 63      """ 
 64      # Code brought with minor podifications from mpl.backends.backend_wx 
 65      # Copyright (C) Jeremy O'Donoghue & John Hunter, 2003-4 
 66      dialog =wx.MessageDialog(parent  = parent, 
 67                               message = msg, 
 68                               caption = 'Polplot error', 
 69                               style=wx.OK | wx.CENTRE) 
 70      dialog.ShowModal() 
 71      dialog.Destroy() 
 72      return None 
 73   
 74   
 75 -def save_canvas(canvas,filebase): 
 76      """ 
 77      Given a canvas and a a base filename, construct a file save 
 78      dialog to save the canvas. 
 79      """ 
 80      # Code brought with minor modifications from mpl.backends.backend_wx 
 81      # Copyright (C) Jeremy O'Donoghue & John Hunter, 2003-4 
 82      # Allows the programmer to specify the base for the filename. 
 83      filetypes, exts, filter_index = canvas._get_imagesave_wildcards() 
 84      default_file = filebase + "." + canvas.get_default_filetype() 
 85      dlg = wx.FileDialog(canvas, "Save to file", "", default_file, filetypes, 
 86                          wx.SAVE|wx.OVERWRITE_PROMPT|wx.CHANGE_DIR) 
 87      dlg.SetFilterIndex(filter_index) 
 88      if dlg.ShowModal() == wx.ID_OK: 
 89          dirname  = dlg.GetDirectory() 
 90          filename = dlg.GetFilename() 
 91          #DEBUG_MSG('Save file dir:%s name:%s' % (dirname, filename), 3, self) 
 92          format = exts[dlg.GetFilterIndex()] 
 93          # Explicitly pass in the selected filetype to override the 
 94          # actual extension if necessary 
 95          try: 
 96              canvas.print_figure(os.path.join(dirname, filename), format=format) 
 97          except Exception, e: 
 98              error_msg(str(e)) 
 99      return 
100   
101 -def bbox_union(bboxes): 
102      """ 
103      Return a Bbox that contains all of the given bboxes. 
104      """ 
105      from matplotlib.transforms import Bbox 
106      if len(bboxes) == 1: 
107          return bboxes[0] 
108   
109      x0 = numpy.inf 
110      y0 = numpy.inf 
111      x1 = -numpy.inf 
112      y1 = -numpy.inf 
113   
114      for bbox in bboxes: 
115          xs = bbox.intervalx 
116          ys = bbox.intervaly 
117          x0 = min(x0, numpy.min(xs)) 
118          y0 = min(y0, numpy.min(ys)) 
119          x1 = max(x1, numpy.max(xs)) 
120          y1 = max(y1, numpy.max(ys)) 
121   
122      return Bbox([[x0, y0], [x1, y1]]) 
123   
124   
125 -class Plotter(wx.Panel): 
126 -    def __init__(self, parent, id = -1, dpi = None, **kwargs): 
127          wx.Panel.__init__(self, parent, id=id, **kwargs) 
128          self.figure = mpl.figure.Figure(dpi=dpi, figsize=(2,2)) 
129          self.canvas = Canvas(self, -1, self.figure) 
130          sizer = wx.BoxSizer(wx.VERTICAL) 
131          sizer.Add(self.canvas,1,wx.EXPAND) 
132          self.SetSizer(sizer) 
133   
134 -class Plotter4(wx.Panel): 
135      """ 
136      The PlotPanel has a Figure and a Canvas. OnSize events simply set a 
137      flag, and the actually redrawing of the 
138      figure is triggered by an Idle event. 
139      """ 
140 -    def __init__(self, parent, id = -1, dpi = None, color=None, **kwargs): 
141          # TODO: inherit directly from Canvas --- it is after all a panel. 
142   
143          # Set up the panel parameters 
144          #style = wx.NO_FULL_REPAINT_ON_RESIZE 
145          wx.Panel.__init__(self, parent, id = id, **kwargs) 
146          self.Bind(wx.EVT_CONTEXT_MENU, self.onContextMenu) 
147   
148          # Create the figure 
149          self.figure = mpl.figure.Figure(dpi=dpi, figsize=(2,2)) 
150          #self.canvas = NoRepaintCanvas(self, -1, self.figure) 
151          self.canvas = Canvas(self, -1, self.figure) 
152          self.canvas.mpl_connect('key_press_event',self.onKeyPress) 
153          self.canvas.mpl_connect('button_press_event',self.onButtonPress) 
154          self.canvas.mpl_connect('scroll_event',self.onWheel) 
155          # TODO use something for pan events 
156          sizer = wx.BoxSizer(wx.VERTICAL) 
157          sizer.Add(self.canvas,1,wx.EXPAND) 
158          self.SetSizer(sizer) 
159          # Construct an idle timer.  This won't be needed when matplotlib 
160          # supports draw_idle for wx. 
161          #self.idle_timer = wx.CallLater(1,self.onDrawIdle) 
162          #self.Fit() 
163   
164          # Create four subplots with no space between them. 
165          # Leave space for the colorbar. 
166          # Use a common set of axes. 
167          self.pp = self.figure.add_subplot(221) 
168          self.mm = self.figure.add_subplot(222,sharex=self.pp,sharey=self.pp) 
169          self.pm = self.figure.add_subplot(223,sharex=self.pp,sharey=self.pp) 
170          self.mp = self.figure.add_subplot(224,sharex=self.pp,sharey=self.pp) 
171          self.figure.subplots_adjust(left=.1, bottom=.1, top=.9, right=0.85, 
172                                      wspace=0.0, hspace=0.0) 
173   
174          self.axes = [self.pp, self.mm, self.pm, self.mp] 
175          self.grid = True 
176   
177          # Create the colorbar 
178          # Provide an empty handle to attach colormap properties 
179          self.coloraxes = self.figure.add_axes([0.88, 0.2, 0.04, 0.6]) 
180          self.colormapper = mpl.image.FigureImage(self.figure) 
181          self.colormapper.set_array(numpy.ones(1)) 
182          self.colorbar = self.figure.colorbar(self.colormapper,self.coloraxes) 
183   
184          # Provide slots for the graph labels 
185          self.tbox = self.figure.text(0.5, 0.95, '', 
186                                       horizontalalignment='center', 
187                                       fontproperties=FontProperties(size=16)) 
188          self.xbox = self.figure.text(0.5,0.01,'', 
189                                       verticalalignment='bottom', 
190                                       horizontalalignment='center', 
191                                       rotation='horizontal') 
192          self.ybox = self.figure.text(0.01,0.5,'', 
193                                       horizontalalignment='left', 
194                                       verticalalignment='center', 
195                                       rotation='vertical') 
196          self.zbox = self.figure.text(0.99,0.5,'', 
197                                       horizontalalignment='right', 
198                                       verticalalignment='center', 
199                                       rotation='vertical') 
200   
201          self.xscale = 'linear' 
202          self.yscale = 'linear' 
203          self.zscale = 'linear' 
204          self.set_vscale('log') 
205   
206          # Set up the default plot 
207          self.clear() 
208   
209          self._sets = {} # Container for all plotted objects 
210   
211 -    def autoaxes(self): 
212          return 
213          bbox = bbox_union([ax.dataLim for ax in self.axes]) 
214          xlims = bbox.intervalx 
215          ylims = bbox.intervaly 
216          self.pp.axis(list(xlims)+list(ylims)) 
217   
218 -    def onKeyPress(self,event): 
219          #if not event.inaxes: return 
220          # Let me zoom and unzoom even without the scroll wheel 
221          if event.key == 'a': 
222              setattr(event,'step',5.) 
223              self.onWheel(event) 
224          elif event.key == 'z': 
225              setattr(event,'step',-5.) 
226              self.onWheel(event) 
227   
228 -    def onButtonPress(self,event): 
229          # TODO: finish binder so that it allows tagging of fully qualified 
230          # events on an artist by artist basis. 
231          if event.inaxes == self.coloraxes: 
232              self.colormapper.set_clim(vmin=self.vmin,vmax=self.vmax) 
233              self.canvas.draw_idle() 
234          elif event.inaxes != None: 
235              self.autoaxes() 
236              self.canvas.draw_idle() 
237   
238 -    def onWheel(self, event): 
239          """ 
240          Process mouse wheel as zoom events 
241          """ 
242          ax = event.inaxes 
243          step = event.step 
244   
245          # Icky hardcoding of colorbar zoom. 
246          if ax == self.coloraxes: 
247              # rescale colormap: the axes are already scaled to 0..1, 
248              # so use bal instead of pt for centering 
249              lo,hi = self.colormapper.get_clim() 
250              lo,hi = _rescale(lo,hi,step,bal=event.ydata,scale=self.vscale) 
251              self.colormapper.set_clim(lo,hi) 
252          elif ax != None: 
253              # Event occurred inside a plotting area 
254              lo,hi = ax.get_xlim() 
255              lo,hi = _rescale(lo,hi,step,pt=event.xdata) 
256              ax.set_xlim((lo,hi)) 
257   
258              lo,hi = ax.get_ylim() 
259              lo,hi = _rescale(lo,hi,step,pt=event.ydata) 
260              ax.set_ylim((lo,hi)) 
261          else: 
262              # Check if zoom happens in the axes 
263              xdata,ydata = None,None 
264              x,y = event.x,event.y 
265              for ax in self.axes: 
266                  insidex,_ = ax.xaxis.contains(event) 
267                  if insidex: 
268                      xdata,_ = ax.transAxes.inverse_xy_tup((x,y)) 
269                      #print "xaxis",x,"->",xdata 
270                  insidey,_ = ax.yaxis.contains(event) 
271                  if insidey: 
272                      _,ydata = ax.transAxes.inverse_xy_tup((x,y)) 
273                      #print "yaxis",y,"->",ydata 
274              if xdata is not None: 
275                  lo,hi = ax.get_xlim() 
276                  lo,hi = _rescale(lo,hi,step,bal=xdata) 
277                  ax.set_xlim((lo,hi)) 
278              if ydata is not None: 
279                  lo,hi = ax.get_ylim() 
280                  lo,hi = _rescale(lo,hi,step,bal=ydata) 
281                  ax.set_ylim((lo,hi)) 
282   
283          self.canvas.draw_idle() 
284   
285   
286      # These are properties which the user should control but for which 
287      # a particular plottable might want to set a reasonable default. 
288      # For now leave control with the plottable. 
289 -    def set_xscale(self, scale='linear'): 
290          for axes in self.axes: axes.set_xscale(scale) 
291          self.xscale = scale 
292   
293 -    def get_xscale(self): 
294          return self.xscale 
295   
296 -    def set_yscale(self, scale='linear'): 
297          for axes in self.axes: axes.set_yscale(scale) 
298          self.yscale = scale 
299   
300 -    def get_yscale(self): 
301          return self.yscale 
302   
303 -    def set_vscale(self, scale='linear'): 
304          """Alternate between log and linear colormap""" 
305          if scale == 'linear': 
306              vmapper = Normalize(*self.colormapper.get_clim()) 
307              vformat = mpl.ticker.ScalarFormatter() 
308              vlocate = mpl.ticker.AutoLocator() 
309          else: 
310              vmapper = LogNorm(*self.colormapper.get_clim()) 
311              vformat = mpl.ticker.LogFormatterMathtext(base=10,labelOnlyBase=False) 
312              vlocate = mpl.ticker.LogLocator(base=10) 
313          self.colormapper.set_norm(vmapper) 
314          self.colorbar.formatter = vformat 
315          self.colorbar.locator = vlocate 
316   
317          self.vscale = scale 
318   
319 -    def get_vscale(self): 
320          return self.vscale 
321   
322      # Context menu implementation 
323 -    def onSaveImage(self, evt): 
324          """ 
325          Figure save dialog 
326          """ 
327          save_canvas(self.canvas,"polplot") 
328   
329      # Context menu implementation 
330 -    def onGridToggle(self, event): 
331          self.grid = not self.grid 
332          for ax in self.axes: ax.grid(alpha=0.4,visible=self.grid) 
333          self.draw() 
334   
335 -    def onContextMenu(self, event): 
336          """ 
337          Default context menu for a plot panel 
338          """ 
339          # Define a location event for the position of the popup 
340   
341          # TODO: convert from screen coords to window coords 
342          x,y = event.GetPosition() 
343          self.menuevent = LocationEvent("context",self.canvas, 
344                                         x,y,guiEvent=event) 
345   
346          popup = wx.Menu() 
347          item = popup.Append(wx.ID_ANY,'&Save image', 'Save image as PNG') 
348          wx.EVT_MENU(self, item.GetId(), self.onSaveImage) 
349          item = popup.Append(wx.ID_ANY,'&Grid on/off', 'Toggle grid lines') 
350          wx.EVT_MENU(self, item.GetId(), self.onGridToggle) 
351          cmaps = CMapMenu(self, mapper=self.colormapper, canvas=self.canvas) 
352          item = popup.AppendMenu(wx.ID_ANY, "Colourmaps", cmaps) 
353          #popup.Append(315,'&Properties...','Properties editor for the graph') 
354          #wx.EVT_MENU(self, 315, self.onProp) 
355   
356          pos = event.GetPosition() 
357          pos = self.ScreenToClient(pos) 
358          self.PopupMenu(popup, pos) 
359   
360   
361      ## The following is plottable functionality 
362 -    def properties(self,prop): 
363          """Set some properties of the graph. 
364   
365          The set of properties is not yet determined. 
366          """ 
367          # The particulars of how they are stored and manipulated (e.g., do 
368          # we want an inventory internally) is not settled.  I've used a 
369          # property dictionary for now. 
370          # 
371          # How these properties interact with a user defined style file is 
372          # even less clear. 
373   
374          # Properties defined by plot 
375          self.xbox.set_text(r"$%s$" % prop["xlabel"]) 
376          self.ybox.set_text(r"$%s$" % prop["ylabel"]) 
377          self.vbox.set_text(r"$%s$" % prop["vlabel"]) 
378          self.tbox.set_text(r"$%s$" % prop["title"]) 
379   
380          # Properties defined by user 
381          #self.axes.grid(True) 
382   
383 -    def clear(self): 
384          """Reset the plot""" 
385   
386          # TODO: Redraw is brutal.  Render to a backing store and swap in 
387          # TODO: rather than redrawing on the fly? 
388          # TODO: Want to retain graph properties such as grids and limits 
389          for ax in self.axes: 
390              ax.clear() 
391              ax.hold(True) 
392   
393          # Label the four cross sections 
394          props = dict(va='center',ha='left', 
395                       bbox=dict(facecolor='yellow',alpha=0.67)) 
396          self.pp.text( 0.05, 0.9, '$++$', props, transform=self.pp.transAxes) 
397          self.pm.text( 0.05, 0.9, '$+-$', props, transform=self.pm.transAxes) 
398          self.mp.text( 0.05, 0.9, '$-+$', props, transform=self.mp.transAxes) 
399          self.mm.text( 0.05, 0.9, '$--$', props, transform=self.mm.transAxes) 
400   
401          #self.pp.get_yticklabels()[0].set_visible(False) 
402          #self.mp.get_xticklabels()[0].set_visible(False) 
403          self.vmin, self.vmax = numpy.inf,-numpy.inf 
404          self.colormapper.set_clim(vmin=1,vmax=10) 
405   
406          # Hide tick labels for interior axes 
407          for l in self.pp.get_xticklabels(): l.set_visible(False) 
408          for l in self.mm.get_xticklabels(): l.set_visible(False) 
409          for l in self.mm.get_yticklabels(): l.set_visible(False) 
410          for l in self.mp.get_yticklabels(): l.set_visible(False) 
411   
412          # TODO: Hide tick labels from pm that might overlap neighbours 
413          # The proper algorithm would be to suppress all tick labels 
414          # which are within half the font height of the end of the 
415          # axis.  Second best is to suppress those within 10% of the 
416          # end of the axis. The current method is to remove the last 
417          # two, but it doesn't work very well. 
418          #self.pm.yaxis.get_majorticklabels()[-1].set_visible(False) 
419          #self.pm.xaxis.get_majorticklabels()[-1].set_visible(False) 
420          #self.pm.yaxis.get_majorticklabels()[-2].set_visible(False) 
421          #self.pm.xaxis.get_majorticklabels()[-2].set_visible(False) 
422   
423          # Set the limits on the colormap 
424          #vmin = self.vmin if self.vmin < self.vmax else 1 
425          #vmax = self.vmax if self.vmin < self.vmax else 2 
426          #self.colormapper.set_cmap(mpl.cm.cool) 
427          #self.coloraxes.xaxis.set_major_formatter(mpl.ticker.LogFormatterMathtext()) 
428          #print "setting formatter---doesn't seem to work" 
429          #self.coloraxes.yaxis.set_major_formatter(mpl.ticker.NullFormatter()) 
430   
431          for ax in self.axes: ax.grid(alpha=0.4,visible=self.grid) 
432          pass 
433   
434 -    def xaxis(self,label,units): 
435          """xaxis label and units. 
436   
437          Axis labels know about units. 
438   
439          We need to do this so that we can detect when axes are not 
440          commesurate.  Currently this is ignored other than for formatting 
441          purposes. 
442          """ 
443          if units != "": label = label + " (" + units + ")" 
444          self.xbox.set_text(r"$%s$" % (label)) 
445          pass 
446   
447 -    def yaxis(self,label,units): 
448          """yaxis label and units.""" 
449          if units != "": label = label + " (" + units + ")" 
450          self.ybox.set_text(r"$%s$" % (label)) 
451          pass 
452   
453 -    def vaxis(self,label,units): 
454          """vaxis label and units.""" 
455          if units != "": label = label + " (" + units + ")" 
456          self.vbox.set_text(r"$%s$" % (label)) 
457          pass 
458   
459 -    def surfacepol(self,poldata): 
460          for slice,data in [('++',poldata.pp),('--',poldata.mm), 
461                          ('+-',poldata.pm),('-+',poldata.mp)]: 
462              self.surface(slice,data) 
463          pass 
464   
465 -    def surface(self,slice,data): 
466          if slice == '++': ax = self.pp 
467          elif slice == '+-': ax = self.pm 
468          elif slice == '-+': ax = self.mp 
469          elif slice == '--': ax = self.mm 
470          else: raise ValueError, "expected polarization crosssection" 
471   
472          # Should be the following: 
473          #    im = ax.pcolor(data.xedges,data.yedges,data.v,shading='flat') 
474          # but this won't work in all versions of mpl, so first figure out 
475          # if we are using pcolorfast or pcolormesh then adjust the kwargs. 
476          # TODO Ack! adding 1 for the purposes of visualization! Fix log plots! 
477          x,y,v = data.xedges,data.yedges,data.v+1 
478          try: 
479              im = ax.pcolorfast(x,y,v) 
480          except AttributeError: 
481              im = ax.pcolormesh(x,y,v,shading='flat') 
482          #self.colormapper.add_observer(im) 
483          def on_changed(m): 
484              print "changed",m 
485              self.colorbar.set_cmap(m.get_cmap()) 
486              self.colorbar.set_clim(m.get_clim()) 
487              self.colorbar.update_bruteforce(m) 
488              im.set_cmap(m.get_cmap()) 
489              im.set_clim(m.get_clim()) 
490              self.canvas.draw_idle() 
491          self.colormapper.callbacksSM.connect('changed',on_changed) 
492   
493          self.vmin = min(self.vmin, numpy.min(v)) 
494          self.vmax = max(self.vmax, numpy.max(v)) 
495          self.colormapper.set_clim(vmin=self.vmin,vmax=self.vmax) 
496          self.autoaxes() 
497          self.canvas.draw_idle() 
498   
499          return im 
500   
501   
502 -def demo(): 
503      import data 
504      # Get some data 
505      d = data.peakspol(n=355) 
506      #d = data.peakspol(n=1000) 
507   
508      # Make a frame to show it 
509      app = wx.PySimpleApp() 
510      frame = wx.Frame(None,-1,'Plottables') 
511      plotter = Plotter4(frame) 
512      frame.Show() 
513   
514      # render the graph to the pylab plotter 
515      plotter.clear() 
516      plotter.set_vscale('linear') 
517      plotter.surfacepol(d) 
518   
519      app.MainLoop() 
520      pass 
521   
522  if __name__ == "__main__": demo() 
523