source: branches/Release-2-fixes/python/asaplot.py@ 2377

Last change on this file since 2377 was 693, checked in by mar637, 19 years ago

added user customisable color and linestyles

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 20.3 KB
RevLine 
[111]1"""
2ASAP plotting class based on matplotlib.
3"""
4
5import sys
6from re import match
7import Tkinter as Tk
8
9import matplotlib
10matplotlib.use("TkAgg")
11
12from matplotlib.backends import new_figure_manager, show
13from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, \
[117]14 FigureManagerTkAgg
[111]15from matplotlib.figure import Figure, Text
[491]16from matplotlib.font_manager import FontProperties
[118]17from matplotlib.numerix import sqrt
[482]18from matplotlib import rc, rcParams
[693]19from asap import rcParams as asaprcParams
[111]20
21# Force use of the newfangled toolbar.
22matplotlib.rcParams['toolbar'] = 'toolbar2'
23
24class ASAPlot:
25 """
26 ASAP plotting class based on matplotlib.
27 """
28
[482]29 def __init__(self, rows=1, cols=0, title='', size=(8,6), buffering=False):
[117]30 """
31 Create a new instance of the ASAPlot plotting class.
[119]32
33 If rows < 1 then a separate call to set_panels() is required to define
34 the panel layout; refer to the doctext for set_panels().
[117]35 """
36 self.window = Tk.Tk()
[482]37 self.is_dead = False
38 def dest_callback():
39 self.is_dead = True
40 self.window.destroy()
[111]41
[482]42 self.window.protocol("WM_DELETE_WINDOW", dest_callback)
43
44 self.figure = Figure(figsize=size, facecolor='#ddddee')
[202]45 self.canvas = FigureCanvasTkAgg(self.figure, master=self.window)
[117]46 self.canvas.get_tk_widget().pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
[111]47
[117]48 # Simply instantiating this is enough to get a working toolbar.
49 self.figmgr = FigureManagerTkAgg(self.canvas, 1, self.window)
50 self.window.wm_title('ASAPlot graphics window')
[111]51
[120]52 self.events = {'button_press':None,
53 'button_release':None,
54 'motion_notify':None}
55
[118]56 self.set_title(title)
[117]57 self.subplots = []
[118]58 if rows > 0:
59 self.set_panels(rows, cols)
[111]60
[693]61 # Set matplotlib default colour sequence.
62 self.colormap = "green red black cyan magenta orange blue purple yellow pink".split()
63
64 c = asaprcParams['plotter.colours']
65 if isinstance(c,str) and len(c) > 0:
66 self.colormap = c.split()
[111]67
[693]68 self.lsalias = {"line": [1,0],
69 "dashdot": [4,2,1,2],
70 "dashed" : [4,2,4,2],
71 "dotted" : [1,2],
72 "dashdotdot": [4,2,1,2,1,2],
73 "dashdashdot": [4,2,4,2,1,2]
74 }
75
76 styles = "line dashed dotted dashdot".split()
77 c = asaprcParams['plotter.linestyles']
78 if isinstance(c,str) and len(c) > 0:
79 styles = c.split()
80 s = []
81 for ls in styles:
82 if self.lsalias.has_key(ls):
83 s.append(self.lsalias.get(ls))
84 else:
85 s.append('-')
86 self.linestyles = s
87
[652]88 self.color = 0;
[693]89 self.linestyle = 0;
[117]90 self.attributes = {}
[620]91 self.loc = 0
[111]92
[117]93 matplotlib.interactive = True
94 self.buffering = buffering
[482]95
[117]96 self.canvas.show()
[111]97
98
[117]99 def clear(self):
100 """
101 Delete all lines from the plot. Line numbering will restart from 1.
102 """
[111]103
[652]104 for i in range(len(self.lines)):
[117]105 self.delete(i)
106 self.axes.clear()
[652]107 self.color = 0
[693]108 self.linestyle = 0
[117]109 self.lines = []
[111]110
[693]111 def palette(self, color, colormap=None, linestyle=0, linestyles=None):
[652]112 if colormap:
[693]113 if isinstance(colormap,list):
114 self.colormap = colormap
115 elif isinstance(colormap,str):
116 self.colormap = colormap.split()
[652]117 if 0 <= color < len(self.colormap):
118 self.color = color
[693]119 if linestyles:
120 self.linestyles = []
121 if isinstance(linestyles,list):
122 styles = linestyles
123 elif isinstance(linestyles,str):
124 styles = linestyles.split()
125 for ls in styles:
126 if self.lsalias.has_key(ls):
127 self.linestyles.append(self.lsalias.get(ls))
128 else:
129 self.linestyles.append(self.lsalias.get('line'))
130 if 0 <= linestyle < len(self.linestyles):
131 self.linestyle = linestyle
[652]132
[111]133 def delete(self, numbers=None):
[117]134 """
135 Delete the 0-relative line number, default is to delete the last.
136 The remaining lines are NOT renumbered.
137 """
[111]138
[117]139 if numbers is None: numbers = [len(self.lines)-1]
[111]140
[117]141 if not hasattr(numbers, '__iter__'):
142 numbers = [numbers]
[111]143
[117]144 for number in numbers:
145 if 0 <= number < len(self.lines):
146 if self.lines[number] is not None:
147 for line in self.lines[number]:
148 line.set_linestyle('None')
149 self.lines[number] = None
[620]150 self.show()
[111]151
152 def get_line(self):
[117]153 """
154 Get the current default line attributes.
155 """
156 return self.attributes
[111]157
158
[119]159 def hist(self, x=None, y=None, fmt=None):
160 """
161 Plot a histogram. N.B. the x values refer to the start of the
162 histogram bin.
163
164 fmt is the line style as in plot().
165 """
166
167 if x is None:
168 if y is None: return
169 x = range(0,len(y))
170
171 if len(x) != len(y):
172 return
173
174 l2 = 2*len(x)
175 x2 = range(0,l2)
176 y2 = range(0,l2)
177
178 for i in range(0,l2):
179 x2[i] = x[i/2]
180
181 y2[0] = 0
182 for i in range(1,l2):
183 y2[i] = y[(i-1)/2]
184
185 self.plot(x2, y2, fmt)
186
187
[111]188 def hold(self, hold=True):
[117]189 """
190 Buffer graphics until subsequently released.
191 """
192 self.buffering = hold
[111]193
194
[620]195 def legend(self, loc=None):
[117]196 """
197 Add a legend to the plot.
[111]198
[117]199 Any other value for loc else disables the legend:
200 1: upper right
201 2: upper left
202 3: lower left
203 4: lower right
204 5: right
205 6: center left
206 7: center right
207 8: lower center
208 9: upper center
209 10: center
[111]210
[117]211 """
[620]212 if isinstance(loc,int):
213 if 0 > loc > 10: loc = 0
214 self.loc = loc
[117]215 self.show()
[111]216
217
218 def map(self):
[117]219 """
220 Reveal the ASAPlot graphics window and bring it to the top of the
221 window stack.
222 """
223 self.window.wm_deiconify()
224 self.window.lift()
[111]225
[117]226
[111]227
228 def plot(self, x=None, y=None, mask=None, fmt=None, add=None):
[117]229 """
230 Plot the next line in the current frame using the current line
231 attributes. The ASAPlot graphics window will be mapped and raised.
[111]232
[117]233 The argument list works a bit like the matlab plot() function.
234 """
[111]235
[117]236 if x is None:
237 if y is None: return
238 x = range(len(y))
[111]239
[117]240 elif y is None:
241 y = x
242 x = range(len(y))
[111]243
[117]244 if mask is None:
245 if fmt is None:
246 line = self.axes.plot(x, y)
247 else:
248 line = self.axes.plot(x, y, fmt)
249 else:
250 segments = []
[111]251
[117]252 mask = list(mask)
253 i = 0
254 while mask[i:].count(1):
255 i += mask[i:].index(1)
256 if mask[i:].count(0):
257 j = i + mask[i:].index(0)
258 else:
259 j = len(mask)
[111]260
[117]261 segments.append(x[i:j])
262 segments.append(y[i:j])
[111]263
[117]264 i = j
[111]265
[117]266 line = self.axes.plot(*segments)
[111]267
[117]268 # Add to an existing line?
269 if add is None or len(self.lines) < add < 0:
[119]270 # Don't add.
[117]271 self.lines.append(line)
272 i = len(self.lines) - 1
273 else:
274 if add == 0: add = len(self.lines)
275 i = add - 1
276 self.lines[i].extend(line)
[111]277
[117]278 # Set/reset attributes for the line.
279 gotcolour = False
280 for k, v in self.attributes.iteritems():
281 if k == 'color': gotcolour = True
282 for segment in self.lines[i]:
283 getattr(segment, "set_%s"%k)(v)
[111]284
[652]285 if not gotcolour and len(self.colormap):
[117]286 for segment in self.lines[i]:
[652]287 getattr(segment, "set_color")(self.colormap[self.color])
[693]288 if len(self.colormap) == 1:
289 getattr(segment, "set_dashes")(self.linestyles[self.linestyle])
[652]290 self.color += 1
291 if self.color >= len(self.colormap):
292 self.color = 0
[111]293
[693]294 if len(self.colormap) == 1:
295 self.linestyle += 1
296 if self.linestyle >= len(self.linestyles):
297 self.linestyle = 0
298
[117]299 self.show()
[111]300
301
[482]302 def position(self):
303 """
304 Use the mouse to get a position from a graph.
305 """
306
307 def position_disable(event):
308 self.register('button_press', None)
309 print '%.4f, %.4f' % (event.xdata, event.ydata)
310
311 print 'Press any mouse button...'
312 self.register('button_press', position_disable)
313
314
[111]315 def quit(self):
[117]316 """
317 Destroy the ASAPlot graphics window.
318 """
319 self.window.destroy()
[111]320
321
[482]322 def region(self):
323 """
324 Use the mouse to get a rectangular region from a plot.
325
326 The return value is [x0, y0, x1, y1] in world coordinates.
327 """
328
329 def region_start(event):
330 height = self.canvas.figure.bbox.height()
331 self.rect = {'fig': None, 'height': height,
332 'x': event.x, 'y': height - event.y,
333 'world': [event.xdata, event.ydata,
334 event.xdata, event.ydata]}
335 self.register('button_press', None)
336 self.register('motion_notify', region_draw)
337 self.register('button_release', region_disable)
338
339 def region_draw(event):
340 self.canvas._tkcanvas.delete(self.rect['fig'])
341 self.rect['fig'] = self.canvas._tkcanvas.create_rectangle(
342 self.rect['x'], self.rect['y'],
343 event.x, self.rect['height'] - event.y)
344
345 def region_disable(event):
346 self.register('motion_notify', None)
347 self.register('button_release', None)
348
349 self.canvas._tkcanvas.delete(self.rect['fig'])
350
351 self.rect['world'][2:4] = [event.xdata, event.ydata]
352 print '(%.2f, %.2f) (%.2f, %.2f)' % (self.rect['world'][0],
353 self.rect['world'][1], self.rect['world'][2],
354 self.rect['world'][3])
355
356 self.register('button_press', region_start)
357
358 # This has to be modified to block and return the result (currently
359 # printed by region_disable) when that becomes possible in matplotlib.
360
361 return [0.0, 0.0, 0.0, 0.0]
362
363
[120]364 def register(self, type=None, func=None):
365 """
366 Register, reregister, or deregister events of type 'button_press',
367 'button_release', or 'motion_notify'.
[482]368
[120]369 The specified callback function should have the following signature:
370
371 def func(event)
372
373 where event is an MplEvent instance containing the following data:
374
375 name # Event name.
376 canvas # FigureCanvas instance generating the event.
377 x = None # x position - pixels from left of canvas.
378 y = None # y position - pixels from bottom of canvas.
379 button = None # Button pressed: None, 1, 2, 3.
380 key = None # Key pressed: None, chr(range(255)), shift,
381 win, or control
382 inaxes = None # Axes instance if cursor within axes.
383 xdata = None # x world coordinate.
384 ydata = None # y world coordinate.
385
386 For example:
387
388 def mouse_move(event):
389 print event.xdata, event.ydata
390
391 a = asaplot()
392 a.register('motion_notify', mouse_move)
393
394 If func is None, the event is deregistered.
395
396 Note that in TkAgg keyboard button presses don't generate an event.
397 """
398
399 if not self.events.has_key(type): return
400
401 if func is None:
402 if self.events[type] is not None:
403 # It's not clear that this does anything.
404 self.canvas.mpl_disconnect(self.events[type])
405 self.events[type] = None
406
407 # It seems to be necessary to return events to the toolbar.
408 if type == 'motion_notify':
409 self.canvas.mpl_connect(type + '_event',
410 self.figmgr.toolbar.mouse_move)
411 elif type == 'button_press':
412 self.canvas.mpl_connect(type + '_event',
413 self.figmgr.toolbar.press)
414 elif type == 'button_release':
415 self.canvas.mpl_connect(type + '_event',
416 self.figmgr.toolbar.release)
417
418 else:
419 self.events[type] = self.canvas.mpl_connect(type + '_event', func)
420
421
[111]422 def release(self):
[117]423 """
424 Release buffered graphics.
425 """
426 self.buffering = False
427 self.show()
[111]428
429
[671]430 def save(self, fname=None, orientation=None):
[482]431 """
432 Save the plot to a file.
433
434 fname is the name of the output file. The image format is determined
435 from the file suffix; 'png', 'ps', and 'eps' are recognized. If no
436 file name is specified 'yyyymmdd_hhmmss.png' is created in the current
437 directory.
438 """
439 if fname is None:
440 from datetime import datetime
441 dstr = datetime.now().strftime('%Y%m%d_%H%M%S')
442 fname = 'asap'+dstr+'.png'
443
444 d = ['png','.ps','eps']
445
446 from os.path import expandvars
447 fname = expandvars(fname)
448
449 if fname[-3:].lower() in d:
450 try:
[671]451 if fname[-3:].lower() == ".ps":
452 w = self.figure.figwidth.get()
453 h = self.figure.figheight.get()
454 a4w = 8.25
455 a4h = 11.25
456
457 if orientation is None:
458 # auto oriented
459 if w > h:
460 orientation = 'landscape'
461 else:
462 orientation = 'portrait'
463 ds = None
464 if orientation == 'landscape':
465 ds = min(a4h/w,a4w/h)
466 #self.figure.set_figsize_inches((a4h,a4w))
467 else:
468 ds = min(a4w/w,a4h/h)
469 ow = ds * w
470 oh = ds * h
471 self.figure.set_figsize_inches((ow,oh))
472 self.canvas.print_figure(fname,orientation=orientation)
473 print 'Written file %s' % (fname)
474 else:
475 self.canvas.print_figure(fname)
476 print 'Written file %s' % (fname)
[482]477 except IOError, msg:
478 print 'Failed to save %s: Error msg was\n\n%s' % (fname, err)
479 return
480 else:
481 print "Invalid image type. Valid types are:"
[671]482 print "'ps', 'eps', 'png'"
[482]483
484
[111]485 def set_axes(self, what=None, *args, **kwargs):
[117]486 """
487 Set attributes for the axes by calling the relevant Axes.set_*()
488 method. Colour translation is done as described in the doctext
489 for palette().
490 """
[111]491
[117]492 if what is None: return
493 if what[-6:] == 'colour': what = what[:-6] + 'color'
[111]494
[117]495 newargs = {}
[652]496
[117]497 for k, v in kwargs.iteritems():
498 k = k.lower()
499 if k == 'colour': k = 'color'
500 newargs[k] = v
[111]501
[117]502 getattr(self.axes, "set_%s"%what)(*args, **newargs)
503 self.show()
[111]504
505
506 def set_figure(self, what=None, *args, **kwargs):
[117]507 """
508 Set attributes for the figure by calling the relevant Figure.set_*()
509 method. Colour translation is done as described in the doctext
510 for palette().
511 """
[111]512
[117]513 if what is None: return
514 if what[-6:] == 'colour': what = what[:-6] + 'color'
[652]515 #if what[-5:] == 'color' and len(args):
516 # args = (get_colour(args[0]),)
[111]517
[117]518 newargs = {}
519 for k, v in kwargs.iteritems():
520 k = k.lower()
521 if k == 'colour': k = 'color'
522 newargs[k] = v
[111]523
[117]524 getattr(self.figure, "set_%s"%what)(*args, **newargs)
525 self.show()
[111]526
527
[482]528 def set_limits(self, xlim=None, ylim=None):
529 """
530 Set x-, and y-limits for each subplot.
531
532 xlim = [xmin, xmax] as in axes.set_xlim().
533 ylim = [ymin, ymax] as in axes.set_ylim().
534 """
535 for s in self.subplots:
536 self.axes = s['axes']
537 self.lines = s['lines']
[620]538 oldxlim = list(self.axes.get_xlim())
539 oldylim = list(self.axes.get_ylim())
540 if xlim is not None:
541 for i in range(len(xlim)):
542 if xlim[i] is not None:
543 oldxlim[i] = xlim[i]
544 if ylim is not None:
545 for i in range(len(ylim)):
546 if ylim[i] is not None:
547 oldylim[i] = ylim[i]
548 self.axes.set_xlim(oldxlim)
549 self.axes.set_ylim(oldylim)
550 return
[482]551
552
[111]553 def set_line(self, number=None, **kwargs):
[117]554 """
555 Set attributes for the specified line, or else the next line(s)
556 to be plotted.
[111]557
[117]558 number is the 0-relative number of a line that has already been
559 plotted. If no such line exists, attributes are recorded and used
560 for the next line(s) to be plotted.
[111]561
[117]562 Keyword arguments specify Line2D attributes, e.g. color='r'. Do
[111]563
[117]564 import matplotlib
565 help(matplotlib.lines)
[111]566
[117]567 The set_* methods of class Line2D define the attribute names and
568 values. For non-US usage, "colour" is recognized as synonymous with
569 "color".
[111]570
[117]571 Set the value to None to delete an attribute.
[111]572
[117]573 Colour translation is done as described in the doctext for palette().
574 """
[111]575
[117]576 redraw = False
577 for k, v in kwargs.iteritems():
578 k = k.lower()
579 if k == 'colour': k = 'color'
[111]580
[117]581 if 0 <= number < len(self.lines):
582 if self.lines[number] is not None:
583 for line in self.lines[number]:
584 getattr(line, "set_%s"%k)(v)
585 redraw = True
586 else:
587 if v is None:
588 del self.attributes[k]
589 else:
590 self.attributes[k] = v
[111]591
[117]592 if redraw: self.show()
[111]593
594
[620]595 def set_panels(self, rows=1, cols=0, n=-1, nplots=-1, ganged=True):
[118]596 """
597 Set the panel layout.
[482]598
[118]599 rows and cols, if cols != 0, specify the number of rows and columns in
600 a regular layout. (Indexing of these panels in matplotlib is row-
601 major, i.e. column varies fastest.)
602
603 cols == 0 is interpreted as a retangular layout that accomodates
604 'rows' panels, e.g. rows == 6, cols == 0 is equivalent to
605 rows == 2, cols == 3.
606
607 0 <= n < rows*cols is interpreted as the 0-relative panel number in
608 the configuration specified by rows and cols to be added to the
609 current figure as its next 0-relative panel number (i). This allows
610 non-regular panel layouts to be constructed via multiple calls. Any
611 other value of n clears the plot and produces a rectangular array of
[482]612 empty panels. The number of these may be limited by nplots.
[118]613 """
614 if n < 0 and len(self.subplots):
615 self.figure.clear()
616 self.set_title()
617
[482]618 if rows < 1: rows = 1
619
620 if cols <= 0:
[118]621 i = int(sqrt(rows))
622 if i*i < rows: i += 1
623 cols = i
624
625 if i*(i-1) >= rows: i -= 1
626 rows = i
[482]627
[118]628 if 0 <= n < rows*cols:
629 i = len(self.subplots)
630 self.subplots.append({})
[482]631
[118]632 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
633 cols, n+1)
634 self.subplots[i]['lines'] = []
635
[119]636 if i == 0: self.subplot(0)
[118]637
[482]638 self.rows = 0
639 self.cols = 0
640
[118]641 else:
642 self.subplots = []
[482]643
644 if nplots < 1 or rows*cols < nplots:
645 nplots = rows*cols
646
647 for i in range(nplots):
[118]648 self.subplots.append({})
[482]649
[118]650 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
651 cols, i+1)
652 self.subplots[i]['lines'] = []
[620]653 xfsize = self.subplots[i]['axes'].xaxis.label.get_size()-cols/2
654 yfsize = self.subplots[i]['axes'].yaxis.label.get_size()-rows/2
655 self.subplots[i]['axes'].xaxis.label.set_size(xfsize)
656 self.subplots[i]['axes'].yaxis.label.set_size(yfsize)
657
658 if ganged:
659 if rows > 1 or cols > 1:
660 # Squeeze the plots together.
661 pos = self.subplots[i]['axes'].get_position()
662 if cols > 1: pos[2] *= 1.2
663 if rows > 1: pos[3] *= 1.2
664 self.subplots[i]['axes'].set_position(pos)
[118]665
[620]666 # Suppress tick labelling for interior subplots.
667 if i <= (rows-1)*cols - 1:
668 if i+cols < nplots:
669 # Suppress x-labels for frames width
670 # adjacent frames
671 for tick in \
672 self.subplots[i]['axes'].xaxis.majorTicks:
673 tick.label1On = False
674 self.subplots[i]['axes'].xaxis.label.set_visible(False)
675 if i%cols:
676 # Suppress y-labels for frames not in the left column.
677 for tick in self.subplots[i]['axes'].yaxis.majorTicks:
678 tick.label1On = False
679 self.subplots[i]['axes'].yaxis.label.set_visible(False)
680
[482]681
682 self.rows = rows
683 self.cols = cols
684
[119]685 self.subplot(0)
[118]686
687 def set_title(self, title=None):
688 """
689 Set the title of the plot window. Use the previous title if title is
690 omitted.
691 """
692 if title is not None:
693 self.title = title
694
695 self.figure.text(0.5, 0.95, self.title, horizontalalignment='center')
696
697
[111]698 def show(self):
[117]699 """
700 Show graphics dependent on the current buffering state.
701 """
702 if not self.buffering:
[620]703 if self.loc is not None:
[482]704 for j in range(len(self.subplots)):
705 lines = []
706 labels = []
707 i = 0
708 for line in self.subplots[j]['lines']:
709 i += 1
710 if line is not None:
711 lines.append(line[0])
712 lbl = line[0].get_label()
713 if lbl == '':
714 lbl = str(i)
715 labels.append(lbl)
[111]716
[482]717 if len(lines):
718 self.subplots[j]['axes'].legend(tuple(lines),
719 tuple(labels),
720 self.loc)
721 else:
722 self.subplots[j]['axes'].legend((' '))
[111]723
[117]724 self.window.wm_deiconify()
725 self.canvas.show()
[111]726
[119]727 def subplot(self, i=None, inc=None):
[117]728 """
[118]729 Set the subplot to the 0-relative panel number as defined by one or
730 more invokations of set_panels().
[117]731 """
[118]732 l = len(self.subplots)
733 if l:
[119]734 if i is not None:
[120]735 self.i = i
[111]736
[119]737 if inc is not None:
[120]738 self.i += inc
[117]739
[119]740 self.i %= l
741 self.axes = self.subplots[self.i]['axes']
742 self.lines = self.subplots[self.i]['lines']
743
744
[111]745 def terminate(self):
[117]746 """
747 Clear the figure.
748 """
749 self.window.destroy()
[111]750
751
752 def text(self, *args, **kwargs):
[117]753 """
754 Add text to the figure.
755 """
756 self.figure.text(*args, **kwargs)
[482]757 self.show()
[111]758
[482]759
[111]760 def unmap(self):
[117]761 """
762 Hide the ASAPlot graphics window.
763 """
764 self.window.wm_withdraw()
Note: See TracBrowser for help on using the repository browser.