source: branches/Release-1-fixes/python/asaplot.py@ 2129

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