source: trunk/python/asaplot.py@ 220

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