source: trunk/python/asaplot.py@ 458

Last change on this file since 458 was 411, checked in by mar637, 20 years ago

Added handling of environment variables throughout.

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