source: trunk/python/asaplot.py@ 657

Last change on this file since 657 was 652, checked in by mar637, 19 years ago

removed color loading as mpl now supports named colors. some minor corrections on pol label handling. Also added orientation option for ps output.

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