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

Last change on this file since 691 was 671, checked in by mar637, 19 years ago

Fixed ps bug, where aspect ratios and sizes weren't handled properly. NOTE. A4 is hardcoded. This needs to be changed.

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 18.5 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, orientation=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 if fname[-3:].lower() == ".ps":
404 w = self.figure.figwidth.get()
405 h = self.figure.figheight.get()
406 a4w = 8.25
407 a4h = 11.25
408
409 if orientation is None:
410 # auto oriented
411 if w > h:
412 orientation = 'landscape'
413 else:
414 orientation = 'portrait'
415 ds = None
416 if orientation == 'landscape':
417 ds = min(a4h/w,a4w/h)
418 #self.figure.set_figsize_inches((a4h,a4w))
419 else:
420 ds = min(a4w/w,a4h/h)
421 ow = ds * w
422 oh = ds * h
423 self.figure.set_figsize_inches((ow,oh))
424 self.canvas.print_figure(fname,orientation=orientation)
425 print 'Written file %s' % (fname)
426 else:
427 self.canvas.print_figure(fname)
428 print 'Written file %s' % (fname)
429 except IOError, msg:
430 print 'Failed to save %s: Error msg was\n\n%s' % (fname, err)
431 return
432 else:
433 print "Invalid image type. Valid types are:"
434 print "'ps', 'eps', 'png'"
435
436
437 def set_axes(self, what=None, *args, **kwargs):
438 """
439 Set attributes for the axes by calling the relevant Axes.set_*()
440 method. Colour translation is done as described in the doctext
441 for palette().
442 """
443
444 if what is None: return
445 if what[-6:] == 'colour': what = what[:-6] + 'color'
446
447 newargs = {}
448
449 for k, v in kwargs.iteritems():
450 k = k.lower()
451 if k == 'colour': k = 'color'
452 newargs[k] = v
453
454 getattr(self.axes, "set_%s"%what)(*args, **newargs)
455 self.show()
456
457
458 def set_figure(self, what=None, *args, **kwargs):
459 """
460 Set attributes for the figure by calling the relevant Figure.set_*()
461 method. Colour translation is done as described in the doctext
462 for palette().
463 """
464
465 if what is None: return
466 if what[-6:] == 'colour': what = what[:-6] + 'color'
467 #if what[-5:] == 'color' and len(args):
468 # args = (get_colour(args[0]),)
469
470 newargs = {}
471 for k, v in kwargs.iteritems():
472 k = k.lower()
473 if k == 'colour': k = 'color'
474 newargs[k] = v
475
476 getattr(self.figure, "set_%s"%what)(*args, **newargs)
477 self.show()
478
479
480 def set_limits(self, xlim=None, ylim=None):
481 """
482 Set x-, and y-limits for each subplot.
483
484 xlim = [xmin, xmax] as in axes.set_xlim().
485 ylim = [ymin, ymax] as in axes.set_ylim().
486 """
487 for s in self.subplots:
488 self.axes = s['axes']
489 self.lines = s['lines']
490 oldxlim = list(self.axes.get_xlim())
491 oldylim = list(self.axes.get_ylim())
492 if xlim is not None:
493 for i in range(len(xlim)):
494 if xlim[i] is not None:
495 oldxlim[i] = xlim[i]
496 if ylim is not None:
497 for i in range(len(ylim)):
498 if ylim[i] is not None:
499 oldylim[i] = ylim[i]
500 self.axes.set_xlim(oldxlim)
501 self.axes.set_ylim(oldylim)
502 return
503
504
505 def set_line(self, number=None, **kwargs):
506 """
507 Set attributes for the specified line, or else the next line(s)
508 to be plotted.
509
510 number is the 0-relative number of a line that has already been
511 plotted. If no such line exists, attributes are recorded and used
512 for the next line(s) to be plotted.
513
514 Keyword arguments specify Line2D attributes, e.g. color='r'. Do
515
516 import matplotlib
517 help(matplotlib.lines)
518
519 The set_* methods of class Line2D define the attribute names and
520 values. For non-US usage, "colour" is recognized as synonymous with
521 "color".
522
523 Set the value to None to delete an attribute.
524
525 Colour translation is done as described in the doctext for palette().
526 """
527
528 redraw = False
529 for k, v in kwargs.iteritems():
530 k = k.lower()
531 if k == 'colour': k = 'color'
532
533 if 0 <= number < len(self.lines):
534 if self.lines[number] is not None:
535 for line in self.lines[number]:
536 getattr(line, "set_%s"%k)(v)
537 redraw = True
538 else:
539 if v is None:
540 del self.attributes[k]
541 else:
542 self.attributes[k] = v
543
544 if redraw: self.show()
545
546
547 def set_panels(self, rows=1, cols=0, n=-1, nplots=-1, ganged=True):
548 """
549 Set the panel layout.
550
551 rows and cols, if cols != 0, specify the number of rows and columns in
552 a regular layout. (Indexing of these panels in matplotlib is row-
553 major, i.e. column varies fastest.)
554
555 cols == 0 is interpreted as a retangular layout that accomodates
556 'rows' panels, e.g. rows == 6, cols == 0 is equivalent to
557 rows == 2, cols == 3.
558
559 0 <= n < rows*cols is interpreted as the 0-relative panel number in
560 the configuration specified by rows and cols to be added to the
561 current figure as its next 0-relative panel number (i). This allows
562 non-regular panel layouts to be constructed via multiple calls. Any
563 other value of n clears the plot and produces a rectangular array of
564 empty panels. The number of these may be limited by nplots.
565 """
566 if n < 0 and len(self.subplots):
567 self.figure.clear()
568 self.set_title()
569
570 if rows < 1: rows = 1
571
572 if cols <= 0:
573 i = int(sqrt(rows))
574 if i*i < rows: i += 1
575 cols = i
576
577 if i*(i-1) >= rows: i -= 1
578 rows = i
579
580 if 0 <= n < rows*cols:
581 i = len(self.subplots)
582 self.subplots.append({})
583
584 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
585 cols, n+1)
586 self.subplots[i]['lines'] = []
587
588 if i == 0: self.subplot(0)
589
590 self.rows = 0
591 self.cols = 0
592
593 else:
594 self.subplots = []
595
596 if nplots < 1 or rows*cols < nplots:
597 nplots = rows*cols
598
599 for i in range(nplots):
600 self.subplots.append({})
601
602 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
603 cols, i+1)
604 self.subplots[i]['lines'] = []
605 xfsize = self.subplots[i]['axes'].xaxis.label.get_size()-cols/2
606 yfsize = self.subplots[i]['axes'].yaxis.label.get_size()-rows/2
607 self.subplots[i]['axes'].xaxis.label.set_size(xfsize)
608 self.subplots[i]['axes'].yaxis.label.set_size(yfsize)
609
610 if ganged:
611 if rows > 1 or cols > 1:
612 # Squeeze the plots together.
613 pos = self.subplots[i]['axes'].get_position()
614 if cols > 1: pos[2] *= 1.2
615 if rows > 1: pos[3] *= 1.2
616 self.subplots[i]['axes'].set_position(pos)
617
618 # Suppress tick labelling for interior subplots.
619 if i <= (rows-1)*cols - 1:
620 if i+cols < nplots:
621 # Suppress x-labels for frames width
622 # adjacent frames
623 for tick in \
624 self.subplots[i]['axes'].xaxis.majorTicks:
625 tick.label1On = False
626 self.subplots[i]['axes'].xaxis.label.set_visible(False)
627 if i%cols:
628 # Suppress y-labels for frames not in the left column.
629 for tick in self.subplots[i]['axes'].yaxis.majorTicks:
630 tick.label1On = False
631 self.subplots[i]['axes'].yaxis.label.set_visible(False)
632
633
634 self.rows = rows
635 self.cols = cols
636
637 self.subplot(0)
638
639 def set_title(self, title=None):
640 """
641 Set the title of the plot window. Use the previous title if title is
642 omitted.
643 """
644 if title is not None:
645 self.title = title
646
647 self.figure.text(0.5, 0.95, self.title, horizontalalignment='center')
648
649
650 def show(self):
651 """
652 Show graphics dependent on the current buffering state.
653 """
654 if not self.buffering:
655 if self.loc is not None:
656 for j in range(len(self.subplots)):
657 lines = []
658 labels = []
659 i = 0
660 for line in self.subplots[j]['lines']:
661 i += 1
662 if line is not None:
663 lines.append(line[0])
664 lbl = line[0].get_label()
665 if lbl == '':
666 lbl = str(i)
667 labels.append(lbl)
668
669 if len(lines):
670 self.subplots[j]['axes'].legend(tuple(lines),
671 tuple(labels),
672 self.loc)
673 else:
674 self.subplots[j]['axes'].legend((' '))
675
676 self.window.wm_deiconify()
677 self.canvas.show()
678
679 def subplot(self, i=None, inc=None):
680 """
681 Set the subplot to the 0-relative panel number as defined by one or
682 more invokations of set_panels().
683 """
684 l = len(self.subplots)
685 if l:
686 if i is not None:
687 self.i = i
688
689 if inc is not None:
690 self.i += inc
691
692 self.i %= l
693 self.axes = self.subplots[self.i]['axes']
694 self.lines = self.subplots[self.i]['lines']
695
696
697 def terminate(self):
698 """
699 Clear the figure.
700 """
701 self.window.destroy()
702
703
704 def text(self, *args, **kwargs):
705 """
706 Add text to the figure.
707 """
708 self.figure.text(*args, **kwargs)
709 self.show()
710
711
712 def unmap(self):
713 """
714 Hide the ASAPlot graphics window.
715 """
716 self.window.wm_withdraw()
Note: See TracBrowser for help on using the repository browser.