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

Last change on this file since 2274 was 693, checked in by mar637, 19 years ago

added user customisable color and linestyles

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