source: trunk/python/asaplot.py@ 695

Last change on this file since 695 was 665, checked in by mar637, 19 years ago

casapmath.py

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 17.4 KB
RevLine 
[111]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, \
[117]14 FigureManagerTkAgg
[111]15from matplotlib.figure import Figure, Text
[491]16from matplotlib.font_manager import FontProperties
[118]17from matplotlib.numerix import sqrt
[482]18from matplotlib import rc, rcParams
[111]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
[482]28 def __init__(self, rows=1, cols=0, title='', size=(8,6), buffering=False):
[117]29 """
30 Create a new instance of the ASAPlot plotting class.
[119]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().
[117]34 """
35 self.window = Tk.Tk()
[482]36 self.is_dead = False
37 def dest_callback():
38 self.is_dead = True
39 self.window.destroy()
[111]40
[482]41 self.window.protocol("WM_DELETE_WINDOW", dest_callback)
42
43 self.figure = Figure(figsize=size, facecolor='#ddddee')
[202]44 self.canvas = FigureCanvasTkAgg(self.figure, master=self.window)
[117]45 self.canvas.get_tk_widget().pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
[111]46
[117]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')
[111]50
[120]51 self.events = {'button_press':None,
52 'button_release':None,
53 'motion_notify':None}
54
[118]55 self.set_title(title)
[117]56 self.subplots = []
[118]57 if rows > 0:
58 self.set_panels(rows, cols)
[111]59
60
[117]61 # Set matplotlib default colour sequence.
[652]62 self.colormap = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'purple', 'orange', 'pink']
63 self.color = 0;
[117]64 self.attributes = {}
[620]65 self.loc = 0
[111]66
[117]67 matplotlib.interactive = True
68 self.buffering = buffering
[482]69
[117]70 self.canvas.show()
[111]71
72
[117]73 def clear(self):
74 """
75 Delete all lines from the plot. Line numbering will restart from 1.
76 """
[111]77
[652]78 for i in range(len(self.lines)):
[117]79 self.delete(i)
80 self.axes.clear()
[652]81 self.color = 0
[117]82 self.lines = []
[111]83
[482]84
[652]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
[111]91 def delete(self, numbers=None):
[117]92 """
93 Delete the 0-relative line number, default is to delete the last.
94 The remaining lines are NOT renumbered.
95 """
[111]96
[117]97 if numbers is None: numbers = [len(self.lines)-1]
[111]98
[117]99 if not hasattr(numbers, '__iter__'):
100 numbers = [numbers]
[111]101
[117]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
[620]108 self.show()
[111]109
110 def get_line(self):
[117]111 """
112 Get the current default line attributes.
113 """
114 return self.attributes
[111]115
116
[119]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
[111]146 def hold(self, hold=True):
[117]147 """
148 Buffer graphics until subsequently released.
149 """
150 self.buffering = hold
[111]151
152
[620]153 def legend(self, loc=None):
[117]154 """
155 Add a legend to the plot.
[111]156
[117]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
[111]168
[117]169 """
[620]170 if isinstance(loc,int):
171 if 0 > loc > 10: loc = 0
172 self.loc = loc
[117]173 self.show()
[111]174
175
176 def map(self):
[117]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()
[111]183
[117]184
[111]185
186 def plot(self, x=None, y=None, mask=None, fmt=None, add=None):
[117]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.
[111]190
[117]191 The argument list works a bit like the matlab plot() function.
192 """
[111]193
[117]194 if x is None:
195 if y is None: return
196 x = range(len(y))
[111]197
[117]198 elif y is None:
199 y = x
200 x = range(len(y))
[111]201
[117]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 = []
[111]209
[117]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)
[111]218
[117]219 segments.append(x[i:j])
220 segments.append(y[i:j])
[111]221
[117]222 i = j
[111]223
[117]224 line = self.axes.plot(*segments)
[111]225
[117]226 # Add to an existing line?
227 if add is None or len(self.lines) < add < 0:
[119]228 # Don't add.
[117]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)
[111]235
[117]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)
[111]242
[652]243 if not gotcolour and len(self.colormap):
[117]244 for segment in self.lines[i]:
[652]245 getattr(segment, "set_color")(self.colormap[self.color])
[111]246
[652]247 self.color += 1
248 if self.color >= len(self.colormap):
249 self.color = 0
[111]250
[117]251 self.show()
[111]252
253
[482]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
[111]267 def quit(self):
[117]268 """
269 Destroy the ASAPlot graphics window.
270 """
271 self.window.destroy()
[111]272
273
[482]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
[120]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'.
[482]320
[120]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
[111]374 def release(self):
[117]375 """
376 Release buffered graphics.
377 """
378 self.buffering = False
379 self.show()
[111]380
381
[665]382 def save(self, fname=None, orientation='landscape'):
[482]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:
[665]403 self.canvas.print_figure(fname,orientation=orientation)
[482]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
[111]413 def set_axes(self, what=None, *args, **kwargs):
[117]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 """
[111]419
[117]420 if what is None: return
421 if what[-6:] == 'colour': what = what[:-6] + 'color'
[111]422
[117]423 newargs = {}
[652]424
[117]425 for k, v in kwargs.iteritems():
426 k = k.lower()
427 if k == 'colour': k = 'color'
428 newargs[k] = v
[111]429
[117]430 getattr(self.axes, "set_%s"%what)(*args, **newargs)
431 self.show()
[111]432
433
434 def set_figure(self, what=None, *args, **kwargs):
[117]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 """
[111]440
[117]441 if what is None: return
442 if what[-6:] == 'colour': what = what[:-6] + 'color'
[652]443 #if what[-5:] == 'color' and len(args):
444 # args = (get_colour(args[0]),)
[111]445
[117]446 newargs = {}
447 for k, v in kwargs.iteritems():
448 k = k.lower()
449 if k == 'colour': k = 'color'
450 newargs[k] = v
[111]451
[117]452 getattr(self.figure, "set_%s"%what)(*args, **newargs)
453 self.show()
[111]454
455
[482]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']
[620]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
[482]479
480
[111]481 def set_line(self, number=None, **kwargs):
[117]482 """
483 Set attributes for the specified line, or else the next line(s)
484 to be plotted.
[111]485
[117]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.
[111]489
[117]490 Keyword arguments specify Line2D attributes, e.g. color='r'. Do
[111]491
[117]492 import matplotlib
493 help(matplotlib.lines)
[111]494
[117]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".
[111]498
[117]499 Set the value to None to delete an attribute.
[111]500
[117]501 Colour translation is done as described in the doctext for palette().
502 """
[111]503
[117]504 redraw = False
505 for k, v in kwargs.iteritems():
506 k = k.lower()
507 if k == 'colour': k = 'color'
[111]508
[117]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
[111]519
[117]520 if redraw: self.show()
[111]521
522
[620]523 def set_panels(self, rows=1, cols=0, n=-1, nplots=-1, ganged=True):
[118]524 """
525 Set the panel layout.
[482]526
[118]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
[482]540 empty panels. The number of these may be limited by nplots.
[118]541 """
542 if n < 0 and len(self.subplots):
543 self.figure.clear()
544 self.set_title()
545
[482]546 if rows < 1: rows = 1
547
548 if cols <= 0:
[118]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
[482]555
[118]556 if 0 <= n < rows*cols:
557 i = len(self.subplots)
558 self.subplots.append({})
[482]559
[118]560 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
561 cols, n+1)
562 self.subplots[i]['lines'] = []
563
[119]564 if i == 0: self.subplot(0)
[118]565
[482]566 self.rows = 0
567 self.cols = 0
568
[118]569 else:
570 self.subplots = []
[482]571
572 if nplots < 1 or rows*cols < nplots:
573 nplots = rows*cols
574
575 for i in range(nplots):
[118]576 self.subplots.append({})
[482]577
[118]578 self.subplots[i]['axes'] = self.figure.add_subplot(rows,
579 cols, i+1)
580 self.subplots[i]['lines'] = []
[620]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)
[118]593
[620]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
[482]609
610 self.rows = rows
611 self.cols = cols
612
[119]613 self.subplot(0)
[118]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
[111]626 def show(self):
[117]627 """
628 Show graphics dependent on the current buffering state.
629 """
630 if not self.buffering:
[620]631 if self.loc is not None:
[482]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)
[111]644
[482]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((' '))
[111]651
[117]652 self.window.wm_deiconify()
653 self.canvas.show()
[111]654
[119]655 def subplot(self, i=None, inc=None):
[117]656 """
[118]657 Set the subplot to the 0-relative panel number as defined by one or
658 more invokations of set_panels().
[117]659 """
[118]660 l = len(self.subplots)
661 if l:
[119]662 if i is not None:
[120]663 self.i = i
[111]664
[119]665 if inc is not None:
[120]666 self.i += inc
[117]667
[119]668 self.i %= l
669 self.axes = self.subplots[self.i]['axes']
670 self.lines = self.subplots[self.i]['lines']
671
672
[111]673 def terminate(self):
[117]674 """
675 Clear the figure.
676 """
677 self.window.destroy()
[111]678
679
680 def text(self, *args, **kwargs):
[117]681 """
682 Add text to the figure.
683 """
684 self.figure.text(*args, **kwargs)
[482]685 self.show()
[111]686
[482]687
[111]688 def unmap(self):
[117]689 """
690 Hide the ASAPlot graphics window.
691 """
692 self.window.wm_withdraw()
Note: See TracBrowser for help on using the repository browser.