source: trunk/python/asaplot.py @ 652

Last change on this file since 652 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.