source: trunk/python/asaplot.py @ 620

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

various plotter fixes from Release-1

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