source: trunk/python/asaplot.py @ 491

Last change on this file since 491 was 491, checked in by cal103, 19 years ago

Suppressing tick labels by deleting the associated text kills position
tracking! It has to be done by setting label1On = False for the underlying
Tick objects.

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