source: trunk/python/asaplot.py @ 482

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

In set_panels(), disabled labelling of interior subplots and squeezed out
excess space between rows and columns in a rectangular asaplot plot array.
Implemented cursor position readback and interactive region selection via
the new position() and region() methods, though neither of these will work
properly until a blocking mechanism is added to matplotlib (a known
deficiency).

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