source: trunk/python/asaplot.py @ 371

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

added rcParams to support rc style default parameters, read from .asaprc

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 15.9 KB
Line 
1"""
2ASAP plotting class based on matplotlib.
3"""
4
5import sys
6from re import match
7import Tkinter as Tk
8
9#print "Importing matplotlib with TkAgg backend."
10import matplotlib
11matplotlib.use("TkAgg")
12
13from matplotlib.backends import new_figure_manager, show
14from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, \
15        FigureManagerTkAgg
16from matplotlib.figure import Figure, Text
17from matplotlib.numerix import sqrt
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=(7,5), 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        #self.frame1 = Tk.Frame(self.window, relief=Tk.RIDGE, borderwidth=4)
45        #self.frame1.pack(fill=Tk.BOTH)
46
47        self.figure = Figure(figsize=size,facecolor='#ddddee')
48        self.canvas = FigureCanvasTkAgg(self.figure, master=self.window)
49        self.canvas.get_tk_widget().pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
50
51        # Simply instantiating this is enough to get a working toolbar.
52        self.figmgr = FigureManagerTkAgg(self.canvas, 1, self.window)
53        self.window.wm_title('ASAPlot graphics window')
54
55        self.events = {'button_press':None,
56                       'button_release':None,
57                       'motion_notify':None}
58
59        self.set_title(title)
60        self.subplots = []
61        if rows > 0:
62            self.set_panels(rows, cols)
63
64
65        # Set matplotlib default colour sequence.
66        self.colours = [1, 'b', 'g', 'r', 'c', 'm', 'y', 'k']
67        self.attributes = {}
68        self.loc = 1
69
70        matplotlib.interactive = True
71        self.buffering = buffering
72
73        self.canvas.show()
74
75
76    def clear(self):
77        """
78        Delete all lines from the plot.  Line numbering will restart from 1.
79        """
80
81        for i in range(1,len(self.lines)+1):
82           self.delete(i)
83
84        self.axes.clear()
85        self.colours[0] = 1
86        self.lines = []
87       
88
89    def delete(self, numbers=None):
90        """
91        Delete the 0-relative line number, default is to delete the last.
92        The remaining lines are NOT renumbered.
93        """
94
95        if numbers is None: numbers = [len(self.lines)-1]
96
97        if not hasattr(numbers, '__iter__'):
98            numbers = [numbers]
99
100        for number in numbers:
101            if 0 <= number < len(self.lines):
102                if self.lines[number] is not None:
103                    for line in self.lines[number]:
104                        line.set_linestyle('None')
105                        self.lines[number] = None
106
107        self.show()
108
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=1):
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 1 > loc > 10: loc = 0
171        self.loc = loc
172        self.show()
173
174
175    def map(self):
176        """
177        Reveal the ASAPlot graphics window and bring it to the top of the
178        window stack.
179        """
180        self.window.wm_deiconify()
181        self.window.lift()
182
183
184    def palette(self, pen=None, colours=None):
185        """
186        Redefine the colour sequence.
187
188        pen is the pen number to use for the next plot; this will be auto-
189        incremented.
190
191        colours is the list of pen colours.  Colour may be specified via
192        the single letter values understood by matplotlib:
193
194            b: blue
195            g: green
196            r: red
197            c: cyan
198            m: magenta
199            y: yellow
200            k: black
201            w: white
202
203        or via the full name as listed in the colour dictionary which is
204        loaded by default by load_colours() from rgb.txt and listed by
205        list_colours().
206        """
207
208        if pen is None and colours is None:
209            self.colours = []
210            return
211
212        if pen is None:
213            if not len(self.colours):
214                self.colours = [1]
215        else:
216            self.colours[0] = pen
217
218        if colours is None:
219            return
220
221        cols = []
222        for col in colours:
223            cols.append(get_colour(col))
224
225        self.colours[1:] = cols
226
227        if 0 > self.colours[0] > len(self.colours):
228            self.colours[0] = 1
229
230
231    def plot(self, x=None, y=None, mask=None, fmt=None, add=None):
232        """
233        Plot the next line in the current frame using the current line
234        attributes.  The ASAPlot graphics window will be mapped and raised.
235
236        The argument list works a bit like the matlab plot() function.
237        """
238
239        if x is None:
240            if y is None: return
241            x = range(len(y))
242
243        elif y is None:
244            y = x
245            x = range(len(y))
246
247        if mask is None:
248            if fmt is None:
249                line = self.axes.plot(x, y)
250            else:
251                line = self.axes.plot(x, y, fmt)
252        else:
253            segments = []
254
255            mask = list(mask)
256            i = 0
257            while mask[i:].count(1):
258                i += mask[i:].index(1)
259                if mask[i:].count(0):
260                    j = i + mask[i:].index(0)
261                else:
262                    j = len(mask)
263
264                segments.append(x[i:j])
265                segments.append(y[i:j])
266
267                i = j
268
269            line = self.axes.plot(*segments)
270
271        # Add to an existing line?
272        if add is None or len(self.lines) < add < 0:
273            # Don't add.
274            self.lines.append(line)
275            i = len(self.lines) - 1
276        else:
277            if add == 0: add = len(self.lines)
278            i = add - 1
279            self.lines[i].extend(line)
280
281        # Set/reset attributes for the line.
282        gotcolour = False
283        for k, v in self.attributes.iteritems():
284            if k == 'color': gotcolour = True
285            for segment in self.lines[i]:
286                getattr(segment, "set_%s"%k)(v)
287
288        if not gotcolour and len(self.colours):
289            for segment in self.lines[i]:
290                getattr(segment, "set_color")(self.colours[self.colours[0]])
291
292            self.colours[0] += 1
293            if self.colours[0] >= len(self.colours):
294                self.colours[0] = 1
295
296        self.show()
297
298
299    def quit(self):
300        """
301        Destroy the ASAPlot graphics window.
302        """
303        self.window.destroy()
304
305
306    def register(self, type=None, func=None):
307        """
308        Register, reregister, or deregister events of type 'button_press',
309        'button_release', or 'motion_notify'.
310       
311        The specified callback function should have the following signature:
312
313            def func(event)
314
315        where event is an MplEvent instance containing the following data:
316
317            name                # Event name.
318            canvas              # FigureCanvas instance generating the event.
319            x      = None       # x position - pixels from left of canvas.
320            y      = None       # y position - pixels from bottom of canvas.
321            button = None       # Button pressed: None, 1, 2, 3.
322            key    = None       # Key pressed: None, chr(range(255)), shift,
323                                  win, or control
324            inaxes = None       # Axes instance if cursor within axes.
325            xdata  = None       # x world coordinate.
326            ydata  = None       # y world coordinate.
327
328        For example:
329
330            def mouse_move(event):
331                print event.xdata, event.ydata
332
333            a = asaplot()
334            a.register('motion_notify', mouse_move)
335
336        If func is None, the event is deregistered.
337
338        Note that in TkAgg keyboard button presses don't generate an event.
339        """
340
341        if not self.events.has_key(type): return
342
343        if func is None:
344            if self.events[type] is not None:
345                # It's not clear that this does anything.
346                self.canvas.mpl_disconnect(self.events[type])
347                self.events[type] = None
348
349                # It seems to be necessary to return events to the toolbar.
350                if type == 'motion_notify':
351                    self.canvas.mpl_connect(type + '_event',
352                        self.figmgr.toolbar.mouse_move)
353                elif type == 'button_press':
354                    self.canvas.mpl_connect(type + '_event',
355                        self.figmgr.toolbar.press)
356                elif type == 'button_release':
357                    self.canvas.mpl_connect(type + '_event',
358                        self.figmgr.toolbar.release)
359
360        else:
361            self.events[type] = self.canvas.mpl_connect(type + '_event', func)
362
363
364    def release(self):
365        """
366        Release buffered graphics.
367        """
368        self.buffering = False
369        self.show()
370
371
372    def set_axes(self, what=None, *args, **kwargs):
373        """
374        Set attributes for the axes by calling the relevant Axes.set_*()
375        method.  Colour translation is done as described in the doctext
376        for palette().
377        """
378
379        if what is None: return
380        if what[-6:] == 'colour': what = what[:-6] + 'color'
381
382        newargs = {}
383        for k, v in kwargs.iteritems():
384            k = k.lower()
385            if k == 'colour': k = 'color'
386
387            if k == 'color':
388                v = get_colour(v)
389
390            newargs[k] = v
391
392        getattr(self.axes, "set_%s"%what)(*args, **newargs)
393        self.show()
394
395
396    def set_figure(self, what=None, *args, **kwargs):
397        """
398        Set attributes for the figure by calling the relevant Figure.set_*()
399        method.  Colour translation is done as described in the doctext
400        for palette().
401        """
402
403        if what is None: return
404        if what[-6:] == 'colour': what = what[:-6] + 'color'
405        if what[-5:] == 'color' and len(args):
406            args = (get_colour(args[0]),)
407
408        newargs = {}
409        for k, v in kwargs.iteritems():
410            k = k.lower()
411            if k == 'colour': k = 'color'
412
413            if k == 'color':
414                v = get_colour(v)
415
416            newargs[k] = v
417
418        getattr(self.figure, "set_%s"%what)(*args, **newargs)
419        self.show()
420
421
422    def set_line(self, number=None, **kwargs):
423        """
424        Set attributes for the specified line, or else the next line(s)
425        to be plotted.
426
427        number is the 0-relative number of a line that has already been
428        plotted.  If no such line exists, attributes are recorded and used
429        for the next line(s) to be plotted.
430
431        Keyword arguments specify Line2D attributes, e.g. color='r'.  Do
432
433            import matplotlib
434            help(matplotlib.lines)
435
436        The set_* methods of class Line2D define the attribute names and
437        values.  For non-US usage, "colour" is recognized as synonymous with
438        "color".
439
440        Set the value to None to delete an attribute.
441
442        Colour translation is done as described in the doctext for palette().
443        """
444
445        redraw = False
446        for k, v in kwargs.iteritems():
447            k = k.lower()
448            if k == 'colour': k = 'color'
449
450            if k == 'color':
451                v = get_colour(v)
452
453            if 0 <= number < len(self.lines):
454                if self.lines[number] is not None:
455                    for line in self.lines[number]:
456                        getattr(line, "set_%s"%k)(v)
457                    redraw = True
458            else:
459                if v is None:
460                    del self.attributes[k]
461                else:
462                    self.attributes[k] = v
463
464        if redraw: self.show()
465
466
467    def set_panels(self, rows=1, cols=0, n=-1):
468        """
469        Set the panel layout.
470       
471        rows and cols, if cols != 0, specify the number of rows and columns in
472        a regular layout.   (Indexing of these panels in matplotlib is row-
473        major, i.e. column varies fastest.)
474
475        cols == 0 is interpreted as a retangular layout that accomodates
476        'rows' panels, e.g. rows == 6, cols == 0 is equivalent to
477        rows == 2, cols == 3.
478
479        0 <= n < rows*cols is interpreted as the 0-relative panel number in
480        the configuration specified by rows and cols to be added to the
481        current figure as its next 0-relative panel number (i).  This allows
482        non-regular panel layouts to be constructed via multiple calls.  Any
483        other value of n clears the plot and produces a rectangular array of
484        empty panels.
485        """
486        if n < 0 and len(self.subplots):
487            self.figure.clear()
488            self.set_title()
489
490        if rows < 1:
491            rows = 1
492        nel = 1
493        if cols == 0:
494            nel = rows
495            i = int(sqrt(rows))
496            if i*i < rows: i += 1
497            cols = i
498
499            if i*(i-1) >= rows: i -= 1
500            rows = i
501           
502        if 0 <= n < rows*cols:
503            i = len(self.subplots)
504            self.subplots.append({})
505            self.subplots[i]['axes']  = self.figure.add_subplot(rows,
506                                            cols, n+1)
507            self.subplots[i]['lines'] = []
508
509            if i == 0: self.subplot(0)
510
511        else:
512            self.subplots = []
513            for i in range(0,nel):
514                self.subplots.append({})
515                self.subplots[i]['axes']  = self.figure.add_subplot(rows,
516                                                cols, i+1)
517                self.subplots[i]['lines'] = []
518
519            self.subplot(0)
520
521
522    def set_title(self, title=None):
523        """
524        Set the title of the plot window.  Use the previous title if title is
525        omitted.
526        """
527        if title is not None:
528            self.title = title
529
530        self.figure.text(0.5, 0.95, self.title, horizontalalignment='center')
531
532
533    def show(self):
534        """
535        Show graphics dependent on the current buffering state.
536        """
537        if not self.buffering:
538            if self.loc:
539                for j in range(len(self.subplots)):
540                    lines  = []
541                    labels = []
542                    i = 0
543                    for line in self.subplots[j]['lines']:
544                        i += 1
545                        if line is not None:
546                            lines.append(line[0])
547                            lbl = line[0].get_label()
548                            if lbl == '':
549                                lbl = str(i)
550                            labels.append(lbl)
551
552                    if len(lines):
553                        self.subplots[j]['axes'].legend(tuple(lines), tuple(labels), self.loc)
554                    else:
555                        self.subplots[j]['axes'].legend((' '))
556
557            self.window.wm_deiconify()
558            self.canvas.show()
559
560
561    def subplot(self, i=None, inc=None):
562        """
563        Set the subplot to the 0-relative panel number as defined by one or
564        more invokations of set_panels().
565        """
566        l = len(self.subplots)
567        if l:
568            if i is not None:
569                self.i = i
570
571            if inc is not None:
572                self.i += inc
573
574            self.i %= l
575            self.axes  = self.subplots[self.i]['axes']
576            self.lines = self.subplots[self.i]['lines']
577
578
579    def terminate(self):
580        """
581        Clear the figure.
582        """
583        self.window.destroy()
584
585
586    def text(self, *args, **kwargs):
587        """
588        Add text to the figure.
589        """
590        self.figure.text(*args, **kwargs)
591        self.show()
592
593
594    def unmap(self):
595        """
596        Hide the ASAPlot graphics window.
597        """
598        self.window.wm_withdraw()
599
600    def set_limits(self,xlim=None,ylim=None):
601        for s in self.subplots:
602            self.axes  = s['axes']
603            self.lines = s['lines']
604            if xlim is not None:
605                self.axes.set_xlim(xlim)
606            if ylim is not None:
607                self.axes.set_ylim(ylim)
608        return
609
610    def save(self, fname=None):
611        if fname is None:
612            from datetime import datetime
613            dstr = datetime.now().strftime('%Y%m%d_%H%M%S')
614            fname = 'asap'+dstr+'.png'
615           
616        d = ['png','.ps','eps']
617        if fname[-3:].lower() in d:
618            try:
619                self.canvas.print_figure(fname)
620            except IOError, msg:
621                print 'Failed to save %s: Error msg was\n\n%s' % (fname, err)
622                return
623        else:
624            print "Invalid image type. Valid types are:"
625            print d
626
627
628def get_colour(colour='black'):
629    """
630    Look up a colour by name in the colour dictionary.  Matches are
631    case-insensitive, insensitive to blanks, and 'gray' matches 'grey'.
632    """
633
634    if colour is None: return None
635
636    if match('[rgbcmykw]$', colour): return colour
637    if match('#[\da-fA-F]{6}$', colour): return colour
638
639    if len(colours) == 0: load_colours()
640
641    # Try a quick match.
642    if colours.has_key(colour): return colours[colour]
643
644    colour = colour.replace(' ','').lower()
645    colour = colour.replace('gray','grey')
646    for name in colours.keys():
647        if name.lower() == colour:
648            return colours[name]
649
650    return '#000000'
651
652
653def list_colours():
654    """
655    List the contents of the colour dictionary sorted by name.
656    """
657
658    if len(colours) == 0: load_colours()
659
660    names = colours.keys()
661    names.sort()
662    for name in names:
663        print colours[name], name
664
665
666def load_colours(file='/usr/local/lib/rgb.txt'):
667    """
668    Load the colour dictionary from the specified file.
669    """
670    print 'Loading colour dictionary from', file
671    rgb = open(file, 'r')
672
673    while True:
674        line = rgb.readline()
675        if line == '': break
676        tmp = line.split()
677
678        if len(tmp) == 4:
679            if tmp[3][:4] == 'gray': continue
680            if tmp[3].lower().find('gray') != -1: continue
681
682            name = tmp[3][0].upper() + tmp[3][1:]
683            r, g, b = int(tmp[0]), int(tmp[1]), int(tmp[2])
684            colours[name] = '#%2.2x%2.2x%2.2x' % (r, g, b)
Note: See TracBrowser for help on using the repository browser.