source: trunk/python/asaplot.py @ 376

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