source: branches/Release-2-fixes/python/asaplot.py @ 693

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

added user customisable color and linestyles

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