source: trunk/python/asaplotbase.py @ 705

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

split asaplot into gui and non-gui versions inheriting from asaplotbase

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