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

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

Fixed ps bug, where aspect ratios and sizes weren't handled properly. NOTE. A4 is hardcoded. This needs to be changed.

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