source: cows/trunk/cows/service/imps/pywms/figure.py @ 5399

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/cows/trunk/cows/service/imps/pywms/figure.py@5399
Revision 5399, 5.3 KB checked in by domlowe, 12 years ago (diff)

Changing PIL imports in cows

Line 
1# BSD Licence
2# Copyright (c) 2009, Science & Technology Facilities Council (STFC)
3# All rights reserved.
4#
5# See the LICENSE file in the source distribution of this software for
6# the full license text.
7
8"""
9Publication quality plotting engine.
10
11@author: Stephen Pascoe
12"""
13
14import os
15from tempfile import mkstemp
16from cStringIO import StringIO
17
18from matplotlib.toolkits.basemap import Basemap
19import pylab as p
20from matplotlib import cm
21try: 
22    from PIL import Image
23except ImportError:
24    import Image
25import Numeric as N
26import MA
27
28from pywms.utils import word_wrap
29
30class Figure(object):
31    """Render a grid to a publication-quality figure.
32
33    @cvar supportedFormats: A list of supported mime-types
34
35    @ivar grid: The grid to render
36    @ivar cmap: The colourmap to use for type 'colour'
37    @ivar caption: A caption under the figure
38    @ivar type: 'colour' or 'contour'
39    @ivar format: The mime-type to use when rendering
40    """
41
42    supportedFormats = ['image/png', 'image/jpeg', 'application/postscript', 'image/svg+xml']
43
44    def __init__(self, grid=None, cmap=cm.get_cmap, caption='', type='colour', format='image/png',
45                 vmin=None, vmax=None):
46        self.grid = grid
47        self.cmap = cmap
48        self.cmap.set_bad('#000000', alpha=0.0)
49        self.caption = caption
50        self.type = type
51        self.format = format
52        self.vmin=vmin
53        self.vmax=vmax
54
55
56
57    def makeFigure(self):
58        """
59        @todo: Uses pylab so may not be threadsafe.
60        """
61
62        lon = self.grid.lon0 + N.arrayrange(self.grid.nlon)*self.grid.dlon
63        lat = self.grid.lat0 + N.arrayrange(self.grid.nlat)*self.grid.dlat
64
65        ll_lon, ll_lat = lon[0], lat[0]
66        ur_lon, ur_lat = lon[-1], lat[-1]
67
68        # Account for variable lat/lon axis ordering
69        #!TODO: Could we move this (along with the equivilent in render_imp.py into grid.py?
70        if self.grid.ilat < self.grid.ilon:
71            latLonOrdering = True
72        else:
73            latLonOrdering = False
74
75        if latLonOrdering:
76            var = self.grid.value
77        else:
78            var = MA.transpose(self.grid.value)
79           
80
81        fig = p.figure()
82        map = Basemap(projection='cyl', llcrnrlon=ll_lon, llcrnrlat=ll_lat,
83                      urcrnrlon=ur_lon, urcrnrlat=ur_lat, resolution='l')
84
85##         if self.grid.units:
86##             p.title("%s\n(%s)" % (self.grid.long_name, self.grid.units))
87##         else:
88##             p.title(self.grid.long_name)
89        p.title(self.grid.long_name)
90
91        if self.type == 'colour':
92            # transform_scalar doesn't support masked arrays so we must fill then replace the mask.
93            var_dat = map.transform_scalar(var.filled(1.0e20), lon,
94                                           lat, len(lon), len(lat))
95            var = MA.masked_values(var_dat, 1.0e20)
96            map.imshow(var, cmap=self.cmap, vmin=self.vmin, vmax=self.vmax,
97                       interpolation='nearest')
98            cbar = p.colorbar(orientation='horizontal', format='%.2g')
99            if self.grid.units:
100                cbar.ax.set_xlabel(self.grid.units)
101           
102        else:
103            x, y = map(*p.meshgrid(lon, lat))
104            c = map.contour(x, y, var, 12, colors='black')
105            c.clabel(fontsize=8)
106            map.fillcontinents(color='#e0e0e0')
107
108        map.drawcoastlines(color='gray')
109
110        map.drawmeridians(p.arange(-180,180,30), labels=[1,0,0,1], color='gray')
111        map.drawparallels(p.arange(-90,90,15), labels=[1,0,0,1], color='gray')
112
113        # Wrap the caption
114        caption = word_wrap(self.caption, 80)
115
116        fig.text(0.1, 0.08, caption,
117                 fontsize=10, horizontalalignment='left',
118                 verticalalignment='top',
119                 transform=fig.transFigure
120                 )
121
122        return fig
123
124    def render(self, fh):
125        """Render the figure and serialise it to a file stream.
126
127        @param fh: A file object
128        """
129       
130        fig = self.makeFigure()
131       
132        ext = self.getExtension()
133
134        # matplotlib can't write jpeg on my installation
135        if ext == '.jpg':
136            convertTo = 'JPEG'
137            ext = '.png'
138        else:
139            convertTo = None
140
141        (fh1, fn) = mkstemp(ext, 'pywms_tmp_')
142        os.close(fh1)
143        fig.savefig(fn)
144
145        if convertTo:
146            img = Image.open(fn)
147            buf = StringIO()
148            img.save(buf, convertTo)
149            fh.write(buf.getvalue())
150        else:
151            fh.write(open(fn).read())
152
153        os.remove(fn)
154
155    def getExtension(self):
156        (mt, smt) = self.format.split('/')
157        if smt == 'png':
158            ext = '.png'
159        elif smt == 'jpeg':
160            ext = '.jpg'
161        elif smt == 'postscript':
162            ext = '.eps'
163        elif smt == 'svg+xml':
164            ext = '.svg'
165
166        return ext
167
168def makeTimeSeriesFigure(ts):
169    """
170    @todo: time axis labelling
171    @todo: Flesh out TimeSeries interface so that this function doesn't need to use cdms calls.
172    """
173    p.figure()
174    p.title(ts.long_name)
175    try:
176        p.ylabel(ts.units)
177    except AttributeError:
178        pass
179
180    p.plot(ts.value.getTime().getValue(), ts.value.filled(0.0), '-o')
181    (fh, fn) = mkstemp('.png', 'pywms_tmp_')
182    os.close(fh)
183    p.savefig(fn)
184    img = Image.open(fn); img.load()
185    os.remove(fn)
186
187    return img
Note: See TracBrowser for help on using the repository browser.