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

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

adding contents of pywms

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