source: qesdi/geoplot/trunk/lib/geoplot/colour_bar.py @ 6103

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/qesdi/geoplot/trunk/lib/geoplot/colour_bar.py@6103
Revision 6103, 10.1 KB checked in by pnorton, 10 years ago (diff)

Trying to get the different colour bars and plot options to work with the different legend types.

Line 
1
2import math
3import logging
4
5from matplotlib.colors import Normalize, ListedColormap
6
7import matplotlib.ticker
8from matplotlib.figure import Figure
9from matplotlib.colorbar import ColorbarBase
10from matplotlib.patches import Rectangle
11import matplotlib.cm as cm
12import geoplot.config as geoplot_config
13import geoplot
14
15config = geoplot_config.getConfig()
16
17log = logging.getLogger(__name__)
18
19FONTS_SECTION = 'Fonts'
20MAX_CBAR_TICKS = 10
21ADJUSTED_TICK_FORMAT = "%1.2f"
22
23class COLOUR_BAR_STYLES:
24    LEGEND='legend'
25    CONTINUOUS = 'continuous' 
26    LINE = 'line'
27
28    @staticmethod
29    def all():
30        return [COLOUR_BAR_STYLES.CONTINUOUS, COLOUR_BAR_STYLES.LEGEND, COLOUR_BAR_STYLES.LINE]
31
32class COLOUR_SCHEME_SCALE:
33    LINEAR='linear'
34    LOG='log'
35
36class ColourBar(object):
37
38    def __init__(self, colourBarLabel="",colourBarPosition='horizontal', colourBarStyle=COLOUR_BAR_STYLES.CONTINUOUS):
39
40        self._position = None
41
42        self.colourBarLabel = colourBarLabel
43        self.colourBarPosition = colourBarPosition
44        self.colourBarStyle = colourBarStyle
45       
46        self.labelFont = config[FONTS_SECTION]['ColourBarLabel']
47
48    def draw(self, colourBarAxis, colourScheme, fontSize='medium'):
49        """
50        Adds the colour bar to the (and optionally a label) to the figure.
51
52        @param sm: the scalar mappable generated by applying the grid to the axis
53        @type sm: an instance of matplotlib.cm.ScalarMappable
54        @param units: the units of the values on the mesh, these will be used as a label
55            if the colourBarLabel property is not set.
56        @type units: string
57        """
58        log.debug("drawing colour bar")
59       
60        if self.colourBarStyle == COLOUR_BAR_STYLES.LEGEND:
61            self._drawLegendColourBar(colourBarAxis, colourScheme, fontSize)
62        elif self.colourBarStyle == COLOUR_BAR_STYLES.LINE:
63            self._drawLineColourBar(colourBarAxis, colourScheme, fontSize)
64        else:
65            self._drawContiunousColourBar(colourBarAxis, colourScheme, fontSize)
66               
67        log.debug("finished drawing colour bar")
68
69    def _drawContiunousColourBar(self, axes, colourScheme, fontSize):
70       
71        intervalColourbar = isinstance(colourScheme, geoplot.colour_scheme.IntervalColourScheme)
72         
73        kwargs = {}
74        kwargs['cmap'] = colourScheme.colourMap
75        kwargs['norm'] = colourScheme.norm
76       
77        kwargs['orientation'] = self.colourBarPosition
78       
79        if intervalColourbar:
80            kwargs['ticks'] = colourScheme.midpoints
81            kwargs['format'] = matplotlib.ticker.FixedFormatter(colourScheme.labels)
82            kwargs['spacing'] = 'proportional'
83       
84        cb = ColorbarBase(axes, **kwargs)
85
86        if cb.cmap.__class__ == ListedColormap \
87           and not intervalColourbar:
88            ColourBar._repositionColourBarTicks(cb)
89           
90        if self.colourBarLabel != None:
91            labelDictionary = self.labelFont.getDict(fontSize)
92            cb.set_label(self.colourBarLabel, fontdict=labelDictionary)
93       
94        return cb
95   
96    def _drawLegendColourBar(self, colourBarAxis, colourScheme, fontSize):
97        """
98       
99        """
100        kwargs = {}
101       
102        cmap = colourScheme.colourMap
103        norm = colourScheme.norm
104       
105        locations, labels = colourScheme.midpoints, colourScheme.labels
106
107        kwargs['orientation'] = self.colourBarPosition
108 
109        handles = [Rectangle((0,0), 1, 1, fc=cmap(norm(i))) for i in locations]
110        labels = labels
111       
112        if self.colourBarPosition == 'horizontal':
113            if len(handles) < 3:
114                ncol = len(handles)
115            else:
116                ncol = 3
117        else:
118            ncol = 1
119
120        leg = colourBarAxis.legend(handles, labels, loc=10, mode='expand', 
121                                   ncol=ncol, borderaxespad=0)
122       
123        colourBarAxis.set_xticks([])
124        colourBarAxis.set_yticks([])
125       
126        if self.colourBarLabel != None:
127            labelDictionary = self.labelFont.getDict(fontSize)
128           
129            if self.colourBarPosition == 'horizontal':
130                colourBarAxis.set_xlabel(self.colourBarLabel, fontdict=labelDictionary)     
131            else:
132                colourBarAxis.set_ylabel(self.colourBarLabel, fontdict=labelDictionary)
133               
134        return leg
135   
136    def _drawLineColourBar(self, colourBarAxis, colourScheme, fontSize):
137        """
138       
139        """
140        kwargs = {}
141       
142        cmap = colourScheme.colourMap
143        norm = colourScheme.norm
144        levels = colourScheme.bounds
145        colours = [cmap(norm(x)) for x in levels]
146        widths = [1.0 for x in levels]
147       
148       
149        cb = ColorbarBase(colourBarAxis, norm=norm, 
150                          orientation='horizontal', 
151                          filled=False)
152       
153        cb.add_lines(levels, colours, widths)
154
155#        colourBarAxis.set_xticks([])
156#        colourBarAxis.set_yticks([])
157       
158        if self.colourBarLabel != None:
159            labelDictionary = self.labelFont.getDict(fontSize)
160            cb.set_label(self.colourBarLabel, fontdict=labelDictionary)
161       
162        return cb
163       
164    @staticmethod
165    def _repositionColourBarTicks(cb):
166        """
167        reposition the ticks of a ListedColormap so that they appear at the
168        """
169       
170        log.debug("Repositioning colour bar ticks")
171        span = cb.vmax - cb.vmin
172
173        # Define flag for whether or not zero should be added
174        useZero = False
175        if cb.vmin < 0 and cb.vmax > 0:
176            useZero = True
177
178        numColours = len(cb.cmap.colors)
179        interval = float(span) / float(numColours)
180
181        showEvery = 1
182        while float(numColours)/float(showEvery) > float(MAX_CBAR_TICKS):
183            showEvery += 1
184
185        newLocs = []
186        for i in range(0, numColours + 1, showEvery):
187            newLocs.append(cb.vmin + i * interval)
188
189        # If need to add a zero then do so
190        if useZero == True and 0 not in newLocs:
191            newLocsWithZero = []
192         
193            zeroInserted = False 
194            for newLoc in newLocs:
195                if newLoc > 0 and zeroInserted == False:
196                    newLocsWithZero.append(0)
197                    zeroInserted = True                 
198                newLocsWithZero.append(newLoc)
199
200            newLocs = newLocsWithZero
201
202        #change the locator and the formatter (as the locations now have a high number of dp)
203        cb.locator = matplotlib.ticker.FixedLocator(newLocs)
204
205        # Decide the float formatting of the tick labels based on the span
206        if span < 100:
207            # In Excel this worked: =IF(E3<100,(LOG(E3)*-1)+2,0)
208            tickFormatDecPoints = int((math.log(span, 10) * -1) + 2)
209            tickFormatString = "%." + str(tickFormatDecPoints) + "f"
210        else: 
211            tickFormatString = "%d"
212           
213        cb.formatter = matplotlib.ticker.FormatStrFormatter(tickFormatString)  ###ADJUSTED_TICK_FORMAT)
214
215        # The next line removes axis artists as draw_all() adds new ones
216        cb.ax.artists = []
217        cb.draw_all()    # cause the colourbar to be redrawn, otherwise not changes will hapen
218
219        # Hard code line width of colour bar outline
220        cb.outline.set_linewidth(0.5)
221
222        # this can sometimes cause the tick positions to become unknown so
223        # set it back to default.
224        if cb.ax.yaxis.get_ticks_position() == 'unknown':
225            log.debug("Resetting yaxis ticks position to default")
226            cb.ax.yaxis.set_ticks_position('default')
227       
228        if cb.ax.xaxis.get_ticks_position() == 'unknown':
229            log.debug("Resetting xaxis ticks position to default")
230            cb.ax.xaxis.set_ticks_position('default')
231   
232 
233    def __get_position(self): return self._position
234
235    def __set_position(self, value):
236        if value not in ['horizontal', 'vertical', None]:
237            raise ValueError("ColourBar position value must be 'horizontal'" + \
238                             " or 'vertical, value recieved :" + str(value))
239        self._position = value
240
241    colourBarPosition = property(__get_position, __set_position, None,
242                    "colour bar position, 'horizontal' or 'vertical'")
243
244       
245
246def getColourBarImage(width=600, height=100,
247                      label=None, 
248                      cmap=cm.jet, 
249                      colourBarMin=None, 
250                      colourBarMax=None, 
251                      colourBarScale=COLOUR_SCHEME_SCALE.LINEAR,
252                      orientation='horizontal',
253                      intervals=None,
254                      intervalNames=None,
255                      numIntervals=None,
256                      colourBarStyle='continuous',
257                      dpi=100):
258   
259    figsize=(width / float(dpi), height / float(dpi))
260    fig = Figure(figsize=figsize, dpi=dpi, facecolor='w')
261   
262    #build a colour scheme
263    log.debug("colourBarScale = %s" % (colourBarScale,))
264    schemeBuilder = geoplot.colour_scheme.ColourSchemeBuilder(cmap=cmap, 
265                 colourBarMin=colourBarMin, 
266                 colourBarMax=colourBarMax, 
267                 colourBarScale=colourBarScale,
268                 numIntervals=numIntervals,
269                 intervals=intervals,
270                 intervalNames=intervalNames)
271   
272    colourScheme = schemeBuilder.buildScheme(colourBarStyle)   
273    log.debug("colourScheme.norm.__class__ = %s" % (colourScheme.norm.__class__,))
274    #for agg bakcend
275    #need about 40px at the bottom of the axes to draw the labels
276    #x = 40.0
277    #cbMin = 0.5
278   
279    #for cairo backend
280    if colourBarStyle == 'legend':
281        x = 40.0
282        cbMin = 0.6
283    else:
284        x = 70.0
285        cbMin = 0.6
286    cbBottom = x/height
287   
288    if cbBottom < 0.1:
289        cbBottom = 0.1
290   
291    if cbBottom > cbMin:
292        cbBottom = cbMin
293       
294    cbHeight = 0.9 - cbBottom
295    axes = fig.add_axes([0.05, cbBottom, 0.9, cbHeight], frameon=False)
296   
297    cb = ColourBar(colourBarLabel=label, 
298                   colourBarPosition=orientation,
299                   colourBarStyle=colourBarStyle)
300   
301    cb.draw(axes, colourScheme)   
302   
303    return geoplot.utils.figureToImage(fig)
Note: See TracBrowser for help on using the repository browser.