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

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

Improved the colour bar code so it now accepts the data array instead of just a minimum + maximum. It also now defaults to creating a Normalize instance with vmin=None and vmax=None instead of 0 and 1.

Also modified the layer_drawer objects to use the ColourBar? class instead of implementing it themselves.

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