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

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

Fixed a couple of bugs.

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