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

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

Added the ability to set the scale on the colour scheme to a log scale.

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