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

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

Moved QESDI tree from DCIP repository
 http://proj.badc.rl.ac.uk/svn/dcip/qesdi@3900.

Line 
1
2import math
3import logging
4
5import matplotlib.colors
6import matplotlib.ticker
7
8import geoplot.config as geoplot_config
9config = geoplot_config.getConfig()
10
11log = logging.getLogger(__name__)
12log.setLevel(logging.DEBUG)
13
14FONTS_SECTION = 'Fonts'
15
16
17MAX_CBAR_TICKS = 10
18ADJUSTED_TICK_FORMAT = "%1.2f"
19
20class ColourBar(object):
21
22    def __init__(self, label, position, valueRange):
23
24        self._position = None
25        self._range = None
26
27        self.label = label
28        self.position = position
29        self.range = valueRange
30        self.labelFont = config[FONTS_SECTION]['ColourBarLabel']
31
32    def draw(self, figure, colourBarAxis, sm, fontSize):
33        """
34        Adds the colour bar to the (and optionally a label) to the figure.
35
36        @param sm: the scalar mappable generated by applying the grid to the axis
37        @type sm: an instance of matplotlib.cm.ScalarMappable
38        @param units: the units of the values on the mesh, these will be used as a label
39            if the colourBarLabel property is not set.
40        @type units: string
41        """
42
43#        log.debug("Drawing colour bar")
44
45        self._setLimits(sm)
46       
47        cb = figure.colorbar(sm, cax = colourBarAxis, orientation=self.position)
48
49        log.debug("cb.cmap = %s" % (cb.cmap,))
50        if cb.cmap.__class__ == matplotlib.colors.ListedColormap:
51            ColourBar._repositionColourBarTicks(cb)
52           
53        if self.label != None:
54            labelDictionary = self.labelFont.getDict(fontSize)
55            cb.set_label(self.label, fontdict=labelDictionary)
56
57        log.debug("tick position = %s" % (colourBarAxis.yaxis.get_ticks_position(),))
58
59    def _setLimits(self, sm):
60        """
61        Changes the limits of the given scalar mesh to the limits specified in
62        the constructor.
63
64        First checks if the min/max value is not None before using it to set the
65        value on the scalar mesh. If the min/max value is None then the mesh
66        limit is not changed.
67
68        @param sm: the scalar mappable generated by applying the grid to the axis
69        @type sm: an instance of matplotlib.cm.ScalarMappable
70        """
71
72#        log.debug("setting limits, colourbar range is:" + str(self.range))
73
74        log.debug("Setting limits")
75       
76        (dataMin, dataMax) = sm.get_clim()
77        (minVal, maxVal) = self.range
78
79        if minVal != None and maxVal != None:
80            sm.set_clim(minVal, maxVal)
81
82        elif minVal != None:
83            sm.set_clim(minVal, dataMax)
84
85        elif maxVal != None:
86            sm.set_clim(dataMin, maxVal)
87
88    @staticmethod
89    def _repositionColourBarTicks(cb):
90        """
91        reposition the ticks of a ListedColormap so that they appear at the
92        """
93       
94        log.debug("Repositioning colour bar ticks")
95        span = cb.vmax - cb.vmin
96
97        # Define flag for whether or not zero should be added
98        useZero = False
99        if cb.vmin < 0 and cb.vmax > 0:
100            useZero = True
101
102        numColours = len(cb.cmap.colors)
103        interval = float(span) / float(numColours)
104
105        showEvery = 1
106        while float(numColours)/float(showEvery) > float(MAX_CBAR_TICKS):
107            showEvery += 1
108
109        newLocs = []
110        for i in range(0, numColours + 1, showEvery):
111            newLocs.append(cb.vmin + i * interval)
112
113        # If need to add a zero then do so
114        if useZero == True and 0 not in newLocs:
115            newLocsWithZero = []
116         
117            zeroInserted = False 
118            for newLoc in newLocs:
119                if newLoc > 0 and zeroInserted == False:
120                    newLocsWithZero.append(0)
121                    zeroInserted = True                 
122                newLocsWithZero.append(newLoc)
123
124            newLocs = newLocsWithZero
125
126        #change the locator and the formatter (as the locations now have a high number of dp)
127        cb.locator = matplotlib.ticker.FixedLocator(newLocs)
128
129        # Decide the float formatting of the tick labels based on the span
130        if span < 100:
131            # In Excel this worked: =IF(E3<100,(LOG(E3)*-1)+2,0)
132            tickFormatDecPoints = int((math.log(span, 10) * -1) + 2)
133            tickFormatString = "%." + str(tickFormatDecPoints) + "f"
134        else: 
135            tickFormatString = "%d"
136           
137        cb.formatter = matplotlib.ticker.FormatStrFormatter(tickFormatString)  ###ADJUSTED_TICK_FORMAT)
138
139        # The next line removes axis artists as draw_all() adds new ones
140        cb.ax.artists = []
141        cb.draw_all()    # cause the colourbar to be redrawn, otherwise not changes will hapen
142
143        # Hard code line width of colour bar outline
144        cb.outline.set_linewidth(0.5)
145
146        # this can sometimes cause the tick positions to become unknown so
147        # set it back to default.
148        if cb.ax.yaxis.get_ticks_position() == 'unknown':
149            log.debug("Resetting yaxis ticks position to default")
150            cb.ax.yaxis.set_ticks_position('default')
151       
152        if cb.ax.xaxis.get_ticks_position() == 'unknown':
153            log.debug("Resetting xaxis ticks position to default")
154            cb.ax.xaxis.set_ticks_position('default')
155           
156    #properties code
157    def __get_range(self):
158        if self._range == None:
159            return (None, None)
160        else:
161            return self._range
162
163    def __set_range(self, value):
164
165        if value == None:
166            self._range = None
167        else:
168            #check the new value
169            if not (type(value) == list or type(value) == tuple):
170                raise ValueError("Cant set range to a value of " + str(type(value)))
171
172            if type(value) == list:
173                value == tuple(value)
174
175            if len(value) != 2:
176                raise ValueError("Cant set range to a value of length "
177                                + str(len(value)) + ", length should be 2.")
178
179            (minVal, maxVal) = value
180
181            if minVal != None and maxVal != None:
182                #convert the values to floats   
183                minVal = float(minVal)
184                maxVal = float(maxVal)
185               
186                if minVal > maxVal:
187                    raise ValueError("Cant set range to have minimum > maximum, " + str(value))
188                elif minVal == maxVal:
189                    raise ValueError("Cant set range to have minimum == maximum, " + str(value))
190
191            self._range = (minVal, maxVal)
192
193    range = property(__get_range, __set_range, None, "range of colour bar values (min, max)")
194
195    def __get_position(self): return self._position
196
197    def __set_position(self, value):
198        if value not in ['horizontal', 'vertical']:
199            raise ValueError("ColourBar position value must be 'horizontal'" + \
200                             " or 'vertical, value recieved :" + str(value))
201        self._position = value
202
203    position = property(__get_position, __set_position, None,
204                    "colour bar position, 'horizontal' or 'vertical'")
205
206if __name__ == '__main__':
207
208
209
210    from cdms_utils.cdms_compat import N
211    import geoplot.utils
212    import geoplot.log_util
213    from matplotlib.backends.backend_agg import FigureCanvasAgg
214    from matplotlib.figure import Figure
215    import matplotlib
216
217    geoplot.log_util.setupGeoplotConsoleHandler(log)
218
219    cb = ColourBar("label", "vertical", None)
220
221    filename = "out.png"
222
223    xlim = 5
224    ylim = 7
225
226    x = range(0,xlim+1)
227    y = range(0,ylim+1)
228    X,Y = N.meshgrid(x,y)
229
230    X[1][1] = 2
231    X[1][2] = 3
232
233    Z=[[0]*xlim for i in range(ylim)]
234    for i in range(0 , xlim*ylim):
235        vX= i%xlim
236        vY= (i - i%xlim) / xlim
237        Z[vY][vX] = i
238
239    Z[0][1] = 0.0009
240    Z[6][1] = Z[6][2] = Z[6][3] = 30;
241    Z[6][4] = 38.545
242
243    c255 = [( 53, 192, 157), ( 25, 162, 151), ( 28, 126, 114), ( 18,  86,  98),
244            ( 30,  45,  78), ( 15,  68, 121), ( 61,  73, 182), ( 88,  96, 220),
245            (121, 126, 255), (153, 153, 255), (156, 182, 255), (156, 209, 255),
246            (156, 237, 255), (166, 255, 246), (196, 255, 217), (224, 255, 189),
247            (249, 255, 161), (255, 236, 153), (255, 207, 153), (255, 180, 153),
248            (255, 153, 153), (255, 121, 131), (240,  88, 104), (215,  78,  91),
249            (192,  66,  83), (167,  13,  61), (131,  13,  61), ( 98,   0,  48),
250            (121,  45, 116), (151,  83, 169)]
251
252    cmap = geoplot.utils.generateDiscreteCmap(c255, "temp")
253    cmap = matplotlib.cm.jet
254   
255    fig = Figure()
256#    axes = fig.add_axes([0.1, 0.2, 0.8, 0.7])   
257#    cbAxes = fig.add_axes([0.1,0.1,0.8,0.05])
258   
259    axes = fig.add_axes([0.1, 0.1, 0.6, 0.8])
260    cbAxes = fig.add_axes([0.8, 0.1, 0.1, 0.8])
261   
262    sm = axes.pcolor(X,Y,Z,shading='faceted', cmap=cmap)
263    cb.draw(fig, cbAxes, sm, "medium")
264    canvas = FigureCanvasAgg(fig)
265    canvas.print_figure(filename)
266    print "Written:", filename
Note: See TracBrowser for help on using the repository browser.