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

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/qesdi/geoplot/trunk/lib/geoplot/colour_scheme.py@5946
Revision 5946, 11.0 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'''
2Created on 15 Oct 2009
3
4The ColourScheme object is responsible for providing access to the ColourMap and
5normalisation objects.
6
7@author: pnorton
8'''
9
10import logging
11import numpy
12import matplotlib
13
14from geoplot.range import Range
15from geoplot.fixed_boundary_norm import FixedBoundaryNorm
16
17import utils
18
19MAX_INTERVALS = 20
20
21log = logging.getLogger(__name__)
22
23
24class COLOUR_SCHEME_SCALE:
25    LINEAR='linear'
26    LOG='log'
27
28
29class ColourScheme(object):
30   
31    def __init__(self, colourMap, norm, scale):
32        self.colourMap = colourMap
33        self.norm = norm
34        self.scale = scale
35
36class ContinuousColourScheme(ColourScheme):
37    pass
38
39class IntervalColourScheme(ColourScheme):
40    def __init__(self, colourMap, norm, scale, labels, labelLocations):
41        self.labels = labels
42        self.labelLocations = labelLocations
43        ColourScheme.__init__(self, colourMap, norm, scale)
44       
45
46class ColourSchemeBuilder(object):
47   
48    def __init__(self, cmap=None, 
49                 colourBarMin=None, 
50                 colourBarMax=None,
51                 colourBarScale=COLOUR_SCHEME_SCALE.LINEAR, 
52                 hideOutsideBounds=False,
53                 intervals=None,
54                 intervalNames=None):
55
56        self.cmap = cmap
57        self.colourBarMin = colourBarMin
58        self.colourBarMax = colourBarMax
59        self.colourBarScale = colourBarScale
60        self.hideOutsideBounds = hideOutsideBounds
61        self.intervals = intervals
62        self.intervalNames = intervalNames
63       
64    def buildScheme(self, grid=None):
65       
66        cbRange = self._getCbarRange()
67        initialCmap = self._getColourMap()
68       
69        if self.intervals != None:
70            builder = _IntervalSchemeBuidler(cbRange, self.colourBarScale, initialCmap, self.intervals, self.intervalNames)
71        else:
72            builder = _ConinuousSchemeBuilder(cbRange, self.colourBarScale, initialCmap)
73
74        scheme = builder.buildScheme(grid)
75       
76        return scheme
77
78    def _getColourMap(self):
79        """
80        Builds a colour map from the self.cmap variable and applies the set_under
81        and set_over variables if neccesary.
82       
83        If self.cmap is None then the matplotlib default will be used.
84        If self.cmap is a string then the named cmap from matplotlib will
85        be used.
86        """
87        if self.cmap == None:
88            cmap = matplotlib.cm.get_cmap()
89            cmap.set_bad("w")       
90       
91        elif type(self.cmap) in [str, unicode]:
92            cmap = matplotlib.cm.get_cmap(self.cmap)
93            cmap.set_bad("w")     
94        else:
95            cmap = self.cmap
96               
97        if self.hideOutsideBounds:
98            log.debug("self.hideOutsideBounds = %s" % (self.hideOutsideBounds,))
99            cmap.set_under('0.25', alpha=0.0)
100            cmap.set_over('0.75', alpha=0.0)
101           
102        return cmap
103   
104    def _getCbarRange(self):
105        cbMin = self.colourBarMin
106        cbMax = self.colourBarMax
107       
108        if cbMin is not None and cbMax is not None and cbMin > cbMax:
109            log.warning("min(=%s) > max(=%s) reversing values" % (cbMin, cbMax))
110            cbMax, cbMin = cbMin, cbMax
111           
112        return Range(cbMin, cbMax)
113
114
115class _ConinuousSchemeBuilder():
116    """
117    This class knows how to buld a continuous colour scheme.
118    """
119   
120    def __init__(self, cbRange, cbScale, initialCmap):
121        self.cbRange = cbRange
122        self.initialCmap = initialCmap
123        self.cbScale = cbScale
124   
125    def buildScheme(self, grid=None):
126        cmap = self._buildColourMap()
127        norm = self._buildNorm(grid)
128        return ContinuousColourScheme(cmap, norm, self.cbScale)
129   
130    def _buildColourMap(self):
131        return self.initialCmap
132   
133    def _buildNorm(self, grid):
134       
135        norm = None
136       
137        cbMin, cbMax = self.cbRange.minimum, self.cbRange.maximum
138       
139        if not grid is None:
140            vmin = cbMin if cbMin != None else grid.getMinValue()
141            vmax = cbMax if cbMax != None else grid.getMaxValue()
142        else:
143            vmin = cbMin
144            vmax = cbMax
145       
146        if self.cbScale == 'log':
147           
148            if vmin is not None and vmin > 0 and vmax is not None and vmax > 0:
149                norm = matplotlib.colors.LogNorm(vmin, vmax)
150            else:
151                log.warning("Can't create a log colour scheme with min = %s and max = %s" 
152                             % (vmin, vmax))
153               
154       
155        if norm == None:   
156            norm = matplotlib.colors.Normalize(vmin, vmax)
157       
158       
159        # check for masked values in vmin and vmax, can occur when data is completly masked
160        if norm.vmin.__class__ == numpy.ma.MaskedArray and norm.vmin.mask == True:
161            norm.vmin = None
162           
163        if norm.vmax.__class__ == numpy.ma.MaskedArray and norm.vmax.mask == True:
164            norm.vmax = None
165               
166        return norm
167
168class _IntervalSchemeBuidler():
169   
170    def __init__(self, cbRange, cbScale, initialCmap, intervalBoundsString, intervalLabelsString=None):
171        self.cbRange = cbRange
172        self.cbScale = cbScale
173        self.initialCmap = initialCmap
174        self.intervalBoundsString = intervalBoundsString
175        self.intervalLabelsString = intervalLabelsString
176       
177    def buildScheme(self, grid=None):
178        """
179        Builds an interval colour scheme using either the unique values form
180        the grid or the bounds in the intervalBoundsString
181        """
182
183        bounds = numpy.array([float(x) for x in self._splitString(self.intervalBoundsString)])
184       
185        bounds = self._filterBoundsUsingCBMinMax(bounds)
186       
187        csInterval = self._buildCSIntervalFromBounds(bounds)
188       
189        labels = self._getLabels(csInterval)
190
191        cmap = self._buildColourMap(csInterval)
192        norm = self._getNorm(csInterval)
193       
194        return IntervalColourScheme(cmap, norm, self.cbScale, labels, csInterval.midpoints)
195   
196    def _buildColourMap(self, csInterval):
197        """
198        Returns a (possibly) modified version of the initialCmap that corresponds
199        to the values given by the _ColourSchemeInterval.
200        """       
201        if self.initialCmap.__class__ != matplotlib.colors.ListedColormap:
202           
203            if len(csInterval.midpoints) == 1:
204                #only one interval value so pick the middle of the cmap
205                colors = [ self.initialCmap(0.5) ]
206                cmap = matplotlib.colors.ListedColormap(colors)
207            else:
208                n = matplotlib.colors.Normalize(csInterval.midpoints.min(), 
209                                                csInterval.midpoints.max())
210   
211                colors = [self.initialCmap(n(x)) for x in csInterval.midpoints]
212                cmap = matplotlib.colors.ListedColormap(colors)
213        else:
214            cmap = self.initialCmap
215       
216        return cmap     
217                       
218    def _getNorm(self, csInterval):
219        """
220        Returns a Boundary Normalisation object for the given _ColourSchemeInterval.
221        """
222        norm = FixedBoundaryNorm(csInterval.bounds, len(csInterval.bounds) - 1 )
223        return norm
224   
225    def _getLabels(self, csInterval):
226        if self.intervalLabelsString == None:
227            labels = csInterval.defalutLabels
228        else:
229            labels = self._splitString(self.intervalLabelsString)
230            if len(labels) != len(csInterval.midpoints):
231                log.warning("Number of labels found (%s) is != to the number of midpoints (%s), using default labels instead" \
232                            % (len(labels), len(csInterval.midpoints)))
233                labels = csInterval.defalutLabels
234       
235        return labels
236               
237   
238    def _splitString(self, csvString):
239        """
240        Splits a csv string and removes any empty values
241        """
242        stringList = csvString.split(',')
243        stringList = filter(lambda x: x.strip() != '', stringList)       
244        return stringList
245   
246    def _buildCSIntervalFromMidpoints(self, midpoints):
247        """
248        Builds a _ColourSchemeInterval object from a list of midpoint values
249        """
250        bounds = utils.getBounds(midpoints)
251        defaultLabels = [str(utils.round_to_n(x, 2)) for x in midpoints]
252        return _ColourSchemeInterval(midpoints, bounds, defaultLabels)
253   
254    def _buildCSIntervalFromBounds(self, bounds):
255        """
256        Builds a _ColourSchemeInterval object from a list of bounds.
257        """
258        midpoints = utils.getMidpoints(bounds)
259        boundStrings = [str(utils.round_to_n(x,2)) for x in bounds]
260        #log.debug("boundStrings = %s" % (boundStrings,))
261        defaultLabels  = ["%s - %s" % (boundStrings[index], boundStrings[index + 1]) for index in range(len(boundStrings)-1)]
262        return _ColourSchemeInterval(midpoints, bounds, defaultLabels)
263   
264    def _getBoundsFromValues(self, values):
265        """
266        This function uses the minimum and maximum of the values to
267        create MAX_INTERVALS bounds which are returned as a numpy array.
268        """
269        values = self._filterValuesUsingCBMinMax(values)
270       
271        boundsMin, boundsMax = values.min(), values.max()
272       
273        if self.cbMin != None and self.cbMin > boundsMin:
274            boundsMin = self.cbMin
275           
276        if self.cbMax != None and self.cbMax < boundsMax:
277            boundsMax = self.cbMax
278       
279        interval = (boundsMax - boundsMin)/float(MAX_INTERVALS)
280        bounds = numpy.arange(boundsMin, boundsMax, interval)
281       
282        #include the upper bound
283        bounds =numpy.append(bounds, boundsMax)
284       
285        return bounds
286   
287    def _filterValuesUsingCBMinMax(self, values):
288        #get rid of any values above cbMax
289        if self.cbRange.maximum != None:
290            values = numpy.array(filter(lambda x: x <= self.cbRange.maximum, values))
291
292        # get rid of any bounds less than the minimum
293        if self.cbRange.minimum != None:
294            values = numpy.array(filter(lambda x: self.cbRange.minimum <= x , values))
295       
296        return values
297   
298    def _filterBoundsUsingCBMinMax(self, bounds):
299       
300        newBounds = bounds
301       
302        if self.cbRange.minimum != None:
303           
304            #start with the maximum of the first bound
305            for i in range(1, len(bounds)):
306                if bounds[i] <= self.cbRange.minimum:
307                    newBounds = newBounds[1:]
308                else:
309                    break
310                     
311        if self.cbRange.maximum != None:
312           
313            #start with the lower of the top bound
314            for i in range(len(bounds) - 2, 0, -1):
315                if bounds[i] > self.cbRange.maximum:
316                    newBounds = newBounds[:-1]
317                else:
318                    break
319
320        return newBounds
321       
322           
323
324class _ColourSchemeInterval(object):
325    """
326    Represents interval information required to build an interval colour scheme
327    """
328   
329    def __init__(self, midpoints, bounds, defalutLabels):
330        self.midpoints = midpoints
331        self.bounds = bounds
332        self.defalutLabels = defalutLabels
333       
Note: See TracBrowser for help on using the repository browser.