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

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

Fixed a problem with the subsetting in the grid_builder_lat_lon_subset.py file.

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