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

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

Imroved the colour bar code so that a legend colour bar can be used without specifying any intervals.

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
16from geoplot.colour_bar import COLOUR_BAR_STYLES, COLOUR_SCHEME_SCALE
17from geoplot.colour_scheme_intervals import ColourSchemeIntervalBuilder
18
19MAX_INTERVALS = 20
20
21log = logging.getLogger(__name__)
22
23class ColourScheme(object):
24   
25    def __init__(self, colourMap, norm, scale):
26        self.colourMap = colourMap
27        self.norm = norm
28        self.scale = scale
29
30class ContinuousColourScheme(ColourScheme):
31    pass
32
33class IntervalColourScheme(ColourScheme):
34       
35    def __init__(self, colourMap, norm, scale, labels, labelLocations):
36        self.labels = labels
37        self.labelLocations = labelLocations
38        ColourScheme.__init__(self, colourMap, norm, scale)
39       
40
41class ColourSchemeBuilder(object):
42   
43    def __init__(self, cmap=None, 
44                 colourBarMin=None, 
45                 colourBarMax=None,
46                 colourBarScale=COLOUR_SCHEME_SCALE.LINEAR, 
47                 hideOutsideBounds=False,
48                 numIntervals=6,
49                 intervals=None,
50                 intervalNames=None):
51
52        self._numIntervals = None
53       
54        self.cmap = cmap
55        self.colourBarMin = colourBarMin
56        self.colourBarMax = colourBarMax
57        self.colourBarScale = colourBarScale
58        self.hideOutsideBounds = hideOutsideBounds
59        self.numIntervals = numIntervals
60        self.intervals = intervals
61        self.intervalNames = intervalNames
62        self._intervalBuilder = ColourSchemeIntervalBuilder()
63       
64    def buildScheme(self, colourBarStyle, grid=None):
65       
66        cbRange = self._getCbarRange(grid)
67        initialCmap = self._getColourMap()
68       
69        assert colourBarStyle in COLOUR_BAR_STYLES.all()
70       
71        if colourBarStyle == COLOUR_BAR_STYLES.LEGEND:
72            builder = self._getIntervalBuilder(cbRange, initialCmap)
73        else:
74            builder = self._getContinuousBuilder(cbRange, initialCmap)
75
76        scheme = builder.buildScheme()
77       
78        if self.hideOutsideBounds:
79            scheme.norm.clip = False
80       
81        return scheme
82   
83    def _getIntervalBuilder(self,cbRange, initialCmap):
84       
85        interval = self._intervalBuilder.buildCSInterval(cbRange, self.intervals, self.intervalNames, numIntervals=self.numIntervals)
86        log.debug("interval = %s" % (interval,))
87        builder = _IntervalSchemeBuidler(cbRange, self.colourBarScale, initialCmap, interval)
88       
89        return builder
90   
91    def _getContinuousBuilder(self, cbRange, initialCmap):
92        builder = _ConinuousSchemeBuilder(cbRange, self.colourBarScale, initialCmap)
93        return builder
94
95    def _getColourMap(self):
96        """
97        Builds a colour map from the self.cmap variable and applies the set_under
98        and set_over variables if neccesary.
99       
100        If self.cmap is None then the matplotlib default will be used.
101        If self.cmap is a string then the named cmap from matplotlib will
102        be used.
103        """
104        if self.cmap == None:
105            cmap = matplotlib.cm.get_cmap()
106            cmap.set_bad("w")       
107       
108        elif type(self.cmap) in [str, unicode]:
109            cmap = matplotlib.cm.get_cmap(self.cmap)
110            cmap.set_bad("w")     
111        else:
112            cmap = self.cmap
113               
114        if self.hideOutsideBounds:
115            log.debug("self.hideOutsideBounds = %s" % (self.hideOutsideBounds,))
116            cmap.set_under('0.25', alpha=0.0)
117            cmap.set_over('0.75', alpha=0.0)
118           
119        return cmap
120   
121    def _getCbarRange(self, grid=None):
122
123        if not grid is None:
124            cbMin = self.colourBarMin if self.colourBarMin != None else grid.getMinValue()
125            cbMax = self.colourBarMax if self.colourBarMax != None else grid.getMaxValue()
126        else:
127            cbMin = self.colourBarMin if self.colourBarMin != None else 0
128            cbMax = self.colourBarMax if self.colourBarMax != None else 1
129       
130        # check for masked values in vmin and vmax, can occur when data is completly masked
131        if cbMin.__class__ == numpy.ma.MaskedArray and cbMin.mask == True:
132            cbMin = None
133           
134        if cbMax.__class__ == numpy.ma.MaskedArray and cbMax.mask == True:
135            cbMax = None       
136       
137        if cbMin is not None and cbMax is not None and cbMin > cbMax:
138            log.warning("min(=%s) > max(=%s) reversing values" % (cbMin, cbMax))
139            cbMax, cbMin = cbMin, cbMax
140           
141        return Range(cbMin, cbMax)
142   
143   
144   
145    def __set_numIntervals(self, value):
146        if value <= 0:
147            log.warning("Can't set numIntervals to %s in ColourSchemeBuilder" % (value,))
148        else:
149            self._numIntervals = value
150
151    def __get_numIntervals(self):
152        return self._numIntervals
153       
154    numIntervals = property(__get_numIntervals, __set_numIntervals, None, None)
155
156class _ConinuousSchemeBuilder():
157    """
158    This class knows how to buld a continuous colour scheme.
159    """
160   
161    def __init__(self, cbRange, cbScale, initialCmap):
162        self.cbRange = cbRange
163        self.initialCmap = initialCmap
164        self.cbScale = cbScale
165   
166    def buildScheme(self):
167        cmap = self._buildColourMap()
168        norm = self._buildNorm()
169        return ContinuousColourScheme(cmap, norm, self.cbScale)
170   
171    def _buildColourMap(self):
172        return self.initialCmap
173   
174    def _buildNorm(self):
175       
176        norm = None
177       
178        if self.cbScale == 'log':
179           
180            if self.cbRange.minimum is not None and self.cbRange.minimum > 0 \
181               and self.cbRange.maximum is not None and self.cbRange.maximum > 0:
182                norm = matplotlib.colors.LogNorm(self.cbRange.minimum, self.cbRange.maximum)
183            else:
184                log.warning("Can't create a log colour scheme with min = %s and max = %s" 
185                             % (self.cbRange.minimum, self.cbRange.maximum))
186               
187       
188        if norm == None:   
189            norm = matplotlib.colors.Normalize(self.cbRange.minimum, self.cbRange.maximum)
190       
191        return norm
192
193class _IntervalSchemeBuidler():
194   
195    def __init__(self, cbRange, cbScale, initialCmap, csInterval):
196        self.cbRange = cbRange
197        self.cbScale = cbScale
198        self.initialCmap = initialCmap
199        self.csInterval = csInterval
200       
201    def buildScheme(self):
202        """
203        Builds an interval colour scheme using either the unique values form
204        the grid or the bounds in the intervalBoundsString
205        """
206       
207        labels = self.csInterval.labels
208        cmap = self._buildColourMap()
209        norm = self._getNorm()
210       
211        return IntervalColourScheme(cmap, norm, self.cbScale, labels, self.csInterval.midpoints)
212   
213    def _buildColourMap(self):
214        """
215        Returns a (possibly) modified version of the initialCmap that corresponds
216        to the values given by the _ColourSchemeInterval.
217        """       
218        if self.initialCmap.__class__ != matplotlib.colors.ListedColormap:
219           
220            if len(self.csInterval.midpoints) == 1:
221                #only one interval value so pick the middle of the cmap
222                colors = [ self.initialCmap(0.5) ]
223                cmap = matplotlib.colors.ListedColormap(colors)
224            else:
225                n = matplotlib.colors.Normalize(self.csInterval.midpoints.min(), 
226                                                self.csInterval.midpoints.max())
227   
228                colors = [self.initialCmap(n(x)) for x in self.csInterval.midpoints]
229                cmap = matplotlib.colors.ListedColormap(colors)
230        else:
231            cmap = self.initialCmap
232       
233        return cmap     
234                       
235    def _getNorm(self):
236        """
237        Returns a Boundary Normalisation object for the given _ColourSchemeInterval.
238        """
239        norm = FixedBoundaryNorm(self.csInterval.bounds, len(self.csInterval.bounds) - 1 )
240        return norm
241   
Note: See TracBrowser for help on using the repository browser.