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

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

Trying to get the different colour bars and plot options to work with the different legend types.

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