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

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

Fixed lots of small bugs and tidied up the layer drawer code a bit.

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_SCHEME_SCALE
17from geoplot.colour_scheme_intervals import ColourSchemeIntervalBuilder
18
19MAX_INTERVALS = 20
20
21log = logging.getLogger(__name__)
22
23class ColourScheme(object):
24    """
25    A colour scheme object used to keep track of the colourbar, normalization and
26    intervals asociated with a plot
27    """
28   
29    def __init__(self, cmap, norm, intervals, scale):
30        '''
31        @param cmap: the supplied matplotlib colourmap (may be listed if one
32          was given)
33        @type cmap:
34        @param norm: a Normalize instance for the colour bar boundaries given
35        @type norm:
36        @param intervals: the colour bar intervals supplied (or default intervals)
37        @type intervals:
38        '''
39       
40        self.cmap = cmap
41        self.norm = norm
42        self.intervals = intervals
43        self.scale = scale
44       
45    def getListedCmap(self):
46        if self.cmap.__class__ != matplotlib.colors.ListedColormap:
47           
48            if len(self.intervals.midpoints) == 1:
49                #only one interval value so pick the middle of the cmap
50                colors = [ self.cmap(0.5) ]
51                cmap = matplotlib.colors.ListedColormap(colors)
52            else:
53                colors = [self.cmap(self.norm(x)) for x in self.intervals.midpoints]
54                cmap = matplotlib.colors.ListedColormap(colors)
55        else:
56            cmap = self.cmap
57       
58        return cmap     
59   
60    def getBoundaryNorm(self):
61        return FixedBoundaryNorm(self.intervals.bounds, len(self.intervals.bounds) - 1 )
62
63
64class ColourSchemeBuilder(object):
65   
66    def __init__(self, cmap='jet', 
67                 colourBarMin=None, 
68                 colourBarMax=None,
69                 colourBarScale=COLOUR_SCHEME_SCALE.LINEAR,
70                 numIntervals=6,
71                 intervals=None):
72
73        self._numIntervals = 6
74        self.cmap = cmap
75        self.colourBarMin = colourBarMin
76        self.colourBarMax = colourBarMax
77        self.colourBarScale = colourBarScale
78        self.numIntervals = numIntervals
79        self.intervals = intervals
80        self._intervalBuilder = ColourSchemeIntervalBuilder()
81       
82    def buildScheme(self, grid=None):
83       
84        cbRange = self._getCbarRange(grid)
85       
86        if self.colourBarScale == COLOUR_SCHEME_SCALE.LOG:
87            cbRange = self._adjustCBRangeForLog(cbRange)
88       
89        cmap = self._getColourMap()
90       
91        intervals = self._intervalBuilder.buildCSInterval(cbRange, 
92                            self.colourBarScale, 
93                            self.intervals, 
94                            numIntervals=self.numIntervals)
95       
96        norm = self._buildNorm(cbRange)       
97       
98        return ColourScheme(cmap, norm, intervals, self.colourBarScale)
99   
100    def _buildNorm(self, cbRange):
101       
102        norm = None
103       
104        if self.colourBarScale == COLOUR_SCHEME_SCALE.LOG:
105           
106            if cbRange.minimum is not None and cbRange.minimum > 0 \
107               and cbRange.maximum is not None and cbRange.maximum > 0:
108                norm = matplotlib.colors.LogNorm(cbRange.minimum, cbRange.maximum)
109            else:
110                log.warning("Can't create a log colour scheme with min = %s and max = %s" 
111                             % (cbRange.minimum, cbRange.maximum))
112               
113        if norm == None:   
114            norm = matplotlib.colors.Normalize(cbRange.minimum, cbRange.maximum)
115       
116        return norm
117   
118    def _getColourMap(self):
119        """
120        Builds a colour map from the self.cmap variable and applies the set_under
121        and set_over variables if neccesary.
122       
123        If self.cmap is None then the matplotlib default will be used.
124        If self.cmap is a string then the named cmap from matplotlib will
125        be used.
126        """
127        if self.cmap == None:
128            cmap = matplotlib.cm.get_cmap()
129            cmap.set_bad("w")       
130       
131        elif type(self.cmap) in [str, unicode]:
132            cmap = matplotlib.cm.get_cmap(self.cmap)
133            cmap.set_bad("w")     
134        else:
135            cmap = self.cmap
136               
137        return cmap
138   
139    def _getCbarRange(self, grid=None):
140
141        if not grid is None:
142            cbMin = self.colourBarMin if self.colourBarMin != None else grid.getMinValue()
143            cbMax = self.colourBarMax if self.colourBarMax != None else grid.getMaxValue()
144        else:
145            cbMin = self.colourBarMin if self.colourBarMin != None else 0
146            cbMax = self.colourBarMax if self.colourBarMax != None else 1
147       
148        # check for masked values in vmin and vmax, can occur when data is completly masked
149        if cbMin.__class__ == numpy.ma.MaskedArray and cbMin.mask == True:
150            cbMin = None
151           
152        if cbMax.__class__ == numpy.ma.MaskedArray and cbMax.mask == True:
153            cbMax = None       
154       
155        if cbMin is not None and cbMax is not None and cbMin > cbMax:
156            log.warning("min(=%s) > max(=%s) reversing values" % (cbMin, cbMax))
157            cbMax, cbMin = cbMin, cbMax
158           
159        return Range(cbMin, cbMax)
160   
161    def _adjustCBRangeForLog(self, cbRange):
162        adjustedRange = Range(cbRange.minimum, cbRange.maximum)
163       
164        if adjustedRange.maximum <= 0.0:
165            adjustedRange.maximum = 1.0
166       
167        if adjustedRange.minimum <= 0.0:
168                           
169            if adjustedRange.maximum >= 10.0:
170                adjustedRange.minimum = 1.0
171            else:
172                adjustedRange.minimum = 10 ** (numpy.log10(adjustedRange.maximum) - 2)
173       
174        return adjustedRange
175   
176    def __set_numIntervals(self, value):
177        if value <= 0:
178            log.warning("Can't set numIntervals to %s in ColourSchemeBuilder" % (value,))
179        else:
180            self._numIntervals = value
181
182    def __get_numIntervals(self):
183        return self._numIntervals
184       
185    numIntervals = property(__get_numIntervals, __set_numIntervals, None, None)
Note: See TracBrowser for help on using the repository browser.