source: qesdi/geoplot/trunk/lib/geoplot/colour_bar.py @ 5735

import math
import logging
import numpy

from matplotlib.colors import Normalize, ListedColormap

import matplotlib.ticker
from matplotlib.figure import Figure
from matplotlib.colorbar import ColorbarBase
import matplotlib.cm as cm
import operator


from geoplot.fixed_boundary_norm import FixedBoundaryNorm
import geoplot.config as geoplot_config
import geoplot.utils


config = geoplot_config.getConfig()

log = logging.getLogger(__name__)

FONTS_SECTION = 'Fonts'
MAX_CBAR_TICKS = 10
ADJUSTED_TICK_FORMAT = "%1.2f"

class ColourBar(object):

    def __init__(self, colourBarLabel="",colourBarPosition='horizontal', 
                       cmap=None, colourBarMin=None, colourBarMax=None,
                       intervalColourbar=False, intervalNames=None,
                       hideOutsideBounds=False):

        self._position = None
        self._range = None
        self._cmap = None

        self.cmap = cmap
        self.colourBarLabel = colourBarLabel
        self.colourBarPosition = colourBarPosition
        self.colourBarMin = colourBarMin
        self.colourBarMax = colourBarMax
        self.hideOutsideBounds = hideOutsideBounds
        self.intervalColourbar = intervalColourbar
        self.intervalNames = intervalNames
        
        self.labelFont = config[FONTS_SECTION]['ColourBarLabel']

    def draw(self, colourBarAxis, fontSize='medium', grid=None):
        """
        Adds the colour bar to the (and optionally a label) to the figure.

        @param sm: the scalar mappable generated by applying the grid to the axis
        @type sm: an instance of matplotlib.cm.ScalarMappable
        @param units: the units of the values on the mesh, these will be used as a label
            if the colourBarLabel property is not set.
        @type units: string
        """
        log.debug("drawing colour bar")
        kwargs = {}
        
        
        
        kwargs['cmap'] = self.getColourMap(grid)
        kwargs['norm'] = self.getNormalize(grid)
        kwargs['orientation'] = self.colourBarPosition
        
        if self.intervalColourbar:
            
            cbMin, cbMax = self._getCbarMinMax()
            kwargs['ticks'] = grid.getUniqueValues(cbMin, cbMax)
            
            if self.intervalNames != None:
                
                if len(self.intervalNames) == len(kwargs['ticks']):
                    kwargs['format'] = matplotlib.ticker.FixedFormatter(self.intervalNames)
                else:
                    log.warning("Incorrect number of interval names, expected %s but recieved %s" % (len(kwargs['ticks']), len(self.intervalNames)))
            
            kwargs['spacing'] = 'proportional'
            
            log.debug("kwargs['cmap'].N = %s" % (kwargs['cmap'].N,))
        
        cb = ColorbarBase(colourBarAxis, **kwargs)

        if cb.cmap.__class__ == ListedColormap \
           and not self.intervalColourbar:
            ColourBar._repositionColourBarTicks(cb)
            
        if self.colourBarLabel != None:
            labelDictionary = self.labelFont.getDict(fontSize)
            cb.set_label(self.colourBarLabel, fontdict=labelDictionary)
        
        log.debug("finished drawing colour bar")

    def getColourMap(self, grid=None):
        
        if self.cmap == None:
            cmap = matplotlib.cm.get_cmap()
            cmap.set_bad("w")        
        
        elif type(self.cmap) in [str, unicode]:
            cmap = cm.get_cmap(self.cmap)
            cmap.set_bad("w")     
        else:
            cmap = self.cmap
    
        if self.intervalColourbar and  cmap.__class__ != ListedColormap:
            
            cbMin, cbMax = self._getCbarMinMax()
            
            oldCmap = cmap
            uniqueVals = grid.getUniqueValues(cbMin, cbMax)
            n = Normalize(uniqueVals[0], uniqueVals[-1])
            colours = [oldCmap(n(x)) for x in uniqueVals]
            cmap = ListedColormap(colours)
                
        if self.hideOutsideBounds:
            log.debug("self.hideOutsideBounds = %s" % (self.hideOutsideBounds,))
            cmap.set_under('0.25', alpha=0.0)
            cmap.set_over('0.75', alpha=0.0)
            
        return cmap

    @staticmethod
    def _repositionColourBarTicks(cb):
        """
        reposition the ticks of a ListedColormap so that they appear at the
        """
        
        log.debug("Repositioning colour bar ticks")
        span = cb.vmax - cb.vmin

        # Define flag for whether or not zero should be added
        useZero = False
        if cb.vmin < 0 and cb.vmax > 0:
            useZero = True

        numColours = len(cb.cmap.colors)
        interval = float(span) / float(numColours)

        showEvery = 1
        while float(numColours)/float(showEvery) > float(MAX_CBAR_TICKS):
            showEvery += 1

        newLocs = []
        for i in range(0, numColours + 1, showEvery):
            newLocs.append(cb.vmin + i * interval)

        # If need to add a zero then do so
        if useZero == True and 0 not in newLocs:
            newLocsWithZero = []
          
            zeroInserted = False 
            for newLoc in newLocs:
                if newLoc > 0 and zeroInserted == False:
                    newLocsWithZero.append(0)
                    zeroInserted = True                  
                newLocsWithZero.append(newLoc)

            newLocs = newLocsWithZero

        #change the locator and the formatter (as the locations now have a high number of dp)
        cb.locator = matplotlib.ticker.FixedLocator(newLocs)

        # Decide the float formatting of the tick labels based on the span
        if span < 100:
            # In Excel this worked: =IF(E3<100,(LOG(E3)*-1)+2,0)
            tickFormatDecPoints = int((math.log(span, 10) * -1) + 2)
            tickFormatString = "%." + str(tickFormatDecPoints) + "f"
        else: 
            tickFormatString = "%d"
            
        cb.formatter = matplotlib.ticker.FormatStrFormatter(tickFormatString)  ###ADJUSTED_TICK_FORMAT)

        # The next line removes axis artists as draw_all() adds new ones
        cb.ax.artists = []
        cb.draw_all()    # cause the colourbar to be redrawn, otherwise not changes will hapen

        # Hard code line width of colour bar outline
        cb.outline.set_linewidth(0.5)

        # this can sometimes cause the tick positions to become unknown so
        # set it back to default.
        if cb.ax.yaxis.get_ticks_position() == 'unknown':
            log.debug("Resetting yaxis ticks position to default")
            cb.ax.yaxis.set_ticks_position('default')
        
        if cb.ax.xaxis.get_ticks_position() == 'unknown':
            log.debug("Resetting xaxis ticks position to default")
            cb.ax.xaxis.set_ticks_position('default')
    
    def getNormalize(self, grid=None):
        
        cbMin, cbMax = self._getCbarMinMax()
        
        if self.intervalColourbar:
            bounds = grid.getUniqueValueBounds(cbMin, cbMax)
            n = len(grid.getUniqueValues(cbMin, cbMax))
            norm = FixedBoundaryNorm(bounds, n)
        else:
            norm = Normalize(cbMin,  cbMax)
            # this should work event if data is none (as N.ma.maximum(None) = None)
            data = None
            if grid != None:
                data = grid.values
            
            norm.autoscale_None(data)
            
            
            # check for masked values in vmin and vmax, can occur when data is completly masked
            if norm.vmin.__class__ == numpy.ma.MaskedArray and norm.vmin.mask == True:
                norm.vmin = None
                
            if norm.vmax.__class__ == numpy.ma.MaskedArray and norm.vmax.mask == True:
                norm.vmax = None
                
        return norm
    
    def _getCbarMinMax(self):
        cbMin = self.colourBarMin
        cbMax = self.colourBarMax
        
        if cbMin is not None and cbMax is not None and cbMin > cbMax:
            log.warning("min(=%s) > max(=%s) reversing values" % (cbMin, cbMax))
            cbMax, cbMin = cbMin, cbMax
            
        return cbMin, cbMax        
    
    def __get_position(self): return self._position

    def __set_position(self, value):
        if value not in ['horizontal', 'vertical', None]:
            raise ValueError("ColourBar position value must be 'horizontal'" + \
                             " or 'vertical, value recieved :" + str(value))
        self._position = value

    colourBarPosition = property(__get_position, __set_position, None,
                    "colour bar position, 'horizontal' or 'vertical'")

    def __get_cmap(self): 
        if self._cmap == None:
            cmap = cm.jet
            cmap.set_bad('w')
            self.cmap = cmap
            
        return self.cmap
    
    def __set_cmap(self, value):
        self._cmap == value;
        

def getColourBarImage(width=600, height=100,
                      label=None, 
                      cmap=cm.jet, 
                      cmapRange=(0,1), 
                      orientation='horizontal',
                      dpi=100):
    
    figsize=(width / float(dpi), height / float(dpi))
    fig = Figure(figsize=figsize, dpi=dpi, facecolor='w')
    
    #for agg bakcend
    #need about 40px at the bottom of the axes to draw the labels
    #x = 40.0
    #cbMin = 0.5
    
    #for cairo backend
    x = 50.0
    cbMin = 0.6
    
    cbBottom = x/height
    
    if cbBottom < 0.1:
        cbBottom = 0.1
    
    if cbBottom > cbMin:
        cbBottom = cbMin
        
    cbHeight = 0.9 - cbBottom
    axes = fig.add_axes([0.05, cbBottom, 0.9, cbHeight], frameon=False)
    
    log.debug("cbBottom = %s, cbHeight = %s" % (cbBottom, cbHeight,))
    
    if cmapRange[0] > cmapRange[1]:
        log.warning("cmapRange[0] > cmapRange[1], swapping values")
        cmapRange = (cmapRange[1], cmapRange[0])
    
    cb = ColourBar(colourBarLabel=label, 
                   colourBarPosition=orientation, 
                   cmap=cmap, 
                   colourBarMin=cmapRange[0], 
                   colourBarMax=cmapRange[1] )
    
    cb.draw(axes)    
    
    return geoplot.utils.figureToImage(fig)

def getIntervalColourBarImage(grid, width=600, height=100,
                      label=None, 
                      cmap=cm.jet, 
                      cmapRange=(None, None), 
                      orientation='horizontal',
                      intervalNames=None,
                      dpi=100):
    
    figsize=(width / float(dpi), height / float(dpi))
    fig = Figure(figsize=figsize, dpi=dpi, facecolor='w')
    
    #for agg bakcend
    #need about 40px at the bottom of the axes to draw the labels
    #x = 40.0
    #cbMin = 0.5
    
    #for cairo backend
    x = 50.0
    cbMin = 0.6
    
    cbBottom = x/height
    
    if cbBottom < 0.1:
        cbBottom = 0.1
    
    if cbBottom > cbMin:
        cbBottom = cbMin
        
    cbHeight = 0.9 - cbBottom
    axes = fig.add_axes([0.05, cbBottom, 0.9, cbHeight], frameon=False)
    
    log.debug("cbBottom = %s, cbHeight = %s" % (cbBottom, cbHeight,))
    
    if cmapRange[0] > cmapRange[1]:
        log.warning("cmapRange[0] > cmapRange[1], swapping values")
        cmapRange = (cmapRange[1], cmapRange[0])
    
    cb = ColourBar(colourBarLabel=label, 
                   colourBarPosition=orientation, 
                   cmap=cmap, 
                   colourBarMin=cmapRange[0], 
                   colourBarMax=cmapRange[1],
                   intervalNames=intervalNames,
                   intervalColourbar=True)
    
    cb.draw(axes, grid=grid)    
    
    return geoplot.utils.figureToImage(fig)