from pylab import *

import matplotlib.numerix.ma as ma

#
# Has to specialise matplotlib.colors.normalize becasue code in axes.py
# checkes that the norm passed in is an instance of normalize.
#
class LogNorm(matplotlib.colors.normalize):
    def __init__(self, vmin=None, vmax=None, clip = True):
        """
        Normalize a given value to the 0-1 range on a log scale

        If vmin or vmax is not given, they are taken from the input's
        minimum and maximum value respectively.  If clip is True and
        the given value falls outside the range, the returned value
        will be 0 or 1, whichever is closer. Returns 0 if vmin==vmax.
        Works with scalars or arrays, including masked arrays.  If
        clip is True, masked values are set to 1; otherwise they
        remain masked.
        """
        self.vmin = vmin
        self.vmax = vmax
        self.clip = clip

    def __call__(self, value):

        if isinstance(value, (int, float)):
            vtype = 'scalar'
            val = ma.array([value])
        else:
            vtype = 'array'
            val = ma.asarray(value)

        self.autoscale(val)
        vmin, vmax = self.vmin, self.vmax
        if vmin > vmax:
            raise ValueError("minvalue must be less than or equal to maxvalue")
        elif vmin<=0:
            raise ValueError("values must all be positive")
        elif vmin==vmax:
            return 0.*value
        else:
            if self.clip:
                mask = ma.getmaskorNone(val)
                val = ma.array(clip(val.filled(vmax), vmin, vmax),
                                mask=mask)
            result = (log(val)-log(vmin))/float(log(vmax)-log(vmin))
        if vtype == 'scalar':
            result = result[0]
        return result

    def autoscale(self, A):
        if not self.scaled():
            if self.vmin is None: self.vmin = ma.minimum(A)
            if self.vmax is None: self.vmax = ma.maximum(A)

    def scaled(self):
        'return true if vmin and vmax set'
        return (self.vmin is not None and self.vmax is not None)




#class LogNorm(matplotlib.colors.normalize):
#     """
#     Normalize on a log scale.
#     """
#     def __init__(self,vmin=None, vmax=None, clip = True):
#         matplotlib.colors.normalize.__init__(self,vmin=log(vmin),vmax=log(vmax),clip=clip)
#     
#     def __call__(self, value):
#         return matplotlib.colors.normalize.__call__(self, log(value))


delta = 0.025
x = arange(-3.0, 3.0, delta)
y = arange(-2.0, 2.0, delta) 

X, Y = meshgrid(x, y)

# A low hump with a spike coming out of the top right. 
# Needs to have z/colour axis on a log scale so we both hump and spike.
# linear scale only shows the spike.
Z1 = bivariate_normal(X, Y, 0.1, 0.2, 1.0, 1.0) + 0.1*bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)

pcolor(X,Y,Z1,shading='flat',norm=LogNorm(vmin=Z1.min(),vmax=Z1.max()))

colorbar()

show()
#savefig('pcolor_log.png')