'''

_slice_dict = {'neither': slice(0, None),

'both': slice(1, -1),

'min': slice(1, None),

'max': slice(0, -1)}

n_rasterize = 50 # rasterize solids if number of colors >= n_rasterize

def __init__(self, ax, cmap=None,

norm=None,

alpha=None,

values=None,

boundaries=None,

orientation='vertical',

ticklocation='auto',

extend='neither',

spacing='uniform', # uniform or proportional

ticks=None,

format=None,

drawedges=False,

filled=True,

extendfrac=None,

extendrect=False,

label='',

):

#: The axes that this colorbar lives in.

self.ax = ax

self._patch_ax()

if cmap is None:

cmap = cm.get_cmap()

if norm is None:

norm = colors.Normalize()

self.alpha = alpha

cm.ScalarMappable.__init__(self, cmap=cmap, norm=norm)

self.values = values

self.boundaries = boundaries

self.extend = extend

self._inside = self._slice_dict[extend]

self.spacing = spacing

self.orientation = orientation

self.drawedges = drawedges

self.filled = filled

self.extendfrac = extendfrac

self.extendrect = extendrect

self.solids = None

self.lines = list()

self.outline = None

self.patch = None

self.dividers = None

if ticklocation == 'auto':

ticklocation = 'bottom' if orientation == 'horizontal' else 'right'

self.ticklocation = ticklocation

self.set_label(label)

if cbook.iterable(ticks):

self.locator = ticker.FixedLocator(ticks, nbins=len(ticks))

else:

self.locator = ticks # Handle default in _ticker()

if format is None:

if isinstance(self.norm, colors.LogNorm):

self.formatter = ticker.LogFormatterMathtext()

else:

self.formatter = ticker.ScalarFormatter()

elif cbook.is_string_like(format):

self.formatter = ticker.FormatStrFormatter(format)

else:

self.formatter = format # Assume it is a Formatter

# The rest is in a method so we can recalculate when clim changes.

self.config_axis()

self.draw_all()

def _extend_lower(self):

"""Returns whether the lower limit is open ended."""

return self.extend in ('both', 'min')

def _extend_upper(self):

"""Returns whether the uper limit is open ended."""

return self.extend in ('both', 'max')

def _patch_ax(self):

# bind some methods to the axes to warn users

# against using those methods.

self.ax.set_xticks = _set_ticks_on_axis_warn

self.ax.set_yticks = _set_ticks_on_axis_warn

def draw_all(self):

'''

self._process_values()

self._find_range()

X, Y = self._mesh()

C = self._values[:, np.newaxis]

self._config_axes(X, Y)

if self.filled:

self._add_solids(X, Y, C)

def config_axis(self):

ax = self.ax

if self.orientation == 'vertical':

ax.xaxis.set_ticks([])

# location is either one of 'bottom' or 'top'

ax.yaxis.set_label_position(self.ticklocation)

ax.yaxis.set_ticks_position(self.ticklocation)

else:

ax.yaxis.set_ticks([])

# location is either one of 'left' or 'right'

ax.xaxis.set_label_position(self.ticklocation)

ax.xaxis.set_ticks_position(self.ticklocation)

self._set_label()

def update_ticks(self):

"""

ax = self.ax

ticks, ticklabels, offset_string = self._ticker()

if self.orientation == 'vertical':

ax.yaxis.set_ticks(ticks)

ax.set_yticklabels(ticklabels)

ax.yaxis.get_major_formatter().set_offset_string(offset_string)

else:

ax.xaxis.set_ticks(ticks)

ax.set_xticklabels(ticklabels)

ax.xaxis.get_major_formatter().set_offset_string(offset_string)

def set_ticks(self, ticks, update_ticks=True):

"""

if cbook.iterable(ticks):

self.locator = ticker.FixedLocator(ticks, nbins=len(ticks))

else:

self.locator = ticks

if update_ticks:

self.update_ticks()

self.stale = True

def get_ticks(self, minor=False):

"""Return the x ticks as a list of locations"""

return self._tick_data_values

def set_ticklabels(self, ticklabels, update_ticks=True):

"""

if isinstance(self.locator, ticker.FixedLocator):

self.formatter = ticker.FixedFormatter(ticklabels)

if update_ticks:

self.update_ticks()

else:

warnings.warn("set_ticks() must have been called.")

self.stale = True

def _config_axes(self, X, Y):

'''

ax = self.ax

ax.set_frame_on(False)

ax.set_navigate(False)

xy = self._outline(X, Y)

ax.update_datalim(xy)

ax.set_xlim(*ax.dataLim.intervalx)

ax.set_ylim(*ax.dataLim.intervaly)

if self.outline is not None:

self.outline.remove()

self.outline = mpatches.Polygon(

xy, edgecolor=mpl.rcParams['axes.edgecolor'],

facecolor='none',

linewidth=mpl.rcParams['axes.linewidth'],

closed=True,

zorder=2)

ax.add_artist(self.outline)

self.outline.set_clip_box(None)

self.outline.set_clip_path(None)

c = mpl.rcParams['axes.facecolor']

if self.patch is not None:

self.patch.remove()

self.patch = mpatches.Polygon(xy, edgecolor=c,

facecolor=c,

linewidth=0.01,

zorder=-1)

ax.add_artist(self.patch)

self.update_ticks()

def _set_label(self):

if self.orientation == 'vertical':

self.ax.set_ylabel(self._label, **self._labelkw)

else:

self.ax.set_xlabel(self._label, **self._labelkw)

self.stale = True

def set_label(self, label, **kw):

'''

self._label = '%s' % (label, )

self._labelkw = kw

self._set_label()

def _outline(self, X, Y):

'''

N = X.shape[0]

ii = [0, 1, N - 2, N - 1, 2 * N - 1, 2 * N - 2, N + 1, N, 0]

x = np.take(np.ravel(np.transpose(X)), ii)

y = np.take(np.ravel(np.transpose(Y)), ii)

x = x.reshape((len(x), 1))

y = y.reshape((len(y), 1))

if self.orientation == 'horizontal':

return np.hstack((y, x))

return np.hstack((x, y))

def _edges(self, X, Y):

'''

N = X.shape[0]

# Using the non-array form of these line segments is much

# simpler than making them into arrays.

if self.orientation == 'vertical':

return [list(zip(X[i], Y[i])) for i in xrange(1, N - 1)]

else:

return [list(zip(Y[i], X[i])) for i in xrange(1, N - 1)]

def _add_solids(self, X, Y, C):

'''

if self.orientation == 'vertical':

args = (X, Y, C)

else:

args = (np.transpose(Y), np.transpose(X), np.transpose(C))

kw = dict(cmap=self.cmap,

norm=self.norm,

alpha=self.alpha,

edgecolors='None')

# Save, set, and restore hold state to keep pcolor from

# clearing the axes. Ordinarily this will not be needed,

# since the axes object should already have hold set.

_hold = self.ax.ishold()

self.ax.hold(True)

col = self.ax.pcolormesh(*args, **kw)

self.ax.hold(_hold)

#self.add_observer(col) # We should observe, not be observed...

if self.solids is not None:

self.solids.remove()

self.solids = col

if self.dividers is not None:

self.dividers.remove()

self.dividers = None

if self.drawedges:

linewidths = (0.5 * mpl.rcParams['axes.linewidth'],)

self.dividers = collections.LineCollection(self._edges(X, Y),

colors=(mpl.rcParams['axes.edgecolor'],),

linewidths=linewidths)

self.ax.add_collection(self.dividers)

elif len(self._y) >= self.n_rasterize:

self.solids.set_rasterized(True)

def add_lines(self, levels, colors, linewidths, erase=True):

'''

y = self._locate(levels)

igood = (y < 1.001) & (y > -0.001)

y = y[igood]

if cbook.iterable(colors):

colors = np.asarray(colors)[igood]

if cbook.iterable(linewidths):

linewidths = np.asarray(linewidths)[igood]

N = len(y)

x = np.array([0.0, 1.0])

X, Y = np.meshgrid(x, y)

if self.orientation == 'vertical':

xy = [list(zip(X[i], Y[i])) for i in xrange(N)]

else:

xy = [list(zip(Y[i], X[i])) for i in xrange(N)]

col = collections.LineCollection(xy, linewidths=linewidths)

if erase and self.lines:

for lc in self.lines:

lc.remove()

self.lines = []

self.lines.append(col)

col.set_color(colors)

self.ax.add_collection(col)

self.stale = True

def _ticker(self):

'''

locator = self.locator

formatter = self.formatter

if locator is None:

if self.boundaries is None:

if isinstance(self.norm, colors.NoNorm):

nv = len(self._values)

base = 1 + int(nv / 10)

locator = ticker.IndexLocator(base=base, offset=0)

elif isinstance(self.norm, colors.BoundaryNorm):

b = self.norm.boundaries

locator = ticker.FixedLocator(b, nbins=10)

elif isinstance(self.norm, colors.LogNorm):

locator = ticker.LogLocator()

else:

if mpl.rcParams['_internal.classic_mode']:

locator = ticker.MaxNLocator()

else:

locator = ticker.AutoLocator()

else:

b = self._boundaries[self._inside]

locator = ticker.FixedLocator(b, nbins=10)

if isinstance(self.norm, colors.NoNorm) and self.boundaries is None:

intv = self._values[0], self._values[-1]

else:

intv = self.vmin, self.vmax

locator.create_dummy_axis(minpos=intv[0])

formatter.create_dummy_axis(minpos=intv[0])

locator.set_view_interval(*intv)

locator.set_data_interval(*intv)

formatter.set_view_interval(*intv)

formatter.set_data_interval(*intv)

b = np.array(locator())

if isinstance(locator, ticker.LogLocator):

eps = 1e-10

b = b[(b <= intv[1] * (1 + eps)) & (b >= intv[0] * (1 - eps))]

else:

eps = (intv[1] - intv[0]) * 1e-10

b = b[(b <= intv[1] + eps) & (b >= intv[0] - eps)]

self._tick_data_values = b

ticks = self._locate(b)

formatter.set_locs(b)

ticklabels = [formatter(t, i) for i, t in enumerate(b)]

offset_string = formatter.get_offset()

return ticks, ticklabels, offset_string

def _process_values(self, b=None):

'''

if b is None:

b = self.boundaries

if b is not None:

self._boundaries = np.asarray(b, dtype=float)

if self.values is None:

self._values = 0.5 * (self._boundaries[:-1]

+ self._boundaries[1:])

if isinstance(self.norm, colors.NoNorm):

self._values = (self._values + 0.00001).astype(np.int16)

return

self._values = np.array(self.values)

return

if self.values is not None:

self._values = np.array(self.values)

if self.boundaries is None:

b = np.zeros(len(self.values) + 1, 'd')

b[1:-1] = 0.5 * (self._values[:-1] - self._values[1:])

b[0] = 2.0 * b[1] - b[2]

b[-1] = 2.0 * b[-2] - b[-3]

self._boundaries = b

return

self._boundaries = np.array(self.boundaries)

return

# Neither boundaries nor values are specified;

# make reasonable ones based on cmap and norm.

if isinstance(self.norm, colors.NoNorm):

b = self._uniform_y(self.cmap.N + 1) * self.cmap.N - 0.5

v = np.zeros((len(b) - 1,), dtype=np.int16)

v[self._inside] = np.arange(self.cmap.N, dtype=np.int16)

if self._extend_lower():

v[0] = -1

if self._extend_upper():

v[-1] = self.cmap.N

self._boundaries = b

self._values = v

return

elif isinstance(self.norm, colors.BoundaryNorm):

b = list(self.norm.boundaries)

if self._extend_lower():

b = [b[0] - 1] + b

if self._extend_upper():

b = b + [b[-1] + 1]

b = np.array(b)

v = np.zeros((len(b) - 1,), dtype=float)

bi = self.norm.boundaries

v[self._inside] = 0.5 * (bi[:-1] + bi[1:])

if self._extend_lower():

v[0] = b[0] - 1

if self._extend_upper():

v[-1] = b[-1] + 1

self._boundaries = b

self._values = v

return

else:

if not self.norm.scaled():

self.norm.vmin = 0

self.norm.vmax = 1

self.norm.vmin, self.norm.vmax = mtrans.nonsingular(self.norm.vmin,

self.norm.vmax,

expander=0.1)

b = self.norm.inverse(self._uniform_y(self.cmap.N + 1))

if self._extend_lower():

b[0] = b[0] - 1

if self._extend_upper():

b[-1] = b[-1] + 1

self._process_values(b)

def _find_range(self):

'''

b = self._boundaries[self._inside]

self.vmin = b[0]

self.vmax = b[-1]

def _central_N(self):

'''number of boundaries **before** extension of ends'''

nb = len(self._boundaries)

if self.extend == 'both':

nb -= 2

elif self.extend in ('min', 'max'):

nb -= 1

return nb

def _extended_N(self):

'''

N = self.cmap.N + 1

if self.extend == 'both':

N += 2

elif self.extend in ('min', 'max'):

N += 1

return N

def _get_extension_lengths(self, frac, automin, automax, default=0.05):

'''

# Set the default value.

extendlength = np.array([default, default])

if isinstance(frac, six.string_types):

if frac.lower() == 'auto':

# Use the provided values when 'auto' is required.

extendlength[0] = automin

extendlength[1] = automax

else:

# Any other string is invalid.

raise ValueError('invalid value for extendfrac')

elif frac is not None:

try:

# Try to set min and max extension fractions directly.

extendlength[:] = frac

# If frac is a sequence contaning None then NaN may

# be encountered. This is an error.

if np.isnan(extendlength).any():

raise ValueError()

except (TypeError, ValueError):

# Raise an error on encountering an invalid value for frac.

raise ValueError('invalid value for extendfrac')

return extendlength

def _uniform_y(self, N):

'''

if self.extend == 'neither':

y = np.linspace(0, 1, N)

else:

automin = automax = 1. / (N - 1.)

extendlength = self._get_extension_lengths(self.extendfrac,

automin, automax,

default=0.05)

if self.extend == 'both':

y = np.zeros(N + 2, 'd')

y[0] = 0. - extendlength[0]

y[-1] = 1. + extendlength[1]

elif self.extend == 'min':

y = np.zeros(N + 1, 'd')

y[0] = 0. - extendlength[0]

else:

y = np.zeros(N + 1, 'd')

y[-1] = 1. + extendlength[1]

y[self._inside] = np.linspace(0, 1, N)

return y

def _proportional_y(self):

'''

if isinstance(self.norm, colors.BoundaryNorm):

y = (self._boundaries - self._boundaries[0])

y = y / (self._boundaries[-1] - self._boundaries[0])

else:

y = self.norm(self._boundaries.copy())

if self.extend == 'min':

# Exclude leftmost interval of y.

clen = y[-1] - y[1]

automin = (y[2] - y[1]) / clen

automax = (y[-1] - y[-2]) / clen

elif self.extend == 'max':

# Exclude rightmost interval in y.

clen = y[-2] - y[0]

automin = (y[1] - y[0]) / clen

automax = (y[-2] - y[-3]) / clen

else:

# Exclude leftmost and rightmost intervals in y.

clen = y[-2] - y[1]

automin = (y[2] - y[1]) / clen

automax = (y[-2] - y[-3]) / clen

extendlength = self._get_extension_lengths(self.extendfrac,

automin, automax,

default=0.05)

if self.extend in ('both', 'min'):

y[0] = 0. - extendlength[0]

if self.extend in ('both', 'max'):

y[-1] = 1. + extendlength[1]

yi = y[self._inside]

norm = colors.Normalize(yi[0], yi[-1])

y[self._inside] = norm(yi)

return y

def _mesh(self):

'''

x = np.array([0.0, 1.0])

if self.spacing == 'uniform':

y = self._uniform_y(self._central_N())

else:

y = self._proportional_y()

self._y = y

X, Y = np.meshgrid(x, y)

if self._extend_lower() and not self.extendrect:

X[0, :] = 0.5

if self._extend_upper() and not self.extendrect:

X[-1, :] = 0.5

return X, Y

def _locate(self, x):

'''

if isinstance(self.norm, (colors.NoNorm, colors.BoundaryNorm)):

b = self._boundaries

xn = x

else:

# Do calculations using normalized coordinates so

# as to make the interpolation more accurate.

b = self.norm(self._boundaries, clip=False).filled()

xn = self.norm(x, clip=False).filled()

# The rest is linear interpolation with extrapolation at ends.

ii = np.searchsorted(b, xn)

i0 = ii - 1

itop = (ii == len(b))

ibot = (ii == 0)

i0[itop] -= 1

ii[itop] -= 1

i0[ibot] += 1

ii[ibot] += 1

db = np.take(b, ii) - np.take(b, i0)

y = self._y

dy = np.take(y, ii) - np.take(y, i0)

z = np.take(y, i0) + (xn - np.take(b, i0)) * dy / db

return z

def set_alpha(self, alpha):

self.alpha = alpha

def remove(self):

"""

fig = self.ax.figure

"""

def __init__(self, ax, mappable, **kw):

# Ensure the given mappable's norm has appropriate vmin and vmax set

# even if mappable.draw has not yet been called.

mappable.autoscale_None()

self.mappable = mappable

kw['cmap'] = cmap = mappable.cmap

kw['norm'] = norm = mappable.norm

if isinstance(mappable, contour.ContourSet):

CS = mappable

kw['alpha'] = mappable.get_alpha()

kw['boundaries'] = CS._levels

kw['values'] = CS.cvalues

kw['extend'] = CS.extend

#kw['ticks'] = CS._levels

kw.setdefault('ticks', ticker.FixedLocator(CS.levels, nbins=10))

kw['filled'] = CS.filled

ColorbarBase.__init__(self, ax, **kw)

if not CS.filled:

self.add_lines(CS)

else:

if getattr(cmap, 'colorbar_extend', False) is not False:

kw.setdefault('extend', cmap.colorbar_extend)

if isinstance(mappable, martist.Artist):

kw['alpha'] = mappable.get_alpha()

ColorbarBase.__init__(self, ax, **kw)

def on_mappable_changed(self, mappable):

"""

self.set_cmap(mappable.get_cmap())

self.set_clim(mappable.get_clim())

self.update_normal(mappable)

def add_lines(self, CS, erase=True):

'''

if not isinstance(CS, contour.ContourSet) or CS.filled:

raise ValueError('add_lines is only for a ContourSet of lines')

tcolors = [c[0] for c in CS.tcolors]

tlinewidths = [t[0] for t in CS.tlinewidths]

# The following was an attempt to get the colorbar lines

# to follow subsequent changes in the contour lines,

# but more work is needed: specifically, a careful

# look at event sequences, and at how

# to make one object track another automatically.

#tcolors = [col.get_colors()[0] for col in CS.collections]

#tlinewidths = [col.get_linewidth()[0] for lw in CS.collections]

#print 'tlinewidths:', tlinewidths

ColorbarBase.add_lines(self, CS.levels, tcolors, tlinewidths,

erase=erase)

def update_normal(self, mappable):

'''

self.draw_all()

if isinstance(self.mappable, contour.ContourSet):

CS = self.mappable

if not CS.filled:

self.add_lines(CS)

self.stale = True

def update_bruteforce(self, mappable):

'''

# We are using an ugly brute-force method: clearing and

# redrawing the whole thing. The problem is that if any

# properties have been changed by methods other than the

# colorbar methods, those changes will be lost.

self.ax.cla()

# clearing the axes will delete outline, patch, solids, and lines:

self.outline = None

self.patch = None

self.solids = None

self.lines = list()

self.dividers = None

self.set_alpha(mappable.get_alpha())

self.cmap = mappable.cmap

self.norm = mappable.norm

self.config_axis()

self.draw_all()

if isinstance(self.mappable, contour.ContourSet):

CS = self.mappable

if not CS.filled:

self.add_lines(CS)

#if self.lines is not None:

# tcolors = [c[0] for c in CS.tcolors]

# self.lines.set_color(tcolors)

#Fixme? Recalculate boundaries, ticks if vmin, vmax have changed.

#Fixme: Some refactoring may be needed; we should not

# be recalculating everything if there was a simple alpha

# change.