def zcolor(z):
return color[argmin(abs(levels - z))]
def pyxplotcontour(canvas, data, x=None, y=None, levels=8,
colors='map', colmap=None, color=None,
labels=False, labelsize=0.4,
**kwargs):
"""do contour plot.
data[y, x]: 2d real array containing z values
x, y: None - equidistant point spacing
1d array: rectangular non-equidistant grid
2d array: arbitrary coordinates
levels: int (level count) or list of z values
default: 8 equispaced levels
colors: "map": use colmap (ColorMapper)
"map_invert": use colmap but invert color
(max visibility on top of bitmap plot)
"color_array": color and levels must be arrays of the same
length, level i get's the color i
"color": use color (as in pyxgraph.pyxplot)
labels: draw text labels designating the level value of each contour
labelsize: label text size (relative to normal text size)
kwargs: are passed to pyxplot
"""
zmin, zmax = data.min(), data.max()
# define a function zcolor(z) that returns the color to use
if 'map' in colors:
if colmap==None:
colmap=ColMapper.ColorMapper("gauss")
if 'inv' in colors:
def zcolor(z):
r, g, b = colmap.colfct((z-zmin)/(zmax-zmin))
return 1-r, 1-g, 1-b
else:
def zcolor(z):
return colmap.colfct((z-zmin)/(zmax-zmin))
elif colors == "color_array":
levels = array(levels)
assert hasattr(color, "__iter__") # for this color and levels
assert hasattr(levels, "__iter__") # have to be lists of the
assert len(color)==len(levels) # same length
def zcolor(z):
return color[argmin(abs(levels - z))]
# return most suitable color
else: #use color
def zcolor(z):
return color
# calculate level list
if isinstance(levels, int):
levels = (arange(levels)+0.5)*(zmax-zmin)/levels + zmin
# calculate the levels
if x is None or y is None:
clines = contour_nogrid(data, levels)
elif len(x.shape) == 1 and len(y.shape) == 1:
clines = contour_rectgrid(x, y, data, levels)
elif x.shape == y.shape == data.shape:
clines = contour_gencoord(x, y, data, levels)
else:
raise ValueError, 'pyxplotcontour: x and/or y misshaped'
# clines is clines[levidx][lineidx][0 or 1 for x/y]
# set default args
if 'linetype' not in kwargs and 'lt' not in kwargs:
kwargs['linetype']=pyx.style.linestyle.solid
if 'style' not in kwargs:
kwargs['style'] = 'lines'
# plot
for levidx, level in enumerate(levels):
llabel = '%0.4g'%level
levelcolor = zcolor(level)
for line in clines[levidx]:
canvas.pyxplot((line[0], line[1]),
color=levelcolor, title=False, **kwargs)
if labels:
# find label position
x, y = line
xy = map(canvas.pos, x, y)
canvas_x = array([i[0] for i in xy])
canvas_y = array([i[1] for i in xy])
dcx = canvas_x[1:] - canvas_x[:-1]
dcy = canvas_y[1:] - canvas_y[:-1]
dcx = where(dcx==0, min(dcx)/10, dcx) # avoid singularity problems
ind = argmin(abs(dcy/dcx))
# try to avoid placing labels on plot borders
# FIXME: find better method for that?
relx = canvas_x[ind] - canvas.xpos
if not (labelsize < relx < canvas.width-labelsize): # label at left or right border
#if ind < 3 or ind > len(x)-4: # probably a line ending at border
ind = len(x) / 2
# label position (graphcoords)
lx, ly = x[ind], y[ind]
angle = arctan(dcy[ind]/dcx[ind])*180/pi
if abs(angle) < 5:
angle = 0
# FIXME: do something against labels exceeding plot area
canvas.pyxlabel((lx, ly), llabel, graphcoords=True,
style=[pyx.trafo.translate(0, 0.05),
pyx.trafo.rotate(angle),
pyx.trafo.scale(labelsize),
pyx.text.halign.center,
pyx.text.valign.bottom])