Example #1
0
 def zcolor(z):
     return color[argmin(abs(levels - z))]
Example #2
0
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])