from numpy import *
from enthought.mayavi import mlab

def complex_plot(function, 
                 maxval=9, 
                 xbounds=(-3.0,3.0), 
                 ybounds=(-3.0,3.0), 
                 stepsize=0.01, 
                 text="", 
                 textwidth=0.1, 
                 unitcylrange=(-4,4), 
                 zmax=9, 
                 zmin=-9, 
                 legend=True,
                 signed=True):
    try:
        engine = mayavi.engine
    except NameError:
        from enthought.mayavi.api import Engine
        engine = Engine()
        engine.start()

    f=mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1), size=(500,500))
    x,y = mgrid[xbounds[0]:xbounds[1]:stepsize, ybounds[0]:ybounds[1]:stepsize]
    if signed:
        z = maximum(minimum(abs(function(x+y*1j))*sign(real(function(x+y*1j))), zmax), zmin)
    else:
        z = maximum(minimum(abs(function(x+y*1j)), zmax), zmin)
    s = angle(function(x+y*1j))
    m=mlab.mesh(x,y,z, scalars=s)
    from enthought.mayavi.modules.axes import Axes
    axes = Axes()
    engine.add_filter(axes, m.module_manager)

    mlab.text(0, 0, text, width=textwidth)
    phi, v=mgrid[0.0:2.01*pi:0.1, unitcylrange[0]:unitcylrange[1]:0.1]
    mlab.mesh(cos(phi),sin(phi),v, color=(1,1,1), opacity=0.5)
    if legend:
        mlab.colorbar(title="phase", orientation='vertical')

complex_plot(lambda z: z, text="f(z)=z", zmin=-5, zmax=5, unitcylrange=(-1,1))
mlab.show()