def contour(q, z, res, ax=None, **kwargs):
"""Vectorized wrapper for contour plot.
@param q: coordinates of contour points
@type q: [2 x #points]
@param z: row vector of scalar function
@type z: [1 x #points]
@param res: resolution of contour
@type res: [nx, ny]
@param ax: axes object handle
@param kwargs: passed to contour function as given, see its help
@return: h = handle to contour object created.
"""
if ax is None:
ax = newax()
# multiple axes ?
try:
contours = []
for curax in ax:
cont = contour(curax, q, z, res, **kwargs)
contours.append(cont)
return contours
except:
pass
# enough dimensions ?
ndim = q.shape[0]
if ndim < 2:
raise Exception("space dim = q.shape[0] = 1.")
res = res2meshsize(res)
# Z ?
if z is None:
z = np.zeros(1, q.shape[1])
elif z is np.NaN:
z = 5 * np.random.rand(1, q.shape[1])
else:
z = z * np.ones([1, q.shape[1]])
X, Y = vec2meshgrid(q, res) # nargout=2
Z, = vec2meshgrid(z, res)
# calc
if ndim < 3:
cont = ax.contour(X, Y, Z, 100, **kwargs)
else:
raise Exception("Dimension of vector q is not 2.")
return cont
def surf_color(q, c, res, ax, **kwargs):
X, Y, Z = vec2meshgrid(q, res) # nargout=3
# no color ?
if 0 in c.shape:
h = ax.plot_surface(X, Y, Z, **kwargs)
else:
C, = vec2meshgrid(c, res)
h = ax.plot_surface(X, Y, Z, cmap=C, **kwargs)
return h
def streamplot(x, v, res, ax, **kwargs):
"""Streamplot in 2d or 3d.
@param x: points
@type x: 2d array, each column a point
@param v: vectors based at points in x
@type v: 2d array, each column a vector
@param ax: axes
@param kwargs: passed to mpl streamplot
"""
dim = dimension(x)
if dim == 2:
(X, Y) = vec2meshgrid(x, res)
(U, V) = vec2meshgrid(v, res)
ax.streamplot(X[0, :], Y[:, 1],
U, V, **kwargs)
else:
raise NotImplementedError
def ezcontour(func, ax, domain, resolution, values, **kwargs):
"""Vectorized easy contour,
for functions accepting vector arguments.
@param ax: axes object handle
@param func: function handle
@param domain: rectangular plotting domain
@type domain: [xmin, xmax, ymin, ymax]
@param resolution: grid spacing
@type resolution: [nx, ny]
@param values: level set values
@type values: [v1, v2, ..., vN]
@param kwargs: additional arguments for
input to func
"""
# which axes ?
if 0 in ax.shape:
warn("vezcontour:axes", "Axes object handle ax is empty, no plot.")
return
# which domain ?
if not domain:
domain = np.array([0, 1, 0, 1]).reshape(1, -1)
# at what grid resolution ?
if not resolution:
resolution = np.array([30, 29]).reshape(1, -1)
else:
if 0 in resolution.shape:
resolution = np.array([30, 29]).reshape(1, -1)
# which level sets ?
if not values:
values = np.array([])
# compute surface
q, X, Y = domain2vec(domain, resolution) # nargout=3
f = feval(func, q, **kwargs)
Z = vec2meshgrid(f, X)
# default level set values ?
if 0 in values.shape:
plt.contour(ax, X, Y, Z)
else:
plt.contour(ax, X, Y, Z, values)
return
def surf2(q, z, res, ax, **kwargs):
X, Y = vec2meshgrid(q, res) # nargout=2
Z, = vec2meshgrid(z, res)
h = ax.plot_surface(X, Y, Z, **kwargs)
return h
