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