class MatplotlibStreamViewer(AbstractMatplotlib2DViewer): """Displays a stream plot of a 2D rank-1 `CellVariable` or `FaceVariable` object using Matplotlib_ One issue is that this `Viewer` relies on `scipy.interpolate.griddata`, which interpolates on the convex hull of the data. The results is that streams are plotted across any concavities in the mesh. Another issue is that it does not seem possible to remove the streams without calling `cla()`, which means that different set of streams cannot be overlaid. .. _Matplotlib: http://matplotlib.sourceforge.net/ """ __doc__ += AbstractMatplotlib2DViewer._test2Dvector( viewer="MatplotlibStreamViewer") __doc__ += AbstractMatplotlib2DViewer._test2DvectorIrregular( viewer="MatplotlibStreamViewer") def __init__(self, vars, title=None, log=False, limits={}, axes=None, figaspect='auto', density=1, linewidth=None, color=None, cmap=None, norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.1, **kwlimits): """Creates a `MatplotlibStreamViewer`. :Parameters: vars a rank-1 `CellVariable` or `FaceVariable` object. title displayed at the top of the `Viewer` window log if `True`, arrow length goes at the base-10 logarithm of the magnitude limits : dict a (deprecated) alternative to limit keyword arguments xmin, xmax, ymin, ymax, datamin, datamax displayed range of data. Any limit set to a (default) value of `None` will autoscale. axes if not `None`, `vars` will be plotted into this Matplotlib `Axes` object figaspect desired aspect ratio of figure. If arg is a number, use that aspect ratio. If arg is 'auto', the aspect ratio will be determined from the Variable's mesh. *density* : float or 2-tuple Controls the closeness of streamlines. When `density = 1`, the domain is divided into a 25x25 grid---*density* linearly scales this grid. Each cell in the grid can have, at most, one traversing streamline. For different densities in each direction, use [density_x, density_y]. linewidth : Numeric or rank-0 `MeshVariable` vary linewidth when given a `CellVariable` or `FaceVariable` of same type as vars. *color* : matplotlib color code, or rank-0 `MeshVariable` Streamline color. When given an array with the type as vars, *color* values are converted to colors using *cmap*. *cmap* : :class:`~matplotlib.colors.Colormap` Colormap used to plot streamlines and arrows. Only necessary when using an `MeshVariable` input for *color*. *norm* : :class:`~matplotlib.colors.Normalize` Normalize object used to scale luminance data to 0, 1. If None, stretch (min, max) to (0, 1). Only necessary when *color* is an `MeshVariable`. *arrowsize* : float Factor scale arrow size. *arrowstyle* : str Arrow style specification. See :class:`~matplotlib.patches.FancyArrowPatch`. *minlength* : float Minimum length of streamline in axes coordinates. """ kwlimits.update(limits) AbstractMatplotlib2DViewer.__init__(self, vars=vars, title=title, axes=axes, figaspect=figaspect, **kwlimits) self.log = log self.kwargs = dict(density=density, cmap=cmap, norm=norm, arrowsize=arrowsize, arrowstyle=arrowstyle, minlength=minlength) self.linewidth = linewidth self.color = color self._stream = None self._plot() def _getSuitableVars(self, vars): from fipy.meshes.mesh2D import Mesh2D from fipy.meshes.uniformGrid2D import UniformGrid2D vars = [var for var in AbstractMatplotlib2DViewer._getSuitableVars(self, vars) \ if ((isinstance(var.mesh, Mesh2D) or isinstance(var.mesh, UniformGrid2D))\ and (isinstance(var, FaceVariable) \ or isinstance(var, CellVariable)) and var.rank == 1)] if len(vars) == 0: from fipy.viewers import MeshDimensionError raise MeshDimensionError("The mesh must be a Mesh2D instance") # this viewer can only display one variable return [vars[0]] def _plot(self): from scipy.interpolate import griddata var = self.vars[0] mesh = var.mesh xmin, ymin = mesh.extents['min'] xmax, ymax = mesh.extents['max'] N = 100 X = numerix.linspace(xmin, xmax, N) Y = numerix.linspace(ymin, ymax, N) grid_x, grid_y = numerix.mgrid[xmin:xmax:N * 1j, ymin:ymax:N * 1j] if isinstance(var, FaceVariable): C = mesh.faceCenters elif isinstance(var, CellVariable): C = mesh.cellCenters U = griddata(C.value.T, var.value[0], (grid_x, grid_y), method='cubic') V = griddata(C.value.T, var.value[1], (grid_x, grid_y), method='cubic') lw = self.linewidth if isinstance(lw, (FaceVariable, CellVariable)): lw = griddata(C.value.T, lw.value, (grid_x, grid_y), method='cubic') color = self.color if isinstance(color, (FaceVariable, CellVariable)): color = griddata(C.value.T, color.value, (grid_x, grid_y), method='cubic', fill_value=color.min()) U = U.T V = V.T ang = numerix.arctan2(V, U) mag = numerix.sqrt(U**2 + V**2) datamin, datamax = self._autoscale(vars=(mag, ), datamin=self._getLimit('datamin'), datamax=self._getLimit('datamax')) mag = numerix.where(mag > datamax, numerix.nan, mag) mag = numerix.where(mag < datamin, numerix.nan, mag) if self.log: mag = numerix.log10(mag) U = mag * numerix.cos(ang) V = mag * numerix.sin(ang) # if self._stream is not None: # # the following doesn't work, nor does it help to `add_collection` first # # self._stream.arrows.remove() # self._stream.lines.remove() self.axes.cla() self._stream = self.axes.streamplot(X, Y, U, V, linewidth=lw, color=color, **self.kwargs) self.axes.set_xlim(xmin=self._getLimit('xmin'), xmax=self._getLimit('xmax')) self.axes.set_ylim(ymin=self._getLimit('ymin'), ymax=self._getLimit('ymax'))
class MatplotlibVectorViewer(AbstractMatplotlib2DViewer): """Displays a vector plot of a 2D rank-1 `CellVariable` or `FaceVariable` object using Matplotlib_ .. _Matplotlib: http://matplotlib.sourceforge.net/ """ __doc__ += AbstractMatplotlib2DViewer._test2Dvector( viewer="MatplotlibVectorViewer") __doc__ += """ >>> for sparsity in numerix.arange(5000, 0, -500): ... viewer.quiver(sparsity=sparsity) ... viewer.plot() >>> viewer._promptForOpinion() """ __doc__ += AbstractMatplotlib2DViewer._test2DvectorIrregular( viewer="MatplotlibVectorViewer") def __init__(self, vars, title=None, scale=None, sparsity=None, log=False, limits={}, axes=None, figaspect='auto', **kwlimits): """Creates a `Matplotlib2DViewer`. :Parameters: vars a rank-1 `CellVariable` or `FaceVariable` object. title displayed at the top of the `Viewer` window scale if not `None`, scale all arrow lengths by this value sparsity if not `None`, then this number of arrows will be randomly chosen (weighted by the cell volume or face area) log if `True`, arrow length goes at the base-10 logarithm of the magnitude limits : dict a (deprecated) alternative to limit keyword arguments xmin, xmax, ymin, ymax, datamin, datamax displayed range of data. Any limit set to a (default) value of `None` will autoscale. axes if not `None`, `vars` will be plotted into this Matplotlib `Axes` object figaspect desired aspect ratio of figure. If arg is a number, use that aspect ratio. If arg is 'auto', the aspect ratio will be determined from the Variable's mesh. """ kwlimits.update(limits) AbstractMatplotlib2DViewer.__init__(self, vars=vars, title=title, axes=axes, figaspect=figaspect, **kwlimits) self.quiver(sparsity=sparsity, scale=scale) self.log = log self._plot() def quiver(self, sparsity=None, scale=None): var = self.vars[0] mesh = var.mesh if isinstance(var, FaceVariable): N = mesh.numberOfFaces X, Y = mesh.faceCenters elif isinstance(var, CellVariable): N = mesh.numberOfCells X, Y = mesh.cellCenters if sparsity is not None and N > sparsity: XYrand = numerix.random.random((2, sparsity)) XYrand = numerix.array([ [min(X)], [min(Y)] ]) + XYrand * numerix.array([[max(X) - min(X)], [max(Y) - min(Y)]]) self.indices = numerix.nearest(numerix.array([X, Y]), XYrand) else: self.indices = numerix.arange(N) X = numerix.take(X, self.indices) Y = numerix.take(Y, self.indices) U = V = numerix.ones(X.shape, 'l') if hasattr(self, "_quiver"): self._quiver.remove() self._quiver = self.axes.quiver(X, Y, U, V, scale=scale, pivot='middle') def _getSuitableVars(self, vars): from fipy.meshes.mesh2D import Mesh2D from fipy.meshes.uniformGrid2D import UniformGrid2D vars = [var for var in AbstractMatplotlib2DViewer._getSuitableVars(self, vars) \ if ((isinstance(var.mesh, Mesh2D) or isinstance(var.mesh, UniformGrid2D))\ and (isinstance(var, FaceVariable) \ or isinstance(var, CellVariable)) and var.rank == 1)] if len(vars) == 0: from fipy.viewers import MeshDimensionError raise MeshDimensionError, "The mesh must be a Mesh2D instance" # this viewer can only display one variable return [vars[0]] def _plot(self): var = self.vars[0] mesh = var.mesh U, V = var.numericValue U = numerix.take(U, self.indices) V = numerix.take(V, self.indices) ang = numerix.arctan2(V, U) mag = numerix.sqrt(U**2 + V**2) datamin, datamax = self._autoscale(vars=(mag, ), datamin=self._getLimit('datamin'), datamax=self._getLimit('datamax')) mag = numerix.where(mag > datamax, datamax, mag) mag = numerix.ma.masked_array(mag, mag < datamin) if self.log: mag = numerix.log10(mag) mag = numerix.ma.masked_array(mag, numerix.isnan(mag)) U = mag * numerix.cos(ang) V = mag * numerix.sin(ang) self._quiver.set_UVC(U, V) self.axes.set_xlim(xmin=self._getLimit('xmin'), xmax=self._getLimit('xmax')) self.axes.set_ylim(ymin=self._getLimit('ymin'), ymax=self._getLimit('ymax'))