diff --git a/fipy/viewers/matplotlibViewer/__init__.py b/fipy/viewers/matplotlibViewer/__init__.py index e0286f21ce..8495c7469a 100644 --- a/fipy/viewers/matplotlibViewer/__init__.py +++ b/fipy/viewers/matplotlibViewer/__init__.py @@ -5,6 +5,7 @@ from fipy.viewers.matplotlibViewer.matplotlib2DGridContourViewer import * from fipy.viewers.matplotlibViewer.matplotlib2DViewer import * from fipy.viewers.matplotlibViewer.matplotlibVectorViewer import * +from fipy.viewers.matplotlibViewer.matplotlibStreamViewer import * __all__ = ["MatplotlibViewer"] __all__.extend(matplotlib1DViewer.__all__) @@ -12,6 +13,7 @@ __all__.extend(matplotlib2DGridContourViewer.__all__) __all__.extend(matplotlib2DViewer.__all__) __all__.extend(matplotlibVectorViewer.__all__) +__all__.extend(matplotlibStreamViewer.__all__) def MatplotlibViewer(vars, title=None, limits={}, cmap=None, colorbar='vertical', axes=None, **kwlimits): """Generic function for creating a `MatplotlibViewer`. diff --git a/fipy/viewers/matplotlibViewer/matplotlibStreamViewer.py b/fipy/viewers/matplotlibViewer/matplotlibStreamViewer.py new file mode 100644 index 0000000000..632af021e9 --- /dev/null +++ b/fipy/viewers/matplotlibViewer/matplotlibStreamViewer.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python + +## -*-Pyth-*- + # ################################################################### + # FiPy - Python-based finite volume PDE solver + # + # FILE: "matplotlibStreamViewer.py" + # + # Author: Jonathan Guyer + # Author: Daniel Wheeler + # Author: James Warren + # mail: NIST + # www: http://www.ctcms.nist.gov/fipy/ + # + # ======================================================================== + # This software was developed at the National Institute of Standards + # and Technology by employees of the Federal Government in the course + # of their official duties. Pursuant to title 17 Section 105 of the + # United States Code this software is not subject to copyright + # protection and is in the public domain. FiPy is an experimental + # system. NIST assumes no responsibility whatsoever for its use by + # other parties, and makes no guarantees, expressed or implied, about + # its quality, reliability, or any other characteristic. We would + # appreciate acknowledgement if the software is used. + # + # This software can be redistributed and/or modified freely + # provided that any derivative works bear some notice that they are + # derived from it, and any modified versions bear some notice that + # they have been modified. + # ======================================================================== + # See the file "license.terms" for information on usage and redistribution + # of this file, and for a DISCLAIMER OF ALL WARRANTIES. + # + # ################################################################### + ## + +__docformat__ = 'restructuredtext' + +from fipy.tools import numerix +from fipy.variables.faceVariable import FaceVariable +from fipy.variables.cellVariable import CellVariable + +from fipy.viewers.matplotlibViewer.matplotlib2DViewer import AbstractMatplotlib2DViewer + +__all__ = ["MatplotlibStreamViewer"] + +class MatplotlibStreamViewer(AbstractMatplotlib2DViewer): + """Displays a stream plot of a 2D rank-1 `CellVariable` or + `FaceVariable` object using Matplotlib_ + + .. _Matplotlib: http://matplotlib.sourceforge.net/ + + """ + +# __doc__ += AbstractMatplotlib2DViewer._test2Dvector(viewer="MatplotlibStreamViewer") +# __doc__XXX += """ +# +# >>> for sparsity in numerix.arange(5000, 0, -500): +# ... viewer.stream(sparsity=sparsity) +# ... viewer.plot() +# >>> viewer._promptForOpinion() +# +# """ + __doc__ += AbstractMatplotlib2DViewer._test2DvectorIrregular(viewer="MatplotlibStreamViewer") + + 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 + + vertexIDs = mesh._orderedCellVertexIDs + + vertexCoords = mesh.getVertexCoords() + + X = numerix.take(vertexCoords[0], vertexIDs) + Y = numerix.take(vertexCoords[1], vertexIDs) + + xmin = X.min() + xmax = X.max() + ymin = Y.min() + ymax = Y.max() + + N = 100 + X = numerix.linspace(xmin, xmax, N) + Y = numerix.linspace(ymin, ymax, N) + + U, V = var(numerix.mgrid[xmin:xmax:N*1j, ymin:ymax:N*1j]) + +# 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.axes.streamplot(X, Y, U, V) + +# 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')) + +if __name__ == "__main__": + import fipy.tests.doctestPlus + fipy.tests.doctestPlus.execButNoTest()