added plotting for mesh

tfields/__init__.py

@@ -2,6 +2,7 @@ from . import core
 from . import bases
 from . import lib
 from .lib import *
+from . import plotting
 
 # __all__ = ['core', 'points3D']
 from .core import Tensors, TensorFields, TensorMaps

tfields/mesh3D.py

@@ -723,11 +723,32 @@ class Mesh3D(tfields.TensorMaps):
             return obj, template
         return obj
 
-    def plot(self):
-        import mplTools as mpt
-        mpt.plotMesh(self, self.faces, color=self.maps[0].fields[0], vmin=0,
-                     vmax=20, axis=mpt.gca(3))
-        mpt.plt.show()
+    def plot(self, **kwargs):
+        """
+        Forwarding to plotTools.plotMesh
+        """
+        scalars_demanded = any([v in kwargs for v in ['vmin', 'vmax', 'cmap']])
+        map_index = kwargs.pop('map_index', None if not scalars_demanded else 0)
+        if map_index is not None:
+            if not len(self.maps[0]) == 0:
+                kwargs['color'] = self.maps[0].fields[map_index]
+
+        dim_defined = False
+        if 'axis' in kwargs:
+            dim_defined = True
+        if 'zAxis' in kwargs:
+            if kwargs['zAxis'] is not None:
+                kwargs['dim'] = 3
+            else:
+                kwargs['dim'] = 2
+            dim_defined = True
+        if 'dim' in kwargs:
+            dim_defined = True
+
+        if not dim_defined:
+            kwargs['dim'] = 2
+
+        return tfields.plotting.plotMesh(self, self.faces, **kwargs)
 
 
 if __name__ == '__main__':

tfields/plotting/__init__.py

+"""
+Core plotting tools for tfields library. Especially PlotOptions class
+is basis for many plotting expansions
+"""
+import warnings
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+import numpy as np
+from .mpl import *
+
+
+def setDefault(dictionary, attr, value):
+    """
+    Set defaults to a dictionary
+    """
+    if attr not in dictionary:
+        dictionary[attr] = value
+
+
+def gca(dim=None, **kwargs):
+    """
+    Forwarding to plt.gca but translating the dimension to projection
+    correct dimension
+    """
+    if dim == 3:
+        axis = plt.gca(projection='3d', **kwargs)
+    else:
+        axis = plt.gca(**kwargs)
+        if dim != axisDim(axis):
+            if dim is not None:
+                warnings.warn("You have another dimension set as gca."
+                              "I will force the new dimension to return.")
+                axis = plt.gcf().add_subplot(1, 1, 1, **kwargs)
+    return axis
+
+
+def axisDim(axis):
+    """
+    Returns int: axis dimension
+    """
+    if hasattr(axis, 'get_zlim'):
+        return 3
+    else:
+        return 2
+
+
+def setLabels(axis, *labels):
+    axis.set_xlabel(labels[0])
+    axis.set_ylabel(labels[1])
+    if axisDim(axis) == 3:
+        axis.set_zlabel(labels[2])
+
+
+def autoscale3D(axis, array=None, xLim=None, yLim=None, zLim=None):
+    if array is not None:
+        xMin, yMin, zMin = array.min(axis=0)
+        xMax, yMax, zMax = array.max(axis=0)
+        xLim = (xMin, xMax)
+        yLim = (yMin, yMax)
+        zLim = (zMin, zMax)
+    xLimAxis = axis.get_xlim()
+    yLimAxis = axis.get_ylim()
+    zLimAxis = axis.get_zlim()
+
+    if not False:
+        # not empty axis
+        xMin = min(xLimAxis[0], xLim[0])
+        yMin = min(yLimAxis[0], yLim[0])
+        zMin = min(zLimAxis[0], zLim[0])
+        xMax = max(xLimAxis[1], xLim[1])
+        yMax = max(yLimAxis[1], yLim[1])
+        zMax = max(zLimAxis[1], zLim[1])
+    axis.set_xlim([xMin, xMax])
+    axis.set_ylim([yMin, yMax])
+    axis.set_zlim([zMin, zMax])
+
+
+class PlotOptions(object):
+    """
+    processing kwargs for plotting functions and providing easy
+    access to axis, dimension and plotting method as well as indices
+    for array choice (x..., y..., zAxis)
+    """
+    def __init__(self, kwargs):
+        kwargs = dict(kwargs)
+        self.axis = kwargs.pop('axis', None)
+        self.dim = kwargs.pop('dim', None)
+        self.method = kwargs.pop('methodName', None)
+        self.setXYZAxis(kwargs)
+        self.plotKwargs = kwargs
+
+    @property
+    def method(self):
+        """
+        Method for plotting. Will be callable together with plotKwargs
+        """
+        return self._method
+
+    @method.setter
+    def method(self, methodName):
+        if not isinstance(methodName, str):
+            self._method = methodName
+        else:
+            self._method = getattr(self.axis, methodName)
+
+    @property
+    def dim(self):
+        """
+        axis dimension
+        """
+        return self._dim
+
+    @dim.setter
+    def dim(self, dim):
+        if dim is None:
+            if self._axis is None:
+                dim = 2
+            dim = axisDim(self._axis)
+        elif self._axis is not None:
+            if not dim == axisDim(self._axis):
+                raise ValueError("Axis and dim argument are in conflict.")
+        if dim not in [2, 3]:
+            raise NotImplementedError("Dimensions other than 2 or 3 are not supported.")
+        self._dim = dim
+
+    @property
+    def axis(self):
+        """
+        The plt.Axis object that belongs to this instance
+        """
+        if self._axis is None:
+            return gca(self._dim)
+        else:
+            return self._axis
+
+    @axis.setter
+    def axis(self, axis):
+        self._axis = axis
+
+    def setXYZAxis(self, kwargs):
+        self._xAxis = kwargs.pop('xAxis', 0)
+        self._yAxis = kwargs.pop('yAxis', 1)
+        zAxis = kwargs.pop('zAxis', None)
+        if zAxis is None and self.dim == 3:
+            indicesUsed = [0, 1, 2]
+            indicesUsed.remove(self._xAxis)
+            indicesUsed.remove(self._yAxis)
+            zAxis = indicesUsed[0]
+        self._zAxis = zAxis
+
+    def getXYZAxis(self):
+        return self._xAxis, self._yAxis, self._zAxis
+
+    def setVminVmaxAuto(self, vmin, vmax, scalars):
+        """
+        Automatically set vmin and vmax as min/max of scalars
+        but only if vmin or vmax is None
+        """
+        if scalars is None:
+            return
+        if len(scalars) < 2:
+            warnings.warn("Need at least two scalars to autoset vmin and/or vmax!")
+            return
+        if vmin is None:
+            vmin = min(scalars)
+            self.plotKwargs['vmin'] = vmin
+        if vmax is None:
+            vmax = min(scalars)
+            self.plotKwargs['vmax'] = vmax
+
+    def getNormArgs(self, vminDefault=0, vmaxDefault=1, cmapDefault=None):
+        if cmapDefault is None:
+            cmapDefault = plt.rcParams['image.cmap']
+        cmap = self.get('cmap', cmapDefault)
+        vmin = self.get('vmin', vminDefault)
+        vmax = self.get('vmax', vmaxDefault)
+        return cmap, vmin, vmax
+
+    def formatColors(self, colors, fmt='rgba', length=None):
+        """
+        format colors according to fmt argument
+        Args:
+            colors (list/one value of rgba tuples/int/float/str): This argument will
+                be interpreted as color
+            fmt (str): rgba / norm
+            length (int/None): if not None: corrct colors lenght
+    
+        Returns:
+            colors in fmt
+        """
+        hasIter = True
+        if not hasattr(colors, '__iter__'):
+            # colors is just one element
+            hasIter = False
+            colors = [colors]
+        if hasattr(colors[0], '__iter__') and fmt == 'norm':
+            # rgba given but norm wanted
+            cmap, vmin, vmax = self.getNormArgs(cmapDefault='NotSpecified',
+                                                vminDefault=None,
+                                                vmaxDefault=None)
+            colors = getColorsInverse(colors, cmap, vmin, vmax)
+            self.plotKwargs['vmin'] = vmin
+            self.plotKwargs['vmax'] = vmax
+            self.plotKwargs['cmap'] = cmap
+        elif fmt == 'rgba':
+            if isinstance(colors[0], str) or isinstance(colors[0], unicode):
+                # string color defined
+                colors = map(mpl.colors.to_rgba, colors)
+            else:
+                # norm given rgba wanted
+                cmap, vmin, vmax = self.getNormArgs(cmapDefault='NotSpecified',
+                                                    vminDefault=None,
+                                                    vmaxDefault=None)
+                self.setVminVmaxAuto(vmin, vmax, colors)
+                # update vmin and vmax
+                cmap, vmin, vmax = self.getNormArgs()
+                colors = getColors(colors,
+                                   vmin=vmin,
+                                   vmax=vmax,
+                                   cmap=cmap)
+    
+        if length is not None:
+            # just one colors value given
+            if len(colors) != length:
+                if not len(colors) == 1:
+                    raise ValueError("Can not correct color length")
+                colors = list(colors)
+                colors *= length
+        elif not hasIter:
+            colors = colors[0]
+    
+        colors = np.array(colors)
+        return colors
+
+    def delNormArgs(self):
+        self.plotKwargs.pop('vmin', None)
+        self.plotKwargs.pop('vmax', None)
+        self.plotKwargs.pop('cmap', None)
+
+    def getSortedLabels(self, labels):
+        """
+        Returns the labels corresponding to the axes
+        """
+        return [labels[i] for i in self.getXYZAxis() if i is not None]
+        
+    def get(self, attr, default=None):
+        return self.plotKwargs.get(attr, default)
+
+    def pop(self, attr, default=None):
+        return self.plotKwargs.pop(attr, default)
+
+    def set(self, attr, value):
+        self.plotKwargs[attr] = value
+
+    def setDefault(self, attr, value):
+        setDefault(self.plotKwargs, attr, value)
+
+    def retrieve(self, attr, default=None, keep=True):
+        if keep:
+            return self.get(attr, default)
+        else:
+            return self.pop(attr, default)
+
+    def retrieveChain(self, *args, **kwargs):
+        default = kwargs.pop('default', None)
+        keep = kwargs.pop('keep', True)
+        if len(args) > 1:
+            return self.retrieve(args[0],
+                                 self.retrieveChain(*args[1:],
+                                                    default=default,
+                                                    keep=keep),
+                                 keep=keep)
+        if len(args) != 1:
+            raise ValueError("Invalid number of args ({0})".format(len(args)))
+        return self.retrieve(args[0], default, keep=keep)

tfields/plotting/mpl.py

+import tfields
+import numpy as np
+import warnings
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from matplotlib.patches import Circle
+import mpl_toolkits.mplot3d as plt3D
+
+
+
+
+def plotArray(array, **kwargs):
+    """
+    Points3D plotting method.
+
+    Args:
+        axis (matplotlib.Axis) object
+        xAxis (int): coordinate index that should be on xAxis
+        yAxis (int): coordinate index that should be on yAxis
+        zAxis (int or None): coordinate index that should be on zAxis.
+            If it evaluates to None, 2D plot will be done.
+        methodName (str): method name to use for filling the axis
+
+    Returns:
+        Artist or list of Artists (imitating the axis.scatter/plot behaviour).
+        Better Artist not list of Artists
+    """
+    tfields.plotting.setDefault(kwargs, 'methodName', 'scatter')
+    po = tfields.plotting.PlotOptions(kwargs)
+
+    labelList = po.pop('labelList', ['x (m)', 'y (m)', 'z (m)'])
+    xAxis, yAxis, zAxis = po.getXYZAxis()
+    tfields.plotting.setLabels(po.axis, *po.getSortedLabels(labelList))
+    if zAxis is None:
+        args = [array[:, xAxis],
+                array[:, yAxis]]
+    else:
+        args = [array[:, xAxis],
+                array[:, yAxis],
+                array[:, zAxis]]
+    artist = po.method(*args,
+                       **po.plotKwargs)
+    return artist
+
+
+def plotMesh(vertices, faces, **kwargs):
+    """
+    Args:
+        axis (matplotlib axis)
+        xAxis (int)
+        yAxis (int)
+        zAxis (int)
+        edgecolor (color)
+        color (color): if given, use this color for faces in 2D
+        cmap
+        vmin
+        vmax
+    """
+    if faces.shape[0] == 0:
+        warnings.warn("No faces to plot")
+        return None
+    if max(faces.flat) > vertices.shape[0]:
+        raise ValueError("Some faces point to non existing vertices.")
+    po = tfields.plotting.PlotOptions(kwargs)
+    if po.dim == 2:
+        full = True
+        import npTools as npt
+        import pyTools
+        mesh = npt.Mesh3D(vertices, faces=faces)
+        xAxis, yAxis, zAxis = po.getXYZAxis()
+        facecolors = po.retrieveChain('facecolors', 'color',
+                                      default=0,
+                                      keep=False)
+        if full:
+            # implementation that will sort the triangles by zAxis
+            centroids = mesh.getCentroids()
+            axesIndices = [0, 1, 2]
+            axesIndices.pop(axesIndices.index(xAxis))
+            axesIndices.pop(axesIndices.index(yAxis))
+            zAxis = axesIndices[0]
+            zs = centroids[:, zAxis]
+            zs, faces, facecolors = pyTools.array.getSortedBoth(zs, faces,
+                                                                facecolors)
+            nFacesInitial = len(faces)
+        else:
+            # cut away "back sides" implementation
+            directionVector = np.array([1., 1., 1.])
+            directionVector[xAxis] = 0.
+            directionVector[yAxis] = 0.
+            normVectors = mesh.triangles.getNormVectors()
+            dotProduct = np.dot(normVectors, directionVector)
+            nFacesInitial = len(faces)
+            faces = faces[dotProduct > 0]
+
+        vertices = mesh
+
+        po.plotKwargs['methodName'] = 'tripcolor'
+        po.plotKwargs['triangles'] = faces
+
+        """
+        sort out color arguments
+        """
+        facecolors = po.formatColors(facecolors,
+                                     fmt='norm',
+                                     length=nFacesInitial)
+        if not full:
+            facecolors = facecolors[dotProduct > 0]
+        po.plotKwargs['facecolors'] = facecolors
+
+        d = po.plotKwargs
+        d['xAxis'] = xAxis
+        d['yAxis'] = yAxis
+        artist = plotArray(vertices, **d)
+    elif po.dim == 3:
+        label = po.pop('label', None)
+        color = po.retrieveChain('color', 'c', 'facecolors',
+                                 default='grey',
+                                 keep=False)
+        color = po.formatColors(color,
+                                fmt='rgba',
+                                length=len(faces))
+        nanMask = np.isnan(color)
+        if nanMask.any():
+            warnings.warn("nan found in colors. Removing the corresponding faces!")
+            color = color[~nanMask]
+            faces = faces[~nanMask]
+
+        edgecolor = po.pop('edgecolor', None)
+        alpha = po.pop('alpha', None)
+        po.delNormArgs()
+
+        triangles = np.array([vertices[face] for face in faces])
+        artist = plt3D.art3d.Poly3DCollection(triangles, **po.plotKwargs)
+        po.axis.add_collection3d(artist)
+
+        if edgecolor is not None:
+            artist.set_edgecolor(edgecolor)
+            artist.set_facecolors(color)
+        else:
+            artist.set_color(color)
+
+        if alpha is not None:
+            artist.set_alpha(alpha)
+
+        # for some reason auto-scale does not work
+        tfields.plotting.autoscale3D(po.axis, array=vertices)
+
+        # legend lables do not work at all as an argument
+        if label:
+            artist.set_label(label)
+
+        # when plotting the legend edgecolors/facecolors2d are needed
+        artist._edgecolors2d = None
+        artist._facecolors2d = None
+
+        labelList = ['x (m)', 'y (m)', 'z (m)']
+        tfields.plotting.setLabels(po.axis, *po.getSortedLabels(labelList))
+
+    return artist
+
+
+def plotVectorField(points, vectors, **kwargs):
+    """
+    Args:
+        points (array_like): base vectors
+        vectors (array_like): direction vectors
+    """
+    po = tfields.plotting.PlotOptions(kwargs)
+    if points is None:
+        points = np.full(vectors.shape, 0.)
+    artists = []
+    xAxis, yAxis, zAxis = po.getXYZAxis()
+    for point, vector in zip(points, vectors):
+        if po.dim == 3:
+            artists.append(po.axis.quiver(point[xAxis], point[yAxis], point[zAxis],
+                                          vector[xAxis], vector[yAxis], vector[zAxis],
+                                          **po.plotKwargs))
+        else:
+            artists.append(po.axis.quiver(point[xAxis], point[yAxis],
+                                          vector[xAxis], vector[yAxis],
+                                          **po.plotKwargs))
+    return artists
+
+
+def plotPlane(point, normal, **kwargs):
+
+    def plot_vector(fig, orig, v, color='blue'):
+        axis = fig.gca(projection='3d')
+        orig = np.array(orig)
+        v = np.array(v)
+        axis.quiver(orig[0], orig[1], orig[2], v[0], v[1], v[2], color=color)
+        axis.set_xlim(0, 10)
+        axis.set_ylim(0, 10)
+        axis.set_zlim(0, 10)
+        axis = fig.gca(projection='3d')
+        return fig
+
+    def rotation_matrix(d):
+        sin_angle = np.linalg.norm(d)
+        if sin_angle == 0:
+            return np.identity(3)
+        d /= sin_angle
+        eye = np.eye(3)
+        ddt = np.outer(d, d)
+        skew = np.array([[0, d[2], -d[1]],
+                         [-d[2], 0, d[0]],
+                         [d[1], -d[0], 0]],
+                        dtype=np.float64)
+
+        M = ddt + np.sqrt(1 - sin_angle**2) * (eye - ddt) + sin_angle * skew
+        return M
+
+    def pathpatch_2d_to_3d(pathpatch, z, normal):
+        if type(normal) is str:  # Translate strings to normal vectors
+            index = "xyz".index(normal)
+            normal = np.roll((1.0, 0, 0), index)
+
+        normal /= np.linalg.norm(normal)  # Make sure the vector is normalised
+        path = pathpatch.get_path()  # Get the path and the associated transform
+        trans = pathpatch.get_patch_transform()
+
+        path = trans.transform_path(path)  # Apply the transform
+
+        pathpatch.__class__ = plt3D.art3d.PathPatch3D  # Change the class
+        pathpatch._code3d = path.codes  # Copy the codes
+        pathpatch._facecolor3d = pathpatch.get_facecolor  # Get the face color
+
+        verts = path.vertices  # Get the vertices in 2D
+
+        d = np.cross(normal, (0, 0, 1))  # Obtain the rotation vector
+        M = rotation_matrix(d)  # Get the rotation matrix
+
+        pathpatch._segment3d = np.array([np.dot(M, (x, y, 0)) + (0, 0, z) for x, y in verts])
+
+    def pathpatch_translate(pathpatch, delta):
+        pathpatch._segment3d += delta
+
+    kwargs['alpha'] = kwargs.pop('alpha', 0.5)
+    po = tfields.plotting.PlotOptions(kwargs)
+    patch = Circle((0, 0), **po.plotKwargs)
+    po.axis.add_patch(patch)
+    pathpatch_2d_to_3d(patch, z=0, normal=normal)
+    pathpatch_translate(patch, (point[0], point[1], point[2]))
+
+
+def plotSphere(point, radius, **kwargs):
+    po = tfields.plotting.PlotOptions(kwargs)
+    # Make data
+    u = np.linspace(0, 2 * np.pi, 100)
+    v = np.linspace(0, np.pi, 100)
+    x = point[0] + radius * np.outer(np.cos(u), np.sin(v))
+    y = point[1] + radius * np.outer(np.sin(u), np.sin(v))
+    z = point[2] + radius * np.outer(np.ones(np.size(u)), np.cos(v))
+
+    # Plot the surface
+    return po.axis.plot_surface(x, y, z, **po.plotKwargs)
+
+
+"""
+Color section
+"""
+def getColors(scalars, cmap=None, vmin=None, vmax=None):
+    """
+    retrieve the colors for a list of scalars
+    """
+    if not hasattr(scalars, '__iter__'):
+        scalars = [scalars]
+    if vmin is None:
+        vmin = min(scalars)
+    if vmax is None:
+        vmax = max(scalars)
+    colorMap = plt.get_cmap(cmap)
+    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
+    return colorMap(map(norm, scalars))
+
+
+def getColorsInverse(colors, cmap, vmin, vmax):
+    """
+    Reconstruct the numeric values (0 - 1) of given
+    Args:
+        colors (list or rgba tuple)