Revision: 7206
          http://matplotlib.svn.sourceforge.net/matplotlib/?rev=7206&view=rev
Author:   heeres
Date:     2009-06-09 22:40:35 +0000 (Tue, 09 Jun 2009)

Log Message:
-----------
mplot3d: make Poly3DCollection ScalarMappable, support in Axes3D.plot_surface.
Fix-up and extend contours. Add examples.

Modified Paths:
--------------
    trunk/matplotlib/doc/mpl_toolkits/mplot3d/tutorial.rst
    trunk/matplotlib/examples/mplot3d/surface3d_demo.py
    trunk/matplotlib/lib/mpl_toolkits/mplot3d/art3d.py
    trunk/matplotlib/lib/mpl_toolkits/mplot3d/axes3d.py

Added Paths:
-----------
    trunk/matplotlib/examples/mplot3d/contour3d_demo2.py
    trunk/matplotlib/examples/mplot3d/surface3d_demo2.py

Modified: trunk/matplotlib/doc/mpl_toolkits/mplot3d/tutorial.rst
===================================================================
--- trunk/matplotlib/doc/mpl_toolkits/mplot3d/tutorial.rst      2009-06-09 
16:47:46 UTC (rev 7205)
+++ trunk/matplotlib/doc/mpl_toolkits/mplot3d/tutorial.rst      2009-06-09 
22:40:35 UTC (rev 7206)
@@ -38,12 +38,14 @@
 .. automethod:: Axes3D.plot_surface
 
 .. plot:: mpl_examples/mplot3d/surface3d_demo.py
+.. plot:: mpl_examples/mplot3d/surface3d_demo2.py
 
 Contour plots
 =============
 .. automethod:: Axes3D.contour
 
 .. plot:: mpl_examples/mplot3d/contour3d_demo.py
+.. plot:: mpl_examples/mplot3d/contour3d_demo2.py
 
 Filled contour plots
 ====================

Added: trunk/matplotlib/examples/mplot3d/contour3d_demo2.py
===================================================================
--- trunk/matplotlib/examples/mplot3d/contour3d_demo2.py                        
        (rev 0)
+++ trunk/matplotlib/examples/mplot3d/contour3d_demo2.py        2009-06-09 
22:40:35 UTC (rev 7206)
@@ -0,0 +1,12 @@
+from mpl_toolkits.mplot3d import axes3d
+import pylab
+import random
+
+fig = pylab.figure()
+ax = axes3d.Axes3D(fig)
+X, Y, Z = axes3d.get_test_data(0.05)
+cset = ax.contour(X, Y, Z, 16, extend3d=True)
+ax.clabel(cset, fontsize=9, inline=1)
+
+pylab.show()
+

Modified: trunk/matplotlib/examples/mplot3d/surface3d_demo.py
===================================================================
--- trunk/matplotlib/examples/mplot3d/surface3d_demo.py 2009-06-09 16:47:46 UTC 
(rev 7205)
+++ trunk/matplotlib/examples/mplot3d/surface3d_demo.py 2009-06-09 22:40:35 UTC 
(rev 7206)
@@ -1,16 +1,17 @@
 from mpl_toolkits.mplot3d import Axes3D
+from matplotlib import cm
 import pylab
 import random
 import numpy as np
 
 fig = pylab.figure()
 ax = Axes3D(fig)
-X = np.arange(-5, 5, 0.5)
-Y = np.arange(-5, 5, 0.5)
+X = np.arange(-5, 5, 0.25)
+Y = np.arange(-5, 5, 0.25)
 X, Y = np.meshgrid(X, Y)
 R = np.sqrt(X**2 + Y**2)
 Z = np.sin(R)
-ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='forestgreen')
+ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet)
 
 pylab.show()
 

Added: trunk/matplotlib/examples/mplot3d/surface3d_demo2.py
===================================================================
--- trunk/matplotlib/examples/mplot3d/surface3d_demo2.py                        
        (rev 0)
+++ trunk/matplotlib/examples/mplot3d/surface3d_demo2.py        2009-06-09 
22:40:35 UTC (rev 7206)
@@ -0,0 +1,18 @@
+from mpl_toolkits.mplot3d import Axes3D
+import pylab
+import random
+import numpy as np
+
+fig = pylab.figure()
+ax = Axes3D(fig)
+
+u = np.linspace(0, 2 * np.pi, 100)
+v = np.linspace(0, np.pi, 100)
+
+x = 10 * np.outer(np.cos(u), np.sin(v))
+y = 10 * np.outer(np.sin(u), np.sin(v))
+z = 10 * np.outer(np.ones(np.size(u)), np.cos(v))
+ax.plot_surface(x, y, z,  rstride=4, cstride=4, color='b')
+
+pylab.show()
+

Modified: trunk/matplotlib/lib/mpl_toolkits/mplot3d/art3d.py
===================================================================
--- trunk/matplotlib/lib/mpl_toolkits/mplot3d/art3d.py  2009-06-09 16:47:46 UTC 
(rev 7205)
+++ trunk/matplotlib/lib/mpl_toolkits/mplot3d/art3d.py  2009-06-09 22:40:35 UTC 
(rev 7206)
@@ -10,6 +10,7 @@
 from matplotlib import lines, text as mtext, path as mpath, colors as mcolors
 from matplotlib.collections import Collection, LineCollection, \
         PolyCollection, PatchCollection
+from matplotlib.cm import ScalarMappable
 from matplotlib.patches import Patch
 from matplotlib.colors import Normalize
 from matplotlib.cbook import iterable
@@ -111,26 +112,30 @@
     line.__class__ = Line3D
     line.set_3d_properties(zs, zdir)
 
-def path_to_3d_segment(path, z=0, zdir='z'):
+def path_to_3d_segment(path, zs=0, zdir='z'):
     '''Convert a path to a 3D segment.'''
+
+    if not iterable(zs):
+        zs = [zs] * len(path)
+
     seg = []
-    for (pathseg, code) in path.iter_segments():
-        seg.append(pathseg)
-    seg3d = [juggle_axes(x, y, z, zdir) for (x, y) in seg]
+    pathsegs = path.iter_segments(simplify=False, curves=False)
+    for (((x, y), code), z) in zip(pathsegs, zs):
+        seg.append((x, y, z))
+    seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg]
     return seg3d
 
 def paths_to_3d_segments(paths, zs=0, zdir='z'):
-    '''Convert paths from a collection object to 3D segments.'''
+    '''
+    Convert paths from a collection object to 3D segments.
+    '''
 
-    try:
-        zs = float(zs)
+    if not iterable(zs):
         zs = [zs] * len(paths)
-    except:
-        pass
 
     segments = []
-    for path, z in zip(paths, zs):
-        segments.append(path_to_3d_segment(path, z, zdir))
+    for path, pathz in zip(paths, zs):
+        segments.append(path_to_3d_segment(path, pathz, zdir))
     return segments
 
 class Line3DCollection(LineCollection):
@@ -255,8 +260,17 @@
     '''
 
     def __init__(self, verts, *args, **kwargs):
+        '''
+        Create a Poly3DCollection.
+
+        *verts* should contain 3D coordinates.
+
+        Note that this class does a bit of magic with the _facecolors
+        and _edgecolors properties.
+        '''
+
         PolyCollection.__init__(self, verts, *args, **kwargs)
-        self.set_3d_properties()
+        self._zsort = 1
 
     def get_vector(self, segments3d):
         """Optimize points for projection"""
@@ -276,6 +290,7 @@
         self._sort_zpos = min(zs)
 
     def set_verts(self, verts, closed=True):
+        '''Set 3D vertices.'''
         self.get_vector(verts)
         # 2D verts will be updated at draw time
         PolyCollection.set_verts(self, [], closed)
@@ -283,40 +298,73 @@
     def set_3d_properties(self):
         self._zsort = 1
         self._facecolors3d = PolyCollection.get_facecolors(self)
-        self._edgecolors3d = self.get_edgecolors()
+        self._edgecolors3d = PolyCollection.get_edgecolors(self)
 
     def do_3d_projection(self, renderer):
+        '''
+        Perform the 3D projection for this object.
+        '''
+
+        if self._A is not None:
+            self.update_scalarmappable()
+            self._facecolors3d = self._facecolors
+
         txs, tys, tzs = proj3d.proj_transform_vec(self._vec, renderer.M)
         xyzlist = [(txs[si:ei], tys[si:ei], tzs[si:ei]) \
                 for si, ei in self._segis]
-        colors = self._facecolors3d
 
+        # This extra fuss is to re-order face / edge colors
+        cface = self._facecolors3d
+        if len(self._edgecolors3d) != len(cface):
+            cedge = cface
+        else:
+            cedge = self._edgecolors3d
+
         # if required sort by depth (furthest drawn first)
         if self._zsort:
-            z_segments_2d = [(min(zs), zip(xs, ys), c) for
-                             (xs, ys, zs), c in zip(xyzlist, colors)]
+            z_segments_2d = [(min(zs), zip(xs, ys), fc, ec) for
+                    (xs, ys, zs), fc, ec in zip(xyzlist, cface, cedge)]
             z_segments_2d.sort()
             z_segments_2d.reverse()
         else:
             raise ValueError, "whoops"
-        segments_2d = [s for z, s, c in z_segments_2d]
-        colors = [c for z, s, c in z_segments_2d]
+
+        segments_2d = [s for z, s, fc, ec in z_segments_2d]
         PolyCollection.set_verts(self, segments_2d)
-        self._facecolors2d = colors
 
+        self._facecolors2d = [fc for z, s, fc, ec in z_segments_2d]
+        if len(self._edgecolors3d) == len(cface):
+            self._edgecolors2d = [ec for z, s, fc, ec in z_segments_2d]
+        else:
+            self._edgecolors2d = self._edgecolors3d
+
         # Return zorder value
         zvec = np.array([[0], [0], [self._sort_zpos], [1]])
         ztrans = proj3d.proj_transform_vec(zvec, renderer.M)
         return ztrans[2][0]
 
+    def set_facecolor(self, colors):
+        PolyCollection.set_facecolor(self, colors)
+        self._facecolors3d = PolyCollection.get_facecolor(self)
+    set_facecolors = set_facecolor
+
+    def set_edgecolor(self, colors):
+        PolyCollection.set_edgecolor(self, colors)
+        self._edgecolors3d = PolyCollection.get_edgecolor(self)
+    set_edgecolors = set_edgecolor
+
     def get_facecolors(self):
         return self._facecolors2d
     get_facecolor = get_facecolors
 
+    def get_edgecolors(self):
+        return self._edgecolors2d
+    get_edgecolor = get_edgecolors
+
     def draw(self, renderer):
         return Collection.draw(self, renderer)
 
-def poly_collection_2d_to_3d(col, zs=None, zdir='z'):
+def poly_collection_2d_to_3d(col, zs=0, zdir='z'):
     """Convert a PolyCollection to a Poly3DCollection object."""
     segments_3d = paths_to_3d_segments(col.get_paths(), zs, zdir)
     col.__class__ = Poly3DCollection

Modified: trunk/matplotlib/lib/mpl_toolkits/mplot3d/axes3d.py
===================================================================
--- trunk/matplotlib/lib/mpl_toolkits/mplot3d/axes3d.py 2009-06-09 16:47:46 UTC 
(rev 7205)
+++ trunk/matplotlib/lib/mpl_toolkits/mplot3d/axes3d.py 2009-06-09 22:40:35 UTC 
(rev 7206)
@@ -541,6 +541,10 @@
     def plot_surface(self, X, Y, Z, *args, **kwargs):
         '''
         Create a surface plot.
+
+        By default it will be colored in shades of a solid color,
+        but it also supports color mapping by supplying the *cmap*
+        argument.
         
         ==========  ================================================
         Argument    Description
@@ -550,6 +554,7 @@
         *rstride*   Array row stride (step size)
         *cstride*   Array column stride (step size)
         *color*     Color of the surface patches
+        *cmap*      A colormap for the surface patches.
         ==========  ================================================
         '''
 
@@ -562,9 +567,11 @@
 
         color = kwargs.pop('color', 'b')
         color = np.array(colorConverter.to_rgba(color))
+        cmap = kwargs.get('cmap', None)
 
         polys = []
-        boxes = []
+        normals = []
+        avgz = []
         for rs in np.arange(0, rows-1, rstride):
             for cs in np.arange(0, cols-1, cstride):
                 ps = []
@@ -579,31 +586,53 @@
                     corners.append([ztop[0], ztop[-1], zbase[0], zbase[-1]])
                     z = np.concatenate((ztop, zleft, zbase, zright))
                     ps.append(z)
-                boxes.append(map(np.array, zip(*corners)))
-                polys.append(zip(*ps))
 
-        lines = []
+                # The construction leaves the array with duplicate points, 
which
+                # are removed here.
+                ps = zip(*ps)
+                lastp = np.array([])
+                ps2 = []
+                avgzsum = 0.0
+                for p in ps:
+                    if p != lastp:
+                        ps2.append(p)
+                        lastp = p
+                        avgzsum += p[2]
+                polys.append(ps2)
+                avgz.append(avgzsum / len(ps2))
+
+                v1 = np.array(ps2[0]) - np.array(ps2[1])
+                v2 = np.array(ps2[2]) - np.array(ps2[0])
+                normals.append(np.cross(v1, v2))
+
+        polyc = art3d.Poly3DCollection(polys, *args, **kwargs)
+        if cmap is not None:
+            polyc.set_array(np.array(avgz))
+            polyc.set_linewidth(0)
+        else:
+            colors = self._shade_colors(color, normals)
+            polyc.set_facecolors(colors)
+
+        self.add_collection(polyc)
+        self.auto_scale_xyz(X, Y, Z, had_data)
+
+        return polyc
+
+    def _shade_colors(self, color, normals):
         shade = []
-        for box in boxes:
-            n = proj3d.cross(box[0]-box[1],
-                         box[0]-box[2])
-            n = n/proj3d.mod(n)*5
+        for n in normals:
+            n = n / proj3d.mod(n) * 5
             shade.append(np.dot(n, [-1, -1, 0.5]))
-            lines.append((box[0], n+box[0]))
 
         shade = np.array(shade)
         mask = ~np.isnan(shade)
         norm = Normalize(min(shade[mask]), max(shade[mask]))
 
-        colors = [color * (0.5+norm(v)*0.5) for v in shade]
-        for c in colors: c[3] = 1
-        polyc = art3d.Poly3DCollection(polys, facecolors=colors, \
-                *args, **kwargs)
-        polyc._zsort = 1
-        self.add_collection(polyc)
+        color = color.copy()
+        color[3] = 1
+        colors = [color * (0.5 + norm(v) * 0.5) for v in shade]
 
-        self.auto_scale_xyz(X, Y, Z, had_data)
-        return polyc
+        return colors
 
     def plot_wireframe(self, X, Y, Z, *args, **kwargs):
         '''
@@ -653,21 +682,78 @@
 
         return linec
 
-    def contour(self, X, Y, Z, *args, **kwargs):
+    def _3d_extend_contour(self, cset, stride=5):
         '''
+        Extend a contour in 3D by creating 
+        '''
+
+        levels = cset.levels
+        colls = cset.collections
+        dz = (levels[1] - levels[0]) / 2
+
+        for z, linec in zip(levels, colls):
+            topverts = art3d.paths_to_3d_segments(linec.get_paths(), z - dz)
+            botverts = art3d.paths_to_3d_segments(linec.get_paths(), z + dz)
+
+            color = linec.get_color()[0]
+
+            polyverts = []
+            normals = []
+            nsteps = round(len(topverts[0]) / stride)
+            stepsize = (len(topverts[0]) - 1) / (nsteps - 1)
+            for i in range(int(round(nsteps)) - 1):
+                i1 = int(round(i * stepsize))
+                i2 = int(round((i + 1) * stepsize))
+                polyverts.append([topverts[0][i1],
+                    topverts[0][i2],
+                    botverts[0][i2],
+                    botverts[0][i1]])
+
+                v1 = np.array(topverts[0][i1]) - np.array(topverts[0][i2])
+                v2 = np.array(topverts[0][i1]) - np.array(botverts[0][i1])
+                normals.append(np.cross(v1, v2))
+
+            colors = self._shade_colors(color, normals)
+            colors2 = self._shade_colors(color, normals)
+            polycol = art3d.Poly3DCollection(polyverts, facecolors=colors,
+                    edgecolors=colors2)
+            self.add_collection3d(polycol)
+
+        for col in colls:
+            self.collections.remove(col)
+
+    def contour(self, X, Y, Z, levels=10, **kwargs):
+        '''
         Create a 3D contour plot.
 
-        *X*, *Y*, *Z*: data
+        ==========  ================================================
+        Argument    Description
+        ==========  ================================================
+        *X*, *Y*,   Data values as numpy.arrays
+        *Z*
+        *levels*    Number of levels to use, defaults to 10. Can
+                    also be a tuple of specific levels.
+        *extend3d*  Whether to extend contour in 3D (default: False)
+        *stride*    Stride (step size) for extending contour
+        ==========  ================================================
 
-        Keyword arguments are passed on to
+        Other keyword arguments are passed on to
         :func:`~matplotlib.axes.Axes.contour`
         '''
 
+        extend3d = kwargs.pop('extend3d', False)
+        stride = kwargs.pop('stride', 5)
+        nlevels = kwargs.pop('nlevels', 15)
+
         had_data = self.has_data()
-        cset = Axes.contour(self, X, Y, Z, *args, **kwargs)
-        for z, linec in zip(cset.levels, cset.collections):
-            art3d.line_collection_2d_to_3d(linec, z)
+        cset = Axes.contour(self, X, Y, Z, levels, **kwargs)
 
+        if extend3d:
+            self._3d_extend_contour(cset, stride)
+        else:
+            for z, linec in zip(cset.levels, cset.collections):
+                art3d.line_collection_2d_to_3d(linec, z)
+
         self.auto_scale_xyz(X, Y, Z, had_data)
         return cset
 
@@ -688,11 +774,8 @@
         cset = Axes.contourf(self, X, Y, Z, *args, **kwargs)
         levels = cset.levels
         colls = cset.collections
-
         for z1, z2, linec in zip(levels, levels[1:], colls):
-            zs = [z1] * (len(linec.get_paths()[0]) / 2)
-            zs += [z2] * (len(linec.get_paths()[0]) / 2)
-            art3d.poly_collection_2d_to_3d(linec, zs)
+            art3d.poly_collection_2d_to_3d(linec, z1)
 
         self.auto_scale_xyz(X, Y, Z, had_data)
         return cset


This was sent by the SourceForge.net collaborative development platform, the 
world's largest Open Source development site.

------------------------------------------------------------------------------
Crystal Reports - New Free Runtime and 30 Day Trial
Check out the new simplified licensing option that enables unlimited
royalty-free distribution of the report engine for externally facing 
server and web deployment.
http://p.sf.net/sfu/businessobjects
_______________________________________________
Matplotlib-checkins mailing list
Matplotlib-checkins@lists.sourceforge.net
https://lists.sourceforge.net/lists/listinfo/matplotlib-checkins

Reply via email to