Merge pull request matplotlib#3237 from WeatherGod/mplot3d/3dcollection_fix

tacaswell · tacaswell · commit eb480314b9dd · 2014-07-12T12:01:31.000-04:00
Fix Collection3D. Fixes legend for scatter3d
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -295,7 +295,7 @@ def pathpatch_2d_to_3d(pathpatch, z=0, zdir='z'):
     pathpatch.__class__ = PathPatch3D
     pathpatch.set_3d_properties(mpath, z, zdir)
 
-class Collection3D(object):
+class Patch3DCollection(PatchCollection):
     '''
     A collection of 3D patches.
     '''
@@ -319,8 +319,7 @@ def __init__(self, *args, **kwargs):
         zs = kwargs.pop('zs', 0)
         zdir = kwargs.pop('zdir', 'z')
         self._depthshade = kwargs.pop('depthshade', True)
-        PatchCollection.__init__(self, *args, **kwargs)
-        self._old_draw = lambda x: PatchCollection.draw(self, x)
+        super(self.__class__, self).__init__(*args, **kwargs)
         self.set_3d_properties(zs, zdir)
 
 
@@ -362,16 +361,72 @@ def do_3d_projection(self, renderer):
         else :
             return np.nan
 
-    def draw(self, renderer):
-        self._old_draw(renderer)
 
+class Path3DCollection(PathCollection):
+    '''
+    A collection of 3D paths.
+    '''
+
+    def __init__(self, *args, **kwargs):
+        """
+        Create a collection of flat 3D paths with its normal vector
+        pointed in *zdir* direction, and located at *zs* on the *zdir*
+        axis. 'zs' can be a scalar or an array-like of the same length as
+        the number of paths in the collection.
+
+        Constructor arguments are the same as for
+        :class:`~matplotlib.collections.PathCollection`. In addition,
+        keywords *zs=0* and *zdir='z'* are available.
+
+        Also, the keyword argument "depthshade" is available to
+        indicate whether or not to shade the patches in order to
+        give the appearance of depth (default is *True*).
+        This is typically desired in scatter plots.
+        """
+        zs = kwargs.pop('zs', 0)
+        zdir = kwargs.pop('zdir', 'z')
+        self._depthshade = kwargs.pop('depthshade', True)
+        super(self.__class__, self).__init__(*args, **kwargs)
+        self.set_3d_properties(zs, zdir)
+
+
+    def set_sort_zpos(self,val):
+        '''Set the position to use for z-sorting.'''
+        self._sort_zpos = val
+
+    def set_3d_properties(self, zs, zdir):
+        # Force the collection to initialize the face and edgecolors
+        # just in case it is a scalarmappable with a colormap.
+        self.update_scalarmappable()
+        offsets = self.get_offsets()
+        if len(offsets) > 0:
+            xs, ys = list(zip(*offsets))
+        else:
+            xs = []
+            ys = []
+        self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
+        self._facecolor3d = self.get_facecolor()
+        self._edgecolor3d = self.get_edgecolor()
+
+    def do_3d_projection(self, renderer):
+        xs, ys, zs = self._offsets3d
+        vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
 
-class Patch3DCollection(Collection3D, PatchCollection):
-    pass
+        fcs = (zalpha(self._facecolor3d, vzs) if self._depthshade else
+               self._facecolor3d)
+        fcs = mcolors.colorConverter.to_rgba_array(fcs, self._alpha)
+        self.set_facecolors(fcs)
 
+        ecs = (zalpha(self._edgecolor3d, vzs) if self._depthshade else
+               self._edgecolor3d)
+        ecs = mcolors.colorConverter.to_rgba_array(ecs, self._alpha)
+        self.set_edgecolors(ecs)
+        super(self.__class__, self).set_offsets(list(zip(vxs, vys)))
 
-class Path3DCollection(Collection3D, PathCollection):
-    pass
+        if vzs.size > 0 :
+            return min(vzs)
+        else :
+            return np.nan
 
 
 def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
@@ -386,14 +441,9 @@ def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
                     collection along the *zdir* axis. Defaults to 0.
     *zdir*          The axis in which to place the patches. Default is "z".
     *depthshade*    Whether to shade the patches to give a sense of depth.
-                    Defaults to *True*.
-    
+                    Defaults to *True*. 
 
     """
-    # The tricky part here is that there are several classes that are
-    # derived from PatchCollection. We need to use the right draw method.
-    col._old_draw = col.draw
-
     if isinstance(col, PathCollection):
         col.__class__ = Path3DCollection
     elif isinstance(col, PatchCollection):
