ENH: Added depthshade argument for axes3d scatter

WeatherGod · WeatherGod · commit 923272f46082 · 2014-07-11T12:13:08.000-04:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -311,9 +311,14 @@ def __init__(self, *args, **kwargs):
         :class:`~matplotlib.collections.PatchCollection`. 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)
         PatchCollection.__init__(self, *args, **kwargs)
         self._old_draw = lambda x: PatchCollection.draw(self, x)
         self.set_3d_properties(zs, zdir)
@@ -340,14 +345,16 @@ def set_3d_properties(self, zs, zdir):
     def do_3d_projection(self, renderer):
         xs, ys, zs = self._offsets3d
         vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
-        #FIXME: mpl allows us no way to unset the collection alpha value
-        #self._alpha = None
-        fcs = mcolors.colorConverter.to_rgba_array(self._facecolor3d,
-                                                   self._alpha)
-        mcs = mcolors.colorConverter.to_rgba_array(self._edgecolor3d,
-                                                   self._alpha)
-        self.set_facecolors(zalpha(fcs, vzs))
-        self.set_edgecolors(zalpha(mcs, vzs))
+
+        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)))
 
         if vzs.size > 0 :
@@ -367,9 +374,22 @@ class Path3DCollection(Collection3D, PathCollection):
     pass
 
 
-def patch_collection_2d_to_3d(col, zs=0, zdir='z'):
-    """Convert a PatchCollection to a Patch3DCollection object."""
+def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
+    """
+    Convert a :class:`~matplotlib.collections.PatchCollection` into a
+    :class:`Patch3DCollection` object
+    (or a :class:`~matplotlib.collections.PathCollection` into a
+    :class:`Path3DCollection` object).
+
+    Keywords:
+    *za*            The location or locations to place the patches in the
+                    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*.
+    
 
+    """
     # 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
@@ -378,6 +398,7 @@ def patch_collection_2d_to_3d(col, zs=0, zdir='z'):
         col.__class__ = Path3DCollection
     elif isinstance(col, PatchCollection):
         col.__class__ = Patch3DCollection
+    col._depthshade = depthshade
     col.set_3d_properties(zs, zdir)
 
 class Poly3DCollection(PolyCollection):
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2176,7 +2176,8 @@ def add_collection3d(self, col, zs=0, zdir='z'):
 
         Axes.add_collection(self, col)
 
-    def scatter(self, xs, ys, zs=0, zdir='z', s=20, c='b', *args, **kwargs):
+    def scatter(self, xs, ys, zs=0, zdir='z', s=20, c='b', depthshade=True,
+                *args, **kwargs):
         '''
         Create a scatter plot.
 
@@ -2200,6 +2201,10 @@ def scatter(self, xs, ys, zs=0, zdir='z', s=20, c='b', *args, **kwargs):
                     sequence because that is indistinguishable from an array
                     of values to be colormapped.  *c* can be a 2-D array in
                     which the rows are RGB or RGBA, however.
+
+        *depthshade*
+                    Whether or not to shade the scatter markers to give the
+                    appearance of depth. Default is *True*.
         ==========  ==========================================================
 
         Keyword arguments are passed on to
@@ -2238,7 +2243,8 @@ def scatter(self, xs, ys, zs=0, zdir='z', s=20, c='b', *args, **kwargs):
             zs = np.ones(len(xs)) * zs
         else:
             is_2d = False
-        art3d.patch_collection_2d_to_3d(patches, zs=zs, zdir=zdir)
+        art3d.patch_collection_2d_to_3d(patches, zs=zs, zdir=zdir,
+                                        depthshade=depthshade)
 
         if self._zmargin < 0.05 and xs.size > 0:
             self.set_zmargin(0.05)
