Added rcount/ccount to plot_surface(), providing an alternative to rstride/cstride

WeatherGod · WeatherGod · commit e99f5906b41d · 2016-11-10T11:09:18.000-05:00
diff --git a/examples/mplot3d/surface3d_demo.py b/examples/mplot3d/surface3d_demo.py
@@ -28,7 +28,7 @@
 Z = np.sin(R)
 
 # Plot the surface.
-surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
+surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                        linewidth=0, antialiased=False)
 
 # Customize the z axis.
diff --git a/examples/mplot3d/surface3d_demo2.py b/examples/mplot3d/surface3d_demo2.py
@@ -22,6 +22,6 @@
 z = 10 * np.outer(np.ones(np.size(u)), np.cos(v))
 
 # Plot the surface
-ax.plot_surface(x, y, z, rstride=4, cstride=4, color='b')
+ax.plot_surface(x, y, z, color='b')
 
 plt.show()
diff --git a/examples/mplot3d/surface3d_demo3.py b/examples/mplot3d/surface3d_demo3.py
@@ -34,8 +34,7 @@
         colors[x, y] = colortuple[(x + y) % len(colortuple)]
 
 # Plot the surface with face colors taken from the array we made.
-surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=colors,
-                       linewidth=0)
+surf = ax.plot_surface(X, Y, Z, facecolors=colors, linewidth=0)
 
 # Customize the z axis.
 ax.set_zlim(-1, 1)
diff --git a/examples/mplot3d/surface3d_radial_demo.py b/examples/mplot3d/surface3d_radial_demo.py
@@ -28,7 +28,7 @@
 X, Y = R*np.cos(P), R*np.sin(P)
 
 # Plot the surface.
-ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=plt.cm.YlGnBu_r)
+ax.plot_surface(X, Y, Z, cmap=plt.cm.YlGnBu_r)
 
 # Tweak the limits and add latex math labels.
 ax.set_zlim(0, 1)
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1554,14 +1554,23 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
         The `rstride` and `cstride` kwargs set the stride used to
         sample the input data to generate the graph.  If 1k by 1k
         arrays are passed in the default values for the strides will
-        result in a 100x100 grid being plotted.
+        result in a 100x100 grid being plotted. Defaults to 10.
+        Superceded by `rcount` and `ccount` if both the
+        stride and the count are specified as of v2.0.0.
+
+        The `rcount` and `ccount` kwargs are the new method for
+        sampling the input data to generate the graph, and
+        supercedes `rstride` and `cstride` (unless using the
+        'classic' style) if both are specified. Added in v2.0.0.
 
         ============= ================================================
         Argument      Description
         ============= ================================================
         *X*, *Y*, *Z* Data values as 2D arrays
-        *rstride*     Array row stride (step size), defaults to 10
-        *cstride*     Array column stride (step size), defaults to 10
+        *rstride*     Array row stride (step size)
+        *cstride*     Array column stride (step size)
+        *rcount*      Use at most this many rows, defaults to 50
+        *ccount*      Use at most this many columns, defaults to 50
         *color*       Color of the surface patches
         *cmap*        A colormap for the surface patches.
         *facecolors*  Face colors for the individual patches
@@ -1582,8 +1591,30 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
         X, Y, Z = np.broadcast_arrays(X, Y, Z)
         rows, cols = Z.shape
 
+        has_rstride = 'rstride' in kwargs
+        has_cstride = 'cstride' in kwargs
+        has_rcount = 'rcount' in kwargs
+        has_ccount = 'ccount' in kwargs
+
         rstride = kwargs.pop('rstride', 10)
         cstride = kwargs.pop('cstride', 10)
+        rcount = kwargs.pop('rcount', 50)
+        ccount = kwargs.pop('ccount', 50)
+
+        if rcParams['_internal.classic_mode']:
+            # Strides have priority over counts in classic mode.
+            if not has_rstride and has_rcount:
+                rstride = int(np.ceil(rows / rcount))
+            if not has_cstride and has_ccount:
+                cstride = int(np.ceil(cols / ccount))
+        else:
+            # If the count is provided then it has priority
+            # If neither count or stride is provided then use
+            # the default count.
+            if has_rcount or (not has_rstride and not has_rcount):
+                rstride = int(np.ceil(rows / rcount))
+            if has_ccount or (not has_cstride and not has_ccount):
+                cstride = int(np.ceil(cols / ccount))
 
         if 'facecolors' in kwargs:
             fcolors = kwargs.pop('facecolors')
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -105,7 +105,7 @@ def f(t):
     R = np.sqrt(X ** 2 + Y ** 2)
     Z = np.sin(R)
 
-    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
+    surf = ax.plot_surface(X, Y, Z, rstride=1, ccount=40,
                            linewidth=0, antialiased=False)
 
     ax.set_zlim3d(-1, 1)
@@ -141,7 +141,7 @@ def test_surface3d():
     X, Y = np.meshgrid(X, Y)
     R = np.sqrt(X ** 2 + Y ** 2)
     Z = np.sin(R)
-    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
+    surf = ax.plot_surface(X, Y, Z, rcount=40, cstride=1, cmap=cm.coolwarm,
                            lw=0, antialiased=False)
     ax.set_zlim(-1.01, 1.01)
     fig.colorbar(surf, shrink=0.5, aspect=5)
