Merge pull request #1074 from WeatherGod/mplot3d/input_broadcast2

WeatherGod · WeatherGod · commit 2798e474982e · 2012-08-16T20:13:40.000-07:00
Added broadcasting support in some mplot3d methods
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1352,15 +1352,9 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
         had_data = self.has_data()
 
         Z = np.atleast_2d(Z)
-        rows, cols = Z.shape
         # TODO: Support masked arrays
-        X = np.asarray(X)
-        Y = np.asarray(Y)
-        # Force X and Y to take the same shape.
-        # If they can not be fitted to that shape,
-        # then an exception is automatically thrown.
-        X.shape = (rows, cols)
-        Y.shape = (rows, cols)
+        X, Y, Z = np.broadcast_arrays(X, Y, Z)
+        rows, cols = Z.shape
 
         rstride = kwargs.pop('rstride', 10)
         cstride = kwargs.pop('cstride', 10)
@@ -1526,14 +1520,8 @@ def plot_wireframe(self, X, Y, Z, *args, **kwargs):
         had_data = self.has_data()
         Z = np.atleast_2d(Z)
         # FIXME: Support masked arrays
-        X = np.asarray(X)
-        Y = np.asarray(Y)
+        X, Y, Z = np.broadcast_arrays(X, Y, Z)
         rows, cols = Z.shape
-        # Force X and Y to take the same shape.
-        # If they can not be fitted to that shape,
-        # then an exception is automatically thrown.
-        X.shape = (rows, cols)
-        Y.shape = (rows, cols)
 
         # We want two sets of lines, one running along the "rows" of
         # Z and another set of lines running along the "columns" of Z.
