Merge pull request #5166 from Tillsten/fix_#3116

tacaswell · tacaswell · commit 55110613be5f · 2015-11-04T23:12:34.000-05:00
FIX: Don't allow 1d-arrays in plot_surface.
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1581,7 +1581,8 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
 
         had_data = self.has_data()
 
-        Z = np.atleast_2d(Z)
+        if Z.ndim != 2:
+            raise ValueError("Argument Z must be 2-dimensional.")
         # TODO: Support masked arrays
         X, Y, Z = np.broadcast_arrays(X, Y, Z)
         rows, cols = Z.shape
@@ -1755,7 +1756,8 @@ def plot_wireframe(self, X, Y, Z, *args, **kwargs):
         cstride = kwargs.pop("cstride", 1)
 
         had_data = self.has_data()
-        Z = np.atleast_2d(Z)
+        if Z.ndim != 2:
+            raise ValueError("Argument Z must be 2-dimensional.")
         # FIXME: Support masked arrays
         X, Y, Z = np.broadcast_arrays(X, Y, Z)
         rows, cols = Z.shape
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -311,7 +311,17 @@ def test_axes3d_cla():
     ax.set_axis_off()
     ax.cla()  # make sure the axis displayed is 3D (not 2D)
 
-
+@cleanup
+def test_plotsurface_1d_raises():
+    x = np.linspace(0.5, 10, num=100)
+    y = np.linspace(0.5, 10, num=100)
+    X, Y = np.meshgrid(x, y)
+    z = np.random.randn(100)
+
+    fig = plt.figure(figsize=(14,6))
+    ax = fig.add_subplot(1, 2, 1, projection='3d')
+    assert_raises(ValueError, ax.plot_surface, X, Y, z)
+    
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
