Merge pull request matplotlib#3300 from WeatherGod/mplot3d/quiver_test_fix

tacaswell · tacaswell · commit 69d5ef481e82 · 2014-07-24T15:57:49.000-04:00
Quiver3d fixes
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2523,9 +2523,11 @@ def point_vector_to_line(point, vector, length):
         arrow_length_ratio = kwargs.pop('arrow_length_ratio', 0.3)
 
         # handle args
-        if len(args) < 6:
-            ValueError('Wrong number of arguments')
         argi = 6
+        if len(args) < argi:
+            ValueError('Wrong number of arguments. Expected %d got %d' %
+                       (argi, len(args)))
+
         # first 6 arguments are X, Y, Z, U, V, W
         input_args = args[:argi]
         # if any of the args are scalar, convert into list
@@ -2549,24 +2551,17 @@ def point_vector_to_line(point, vector, length):
 
         if any(len(v) == 0 for v in input_args):
             # No quivers, so just make an empty collection and return early
-            linec = art3d.Line3DCollection([], *args[6:], **kwargs)
+            linec = art3d.Line3DCollection([], *args[argi:], **kwargs)
             self.add_collection(linec)
             return linec
 
-        points = input_args[:3]
-        vectors = input_args[3:]
-
-        # Below assertions must be true before proceed
+        # Following assertions must be true before proceeding
         # must all be ndarray
         assert all(isinstance(k, np.ndarray) for k in input_args)
         # must all in same shape
         assert len(set([k.shape for k in input_args])) == 1
 
-        # X, Y, Z, U, V, W
-        coords = (np.array(k) if not isinstance(k, np.ndarray) else k
-                  for k in args)
-        coords = [k.flatten() for k in coords]
-        xs, ys, zs, us, vs, ws = coords
+        xs, ys, zs, us, vs, ws = input_args[:argi]
         lines = []
 
         # for each arrow
@@ -2578,12 +2573,11 @@ def point_vector_to_line(point, vector, length):
             u = us[i]
             v = vs[i]
             w = ws[i]
-            if any(k is np.ma.masked for k in [x, y, z, u, v, w]):
-                continue
 
             # (u,v,w) expected to be normalized, recursive to fix A=0 scenario.
             if u == 0 and v == 0 and w == 0:
-                raise ValueError("u,v,w can't be all zero")
+                # Just don't make a quiver for such a case.
+                continue
 
             # normalize
             norm = math.sqrt(u ** 2 + v ** 2 + w ** 2)
@@ -2603,6 +2597,7 @@ def point_vector_to_line(point, vector, length):
             ua1, va1, wa1 = d1[0], d1[1], d1[2]
             ua2, va2, wa2 = d2[0], d2[1], d2[2]
 
+            # TODO: num should probably get parameterized
             t = np.linspace(0, length * arrow_length_ratio, num=20)
             la1x = x - t * ua1
             la1y = y - t * va1
@@ -2616,7 +2611,7 @@ def point_vector_to_line(point, vector, length):
             line = list(zip(la2x, la2y, la2z))
             lines.append(line)
 
-        linec = art3d.Line3DCollection(lines, *args[6:], **kwargs)
+        linec = art3d.Line3DCollection(lines, *args[argi:], **kwargs)
         self.add_collection(linec)
 
         self.auto_scale_xyz(xs, ys, zs, had_data)
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -162,8 +162,7 @@ def test_quiver3d():
     fig = plt.figure()
     ax = fig.gca(projection='3d')
 
-    x, y, z = np.meshgrid(np.arange(-1, 1, 0.2), np.arange(-1, 1, 0.2),
-                          np.arange(-1, 1, 0.8))
+    x, y, z = np.ogrid[-1:0.8:10j, -1:0.8:10j, -1:0.6:3j]
 
     u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
     v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
@@ -177,7 +176,7 @@ def test_quiver3d_empty():
     fig = plt.figure()
     ax = fig.gca(projection='3d')
 
-    x, y, z = np.meshgrid([], [], [])
+    x, y, z = np.ogrid[-1:0.8:0j, -1:0.8:0j, -1:0.6:0j]
 
     u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
     v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
@@ -191,8 +190,9 @@ def test_quiver3d_masked():
     fig = plt.figure()
     ax = fig.gca(projection='3d')
 
-    x, y, z = np.meshgrid(np.arange(-1, 1, 0.2), np.arange(-1, 1, 0.2),
-                          np.arange(-1, 1, 0.8))
+    # Using mgrid here instead of ogrid because masked_where doesn't
+    # seem to like broadcasting very much...
+    x, y, z = np.mgrid[-1:0.8:10j, -1:0.8:10j, -1:0.6:3j]
 
     u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
     v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
