Simplifications to quiver3d.

anntzer · anntzer · commit c4ad02b51a61 · 2018-11-12T16:39:29.000+01:00
- No need to convert input_args to lists: broadcast_arrays will handle
  the broadcasting just fine on scalars, and later using np.ravel() will
  ensure that the result is at least 1D.
- The assertions don't really add much (and are clearly always true
  based on the code just above).
- np.linspace(x, y, 2) is a pretty obfuscated way to write np.array([x,
  y]).
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2637,15 +2637,12 @@ def calc_arrow(uvw, angle=15):
 
         # first 6 arguments are X, Y, Z, U, V, W
         input_args = args[:argi]
-        # if any of the args are scalar, convert into list
-        input_args = [[k] if isinstance(k, (int, float)) else k
-                      for k in input_args]
 
         # extract the masks, if any
         masks = [k.mask for k in input_args
                  if isinstance(k, np.ma.MaskedArray)]
         # broadcast to match the shape
-        bcast = np.broadcast_arrays(*(input_args + masks))
+        bcast = np.broadcast_arrays(*input_args, *masks)
         input_args = bcast[:argi]
         masks = bcast[argi:]
         if masks:
@@ -2655,21 +2652,15 @@ def calc_arrow(uvw, angle=15):
             input_args = [np.ma.array(k, mask=mask).compressed()
                           for k in input_args]
         else:
-            input_args = [k.flatten() for k in input_args]
+            input_args = [np.ravel(k) for k in input_args]
 
         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[argi:], **kwargs)
             self.add_collection(linec)
             return linec
 
-        # 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({k.shape for k in input_args}) == 1
-
-        shaft_dt = np.linspace(0, length, num=2)
+        shaft_dt = np.array([0, length])
         arrow_dt = shaft_dt * arrow_length_ratio
 
         if pivot == 'tail':
