Included an option to normalize

DanHickstein · DanHickstein · commit ff2a207da4ca · 2015-11-10T12:27:43.000-07:00
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2543,6 +2543,8 @@ def calc_arrow(uvw, angle=15):
         arrow_length_ratio = kwargs.pop('arrow_length_ratio', 0.3)
         # pivot point
         pivot = kwargs.pop('pivot', 'tail')
+        # normalize
+        pivot = kwargs.pop('normalize', False)
 
         # handle args
         argi = 6
@@ -2603,7 +2605,10 @@ def calc_arrow(uvw, angle=15):
         # If any row of UVW is all zeros, don't make a quiver for it
         mask = norm > 1e-10
         XYZ = XYZ[mask]
-        UVW = UVW[mask]
+        if normalize == True:
+            UVW = UVW[mask]/ norm[mask].reshape((-1, 1))
+        else:
+            UVW = UVW[mask]
 
         if len(XYZ) > 0:
             # compute the shaft lines all at once with an outer product
