Backport PR #27754: fix quiver3d incorrect arrow colors

timhoffm · dstansby · commit 6893bc843d79 · 2024-03-03T18:05:20.000Z
Update figure test image
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2727,15 +2727,10 @@ def calc_arrows(UVW):
         UVW = np.column_stack(input_args[3:]).astype(float)
 
         # Normalize rows of UVW
-        norm = np.linalg.norm(UVW, axis=1)
-
-        # If any row of UVW is all zeros, don't make a quiver for it
-        mask = norm > 0
-        XYZ = XYZ[mask]
         if normalize:
-            UVW = UVW[mask] / norm[mask].reshape((-1, 1))
-        else:
-            UVW = UVW[mask]
+            norm = np.linalg.norm(UVW, axis=1)
+            norm[norm == 0] = 1
+            UVW = UVW / norm.reshape((-1, 1))
 
         if len(XYZ) > 0:
             # compute the shaft lines all at once with an outer product
@@ -2749,7 +2744,7 @@ def calc_arrows(UVW):
             # transpose to get a list of lines
             heads = heads.swapaxes(0, 1)
 
-            lines = [*shafts, *heads]
+            lines = [*shafts, *heads[::2], *heads[1::2]]
         else:
             lines = []
 
diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/quiver3d_colorcoded.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/quiver3d_colorcoded.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -806,7 +806,8 @@ def test_mixedsamplesraises():
         ax.plot_surface(X, Y, Z, cstride=50, rcount=10)
 
 
-@mpl3d_image_comparison(['quiver3d.png'], style='mpl20')
+# remove tolerance when regenerating the test image
+@mpl3d_image_comparison(['quiver3d.png'], style='mpl20', tol=0.003)
 def test_quiver3d():
     fig = plt.figure()
     ax = fig.add_subplot(projection='3d')
@@ -853,6 +854,19 @@ def test_quiver3d_masked():
     ax.quiver(x, y, z, u, v, w, length=0.1, pivot='tip', normalize=True)
 
 
+@mpl3d_image_comparison(['quiver3d_colorcoded.png'], style='mpl20')
+def test_quiver3d_colorcoded():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    x = y = dx = dz = np.zeros(10)
+    z = dy = np.arange(10.)
+
+    color = plt.cm.Reds(dy/dy.max())
+    ax.quiver(x, y, z, dx, dy, dz, colors=color)
+    ax.set_ylim(0, 10)
+
+
 def test_patch_modification():
     fig = plt.figure()
     ax = fig.add_subplot(projection="3d")
@@ -1519,7 +1533,8 @@ def test_minor_ticks():
     ax.set_zticklabels(["half"], minor=True)
 
 
-@mpl3d_image_comparison(['errorbar3d_errorevery.png'], style='mpl20')
+# remove tolerance when regenerating the test image
+@mpl3d_image_comparison(['errorbar3d_errorevery.png'], style='mpl20', tol=0.003)
 def test_errorbar3d_errorevery():
     """Tests errorevery functionality for 3D errorbars."""
     t = np.arange(0, 2*np.pi+.1, 0.01)
