Rotate caps on polar errorbar plots

dstansby · dstansby · commit ea337ef1d467 · 2021-09-13T12:56:55.000+01:00
diff --git a/doc/users/next_whats_new/polar_errorbar_caps.rst b/doc/users/next_whats_new/polar_errorbar_caps.rst
@@ -0,0 +1,4 @@
+Polar errorbar caps are rotated
+-------------------------------
+When plotting errorbars on a polar plot, the caps are now rotated so they are
+perpendicular to the errorbars.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3433,10 +3433,11 @@ def errorbar(self, x, y, yerr=None, xerr=None,
         eb_cap_style['color'] = ecolor
 
         barcols = []
-        caplines = []
+        caplines = {'x': [], 'y': []}
 
         # Vectorized fancy-indexer.
-        def apply_mask(arrays, mask): return [array[mask] for array in arrays]
+        def apply_mask(arrays, mask):
+            return [array[mask] for array in arrays]
 
         # dep: dependent dataset, indep: independent dataset
         for (dep_axis, dep, err, lolims, uplims, indep, lines_func,
@@ -3477,7 +3478,7 @@ def apply_mask(arrays, mask): return [array[mask] for array in arrays]
                     line = mlines.Line2D(indep_masked, indep_masked,
                                          marker=marker, **eb_cap_style)
                     line.set(**{f"{dep_axis}data": lh_masked})
-                    caplines.append(line)
+                    caplines[dep_axis].append(line)
             for idx, (lims, hl) in enumerate([(lolims, high), (uplims, low)]):
                 if not lims.any():
                     continue
@@ -3491,15 +3492,28 @@ def apply_mask(arrays, mask): return [array[mask] for array in arrays]
                 line = mlines.Line2D(x_masked, y_masked,
                                      marker=hlmarker, **eb_cap_style)
                 line.set(**{f"{dep_axis}data": hl_masked})
-                caplines.append(line)
+                caplines[dep_axis].append(line)
                 if capsize > 0:
-                    caplines.append(mlines.Line2D(
+                    caplines[dep_axis].append(mlines.Line2D(
                         x_masked, y_masked, marker=marker, **eb_cap_style))
 
-        for l in caplines:
-            self.add_line(l)
+        for axis in caplines:
+            for l in caplines[axis]:
+                if self.name == 'polar':
+                    # Rotate caps to be perpendicular to the error bars
+                    for theta, r in zip(l.get_xdata(), l.get_ydata()):
+                        rotation = theta
+                        if axis == 'x':
+                            rotation += np.pi / 2
+                        ms = mmarkers.MarkerStyle(marker=marker)
+                        ms._transform = mtransforms.Affine2D().rotate(rotation)
+                        self.add_line(mlines.Line2D([theta], [r], marker=ms,
+                                                    **eb_cap_style))
+                else:
+                    self.add_line(l)
 
         self._request_autoscale_view()
+        caplines = caplines['x'] + caplines['y']
         errorbar_container = ErrorbarContainer(
             (data_line, tuple(caplines), tuple(barcols)),
             has_xerr=(xerr is not None), has_yerr=(yerr is not None),
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/errorbar_polar_caps.png b/lib/matplotlib/tests/baseline_images/test_axes/errorbar_polar_caps.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3351,6 +3351,16 @@ def test_errorbar():
     ax.set_title("Simplest errorbars, 0.2 in x, 0.4 in y")
 
 
+@image_comparison(['errorbar_polar_caps'], extensions=['png'])
+def test_errorbar_polar_caps():
+    fig = plt.figure()
+    ax = plt.subplot(111, projection='polar')
+    theta = np.arange(0, 2*np.pi, np.pi / 4)
+    r = theta / np.pi / 2 + 0.5
+    ax.errorbar(theta, r, xerr=0.15, yerr=0.1)
+    ax.set_rlim(0, 2)
+
+
 def test_errorbar_colorcycle():
 
     f, ax = plt.subplots()
