Curved polar errorbars

deep-jkl · deep-jkl · commit fec9b6098401 · 2022-08-21T12:36:43.000+02:00
- uses _interpolation_steps
 - prefers transform MarkerStyle in init over _transform property
 - adjusted what's new
 - added more tests for overlapping, asymmetric and long errorbars
 - combine all tests to a single figure
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 @@
+Fixed errorbars in polar plots
+------------------------------
+Caps and error lines are now drawn with respect to polar coordinates,
+when plotting errorbars on polar plots.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3522,10 +3522,11 @@ def _upcast_err(err):
         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,
@@ -3556,9 +3557,12 @@ def apply_mask(arrays, mask): return [array[mask] for array in arrays]
             #     return dep - elow * ~lolims, dep + ehigh * ~uplims
             # except that broadcast_to would strip units.
             low, high = dep + np.row_stack([-(1 - lolims), 1 - uplims]) * err
-
             barcols.append(lines_func(
                 *apply_mask([indep, low, high], everymask), **eb_lines_style))
+            if self.name == "polar" and dep_axis == "x":
+                for b in barcols:
+                    for p in b.get_paths():
+                        p._interpolation_steps = 2
             # Normal errorbars for points without upper/lower limits.
             nolims = ~(lolims | uplims)
             if nolims.any() and capsize > 0:
@@ -3571,7 +3575,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
@@ -3585,15 +3589,29 @@ 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)
+        if self.name == 'polar':
+            for axis in caplines:
+                for l in caplines[axis]:
+                    # Rotate caps to be perpendicular to the error bars
+                    for theta, r in zip(l.get_xdata(), l.get_ydata()):
+                        rotation = mtransforms.Affine2D().rotate(theta)
+                        if axis == 'y':
+                            rotation.rotate(-np.pi / 2)
+                        ms = mmarkers.MarkerStyle(marker=marker,
+                                                  transform=rotation)
+                        self.add_line(mlines.Line2D([theta], [r], marker=ms,
+                                                    **eb_cap_style))
+        else:
+            for axis in caplines:
+                for l in caplines[axis]:
+                    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/mixed_errorbar_polar_caps.png b/lib/matplotlib/tests/baseline_images/test_axes/mixed_errorbar_polar_caps.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3600,6 +3600,40 @@ def test_errorbar():
     ax.set_title("Simplest errorbars, 0.2 in x, 0.4 in y")
 
 
+@image_comparison(['mixed_errorbar_polar_caps'], extensions=['png'],
+                  remove_text=True)
+def test_mixed_errorbar_polar_caps():
+    fig = plt.figure()
+    ax = plt.subplot(111, projection='polar')
+    theta = np.linspace(0, np.pi, 3, False)
+    r = [1.0]*3
+    ax.errorbar(theta, r, xerr=0.35, yerr=0.2, fmt="o")
+    _long_errorbar_polar_caps(ax)
+    _asymmetric_polar_caps(ax)
+    _overlapping_polar_caps(ax)
+    ax.set_rlim(0, 4)
+
+
+def _long_errorbar_polar_caps(ax):
+    theta = np.linspace(0, 2 * np.pi, 2, False) + np.pi/2
+    r = [1.8, 2.2]
+    ax.errorbar(theta, r, xerr=0.8 * np.pi, yerr=0.2, fmt="o")
+
+
+def _asymmetric_polar_caps(ax):
+    theta = np.linspace(0, np.pi, 3, False) + np.pi
+    r = [1.0]*3
+    xerr = [[.3, .3, .2], [.2, .3, .3]]
+    yerr = [[.35, .5, .5], [.5, .35, .5]]
+    ax.errorbar(theta, r, xerr=xerr, yerr=yerr, fmt="o")
+
+
+def _overlapping_polar_caps(ax):
+    theta = 3*np.pi/4
+    r = 3
+    ax.errorbar(theta, r, xerr=8.55 * np.pi, yerr=0.2, fmt="o")
+
+
 def test_errorbar_colorcycle():
 
     f, ax = plt.subplots()
