Merge pull request #23592 from deep-jkl/polar-errcaps

timhoffm · web-flow · commit 239386650a2c · 2022-10-07T22:11:07.000+02:00
Polar errcaps
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,7 @@
+Fixed errorbars in polar plots
+------------------------------
+Caps and error lines are now drawn with respect to polar coordinates,
+when plotting errorbars on polar plots.
+
+.. figure:: /gallery/pie_and_polar_charts/images/sphx_glr_polar_error_caps_001.png
+   :target: ../../gallery/pie_and_polar_charts/polar_error_caps.html
diff --git a/examples/pie_and_polar_charts/polar_error_caps.py b/examples/pie_and_polar_charts/polar_error_caps.py
@@ -0,0 +1,53 @@
+"""
+=================================
+Error bar rendering on polar axis
+=================================
+
+Demo of error bar plot in polar coordinates.
+Theta error bars are curved lines ended with caps oriented towards the
+center.
+Radius error bars are straight lines oriented towards center with
+perpendicular caps.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+theta = np.arange(0, 2 * np.pi, np.pi / 4)
+r = theta / np.pi / 2 + 0.5
+
+fig = plt.figure(figsize=(10, 10))
+ax = fig.add_subplot(projection='polar')
+ax.errorbar(theta, r, xerr=0.25, yerr=0.1, capsize=7, fmt="o", c="seagreen")
+ax.set_title("Pretty polar error bars")
+plt.show()
+
+#############################################################################
+# Please acknowledge that large theta error bars will be overlapping.
+# This may reduce readability of the output plot. See example figure below:
+
+fig = plt.figure(figsize=(10, 10))
+ax = fig.add_subplot(projection='polar')
+ax.errorbar(theta, r, xerr=5.25, yerr=0.1, capsize=7, fmt="o", c="darkred")
+ax.set_title("Overlapping theta error bars")
+plt.show()
+
+#############################################################################
+# On the other hand, large radius error bars will never overlap, they just
+# lead to unwanted scale in the data, reducing the displayed range.
+
+fig = plt.figure(figsize=(10, 10))
+ax = fig.add_subplot(projection='polar')
+ax.errorbar(theta, r, xerr=0.25, yerr=10.1, capsize=7, fmt="o", c="orangered")
+ax.set_title("Large radius error bars")
+plt.show()
+
+
+#############################################################################
+#
+# .. admonition:: References
+#
+#    The use of the following functions, methods, classes and modules is shown
+#    in this example:
+#
+#    - `matplotlib.axes.Axes.errorbar` / `matplotlib.pyplot.errorbar`
+#    - `matplotlib.projections.polar`
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3573,10 +3573,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,
@@ -3607,9 +3608,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:
@@ -3622,7 +3626,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
@@ -3636,15 +3640,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
@@ -3687,6 +3687,41 @@ 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():
+    """
+    Mix several polar errorbar use cases in a single test figure.
+
+    It is advisable to position individual points off the grid. If there are
+    problems with reproducibility of this test, consider removing grid.
+    """
+    fig = plt.figure()
+    ax = plt.subplot(111, projection='polar')
+
+    # symmetric errorbars
+    th_sym = [1, 2, 3]
+    r_sym = [0.9]*3
+    ax.errorbar(th_sym, r_sym, xerr=0.35, yerr=0.2, fmt="o")
+
+    # long errorbars
+    th_long = [np.pi/2 + .1, np.pi + .1]
+    r_long = [1.8, 2.2]
+    ax.errorbar(th_long, r_long, xerr=0.8 * np.pi, yerr=0.15, fmt="o")
+
+    # asymmetric errorbars
+    th_asym = [4*np.pi/3 + .1, 5*np.pi/3 + .1, 2*np.pi-0.1]
+    r_asym = [1.1]*3
+    xerr = [[.3, .3, .2], [.2, .3, .3]]
+    yerr = [[.35, .5, .5], [.5, .35, .5]]
+    ax.errorbar(th_asym, r_asym, xerr=xerr, yerr=yerr, fmt="o")
+
+    # overlapping errorbar
+    th_over = [2.1]
+    r_over = [3.1]
+    ax.errorbar(th_over, r_over, xerr=10, yerr=.2, fmt="o")
+
+
 def test_errorbar_colorcycle():
 
     f, ax = plt.subplots()
