fix polar error bar cap orientation

kyracho · kyracho · commit dd8a59847b57 · 2024-10-13T12:57:41.000Z
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -37,6 +37,7 @@
     _AxesBase, _TransformedBoundsLocator, _process_plot_format)
 from matplotlib.axes._secondary_axes import SecondaryAxis
 from matplotlib.container import BarContainer, ErrorbarContainer, StemContainer
+from matplotlib.transforms import ScaledRotation
 
 _log = logging.getLogger(__name__)
 
@@ -3784,13 +3785,16 @@ def apply_mask(arrays, mask):
                     caplines[dep_axis].append(mlines.Line2D(
                         x_masked, y_masked, marker=marker, **eb_cap_style))
         if self.name == 'polar':
+            trans_shift = self.transShift
             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)
+                        rotation = ScaledRotation(
+                            theta=theta,
+                            trans_shift=trans_shift)
                         if axis == 'y':
-                            rotation.rotate(-np.pi / 2)
+                            rotation += mtransforms.Affine2D().rotate(np.pi / 2)
                         ms = mmarkers.MarkerStyle(marker=marker,
                                                   transform=rotation)
                         self.add_line(mlines.Line2D([theta], [r], marker=ms,
diff --git a/lib/matplotlib/tests/baseline_images/test_polar/polar_errorbar.png b/lib/matplotlib/tests/baseline_images/test_polar/polar_errorbar.png
diff --git a/lib/matplotlib/tests/test_polar.py b/lib/matplotlib/tests/test_polar.py
@@ -481,3 +481,21 @@ def test_polar_neg_theta_lims():
     ax.set_thetalim(-np.pi, np.pi)
     labels = [l.get_text() for l in ax.xaxis.get_ticklabels()]
     assert labels == ['-180°', '-135°', '-90°', '-45°', '0°', '45°', '90°', '135°']
+
+
+@pytest.mark.parametrize("order", ["before", "after"])
+@image_comparison(baseline_images=['polar_errorbar'], remove_text=True,
+                  extensions=['png'], style='mpl20')
+def test_polar_errorbar(order):
+    theta = np.arange(0, 2 * np.pi, np.pi / 8)
+    r = theta / np.pi / 2 + 0.5
+    fig = plt.figure(figsize=(5, 5))
+    ax = fig.add_subplot(projection='polar')
+    if order == "before":
+        ax.set_theta_zero_location("N")
+        ax.set_theta_direction(-1)
+        ax.errorbar(theta, r, xerr=0.1, yerr=0.1, capsize=7, fmt="o", c="seagreen")
+    else:
+        ax.errorbar(theta, r, xerr=0.1, yerr=0.1, capsize=7, fmt="o", c="seagreen")
+        ax.set_theta_zero_location("N")
+        ax.set_theta_direction(-1)
diff --git a/lib/matplotlib/tests/test_transforms.py b/lib/matplotlib/tests/test_transforms.py
@@ -9,9 +9,10 @@
 import matplotlib.pyplot as plt
 import matplotlib.patches as mpatches
 import matplotlib.transforms as mtransforms
-from matplotlib.transforms import Affine2D, Bbox, TransformedBbox
+from matplotlib.transforms import Affine2D, Bbox, TransformedBbox, ScaledRotation
 from matplotlib.path import Path
 from matplotlib.testing.decorators import image_comparison, check_figures_equal
+from unittest.mock import MagicMock
 
 
 class TestAffine2D:
@@ -1104,3 +1105,26 @@ def test_interval_contains_open():
     assert not mtransforms.interval_contains_open((0, 1), -1)
     assert not mtransforms.interval_contains_open((0, 1), 2)
     assert mtransforms.interval_contains_open((1, 0), 0.5)
+
+
+def test_initialization():
+    """Test that the ScaledRotation object is initialized correctly."""
+    theta = 1.0  # Arbitrary theta value for testing
+    trans_shift = MagicMock()  # Mock the trans_shift transformation
+    scaled_rot = ScaledRotation(theta, trans_shift)
+    assert scaled_rot._theta == theta
+    assert scaled_rot._trans_shift == trans_shift
+    assert scaled_rot._mtx is None
+
+
+def test_get_matrix_invalid():
+    """Test get_matrix when the matrix is invalid and needs recalculation."""
+    theta = 1.0
+    trans_shift = MagicMock(transform=MagicMock(return_value=[[theta, 0]]))
+    scaled_rot = ScaledRotation(theta, trans_shift)
+    scaled_rot._invalid = True
+    matrix = scaled_rot.get_matrix()
+    trans_shift.transform.assert_called_once_with([[theta, 0]])
+    assert matrix is not None
+    expected_rotation = Affine2D().rotate(theta).get_matrix()[:2, :2]
+    assert (matrix[:2, :2] == expected_rotation).all()
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -49,6 +49,7 @@
 from matplotlib import _api
 from matplotlib._path import (
     affine_transform, count_bboxes_overlapping_bbox, update_path_extents)
+from matplotlib import transforms
 from .path import Path
 
 DEBUG = False
@@ -2685,6 +2686,27 @@ def get_matrix(self):
         return self._mtx
 
 
+class ScaledRotation(Affine2DBase):
+    """
+    A transformation that applies offset and direction
+    based on *trans_shift*.
+    """
+    def __init__(self, theta, trans_shift):
+        super().__init__()
+        self._theta = theta
+        self._trans_shift = trans_shift
+        self._mtx = None
+
+    def get_matrix(self):
+        if self._invalid:
+            transformed_coords = self._trans_shift.transform([[self._theta, 0]])[0]
+            adjusted_theta = transformed_coords[0]
+            rotation = transforms.Affine2D().rotate(adjusted_theta)
+            self._mtx = IdentityTransform._mtx.copy()
+            self._mtx[:2, :2] = rotation.get_matrix()[:2, :2]
+        return self._mtx
+
+
 class AffineDeltaTransform(Affine2DBase):
     r"""
     A transform wrapper for transforming displacements between pairs of points.
diff --git a/lib/matplotlib/transforms.pyi b/lib/matplotlib/transforms.pyi
@@ -334,3 +334,8 @@ def offset_copy(
     y: float = ...,
     units: Literal["inches", "points", "dots"] = ...,
 ) -> Transform: ...
+
+
+class ScaledRotation(Affine2DBase):
+    def __init__(self, theta: float, trans_shift: Transform) -> None: ...
+    def get_matrix(self) -> np.ndarray: ...
