Merge pull request #25224 from ZachDaChampion/secondary-axes-trans-arg

timhoffm · web-flow · commit 5abcae28eebf · 2024-03-08T11:37:34.000+01:00
Allow passing a transformation to secondary_xaxis/_yaxis
diff --git a/galleries/examples/subplots_axes_and_figures/secondary_axis.py b/galleries/examples/subplots_axes_and_figures/secondary_axis.py
@@ -40,6 +40,25 @@ def rad2deg(x):
 secax.set_xlabel('angle [rad]')
 plt.show()
 
+# %%
+# By default, the secondary axis is drawn in the Axes coordinate space.
+# We can also provide a custom transform to place it in a different
+# coordinate space. Here we put the axis at Y = 0 in data coordinates.
+
+fig, ax = plt.subplots(layout='constrained')
+x = np.arange(0, 10)
+np.random.seed(19680801)
+y = np.random.randn(len(x))
+ax.plot(x, y)
+ax.set_xlabel('X')
+ax.set_ylabel('Y')
+ax.set_title('Random data')
+
+# Pass ax.transData as a transform to place the axis relative to our data
+secax = ax.secondary_xaxis(0, transform=ax.transData)
+secax.set_xlabel('Axis at Y = 0')
+plt.show()
+
 # %%
 # Here is the case of converting from wavenumber to wavelength in a
 # log-log scale.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -551,7 +551,7 @@ def indicate_inset_zoom(self, inset_ax, **kwargs):
         return self.indicate_inset(rect, inset_ax, **kwargs)
 
     @_docstring.dedent_interpd
-    def secondary_xaxis(self, location, *, functions=None, **kwargs):
+    def secondary_xaxis(self, location, *, functions=None, transform=None, **kwargs):
         """
         Add a second x-axis to this `~.axes.Axes`.
 
@@ -582,18 +582,30 @@ def invert(x):
             secax = ax.secondary_xaxis('top', functions=(invert, invert))
             secax.set_xlabel('Period [s]')
             plt.show()
+
+        To add a secondary axis relative to your data, you can pass a transform
+        to the new axis.
+
+        .. plot::
+
+            fig, ax = plt.subplots()
+            ax.plot(range(0, 5), range(-1, 4))
+
+            # Pass 'ax.transData' as a transform to place the axis
+            # relative to your data at y=0
+            secax = ax.secondary_xaxis(0, transform=ax.transData)
         """
-        if location in ['top', 'bottom'] or isinstance(location, Real):
-            secondary_ax = SecondaryAxis(self, 'x', location, functions,
-                                         **kwargs)
-            self.add_child_axes(secondary_ax)
-            return secondary_ax
-        else:
+        if not (location in ['top', 'bottom'] or isinstance(location, Real)):
             raise ValueError('secondary_xaxis location must be either '
                              'a float or "top"/"bottom"')
 
+        secondary_ax = SecondaryAxis(self, 'x', location, functions,
+                                     transform, **kwargs)
+        self.add_child_axes(secondary_ax)
+        return secondary_ax
+
     @_docstring.dedent_interpd
-    def secondary_yaxis(self, location, *, functions=None, **kwargs):
+    def secondary_yaxis(self, location, *, functions=None, transform=None, **kwargs):
         """
         Add a second y-axis to this `~.axes.Axes`.
 
@@ -614,16 +626,28 @@ def secondary_yaxis(self, location, *, functions=None, **kwargs):
             secax = ax.secondary_yaxis('right', functions=(np.deg2rad,
                                                            np.rad2deg))
             secax.set_ylabel('radians')
+
+        To add a secondary axis relative to your data, you can pass a transform
+        to the new axis.
+
+        .. plot::
+
+            fig, ax = plt.subplots()
+            ax.plot(range(0, 5), range(-1, 4))
+
+            # Pass 'ax.transData' as a transform to place the axis
+            # relative to your data at x=3
+            secax = ax.secondary_yaxis(3, transform=ax.transData)
         """
-        if location in ['left', 'right'] or isinstance(location, Real):
-            secondary_ax = SecondaryAxis(self, 'y', location,
-                                         functions, **kwargs)
-            self.add_child_axes(secondary_ax)
-            return secondary_ax
-        else:
+        if not (location in ['left', 'right'] or isinstance(location, Real)):
             raise ValueError('secondary_yaxis location must be either '
                              'a float or "left"/"right"')
 
+        secondary_ax = SecondaryAxis(self, 'y', location, functions,
+                                     transform, **kwargs)
+        self.add_child_axes(secondary_ax)
+        return secondary_ax
+
     @_docstring.dedent_interpd
     def text(self, x, y, s, fontdict=None, **kwargs):
         """
diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -94,6 +94,7 @@ class Axes(_AxesBase):
         ]
         | Transform
         | None = ...,
+        transform: Transform | None = ...,
         **kwargs
     ) -> SecondaryAxis: ...
     def secondary_yaxis(
@@ -105,6 +106,7 @@ class Axes(_AxesBase):
         ]
         | Transform
         | None = ...,
+        transform: Transform | None = ...,
         **kwargs
     ) -> SecondaryAxis: ...
     def text(
diff --git a/lib/matplotlib/axes/_secondary_axes.py b/lib/matplotlib/axes/_secondary_axes.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from matplotlib import _api, _docstring
+from matplotlib import _api, _docstring, transforms
 import matplotlib.ticker as mticker
 from matplotlib.axes._base import _AxesBase, _TransformedBoundsLocator
 from matplotlib.axis import Axis
@@ -14,7 +14,8 @@ class SecondaryAxis(_AxesBase):
     General class to hold a Secondary_X/Yaxis.
     """
 
-    def __init__(self, parent, orientation, location, functions, **kwargs):
+    def __init__(self, parent, orientation, location, functions, transform=None,
+                 **kwargs):
         """
         See `.secondary_xaxis` and `.secondary_yaxis` for the doc string.
         While there is no need for this to be private, it should really be
@@ -39,7 +40,7 @@ def __init__(self, parent, orientation, location, functions, **kwargs):
         self._parentscale = None
         # this gets positioned w/o constrained_layout so exclude:
 
-        self.set_location(location)
+        self.set_location(location, transform)
         self.set_functions(functions)
 
         # styling:
@@ -74,7 +75,7 @@ def set_alignment(self, align):
         self._axis.set_ticks_position(align)
         self._axis.set_label_position(align)
 
-    def set_location(self, location):
+    def set_location(self, location, transform=None):
         """
         Set the vertical or horizontal location of the axes in
         parent-normalized coordinates.
@@ -87,8 +88,17 @@ def set_location(self, location):
             orientation='y'. A float indicates the relative position on the
             parent Axes to put the new Axes, 0.0 being the bottom (or left)
             and 1.0 being the top (or right).
+
+        transform : `.Transform`, optional
+            Transform for the location to use. Defaults to
+            the parent's ``transAxes``, so locations are normally relative to
+            the parent axes.
+
+            .. versionadded:: 3.9
         """
 
+        _api.check_isinstance((transforms.Transform, None), transform=transform)
+
         # This puts the rectangle into figure-relative coordinates.
         if isinstance(location, str):
             _api.check_in_list(self._locstrings, location=location)
@@ -106,15 +116,28 @@ def set_location(self, location):
             # An x-secondary axes is like an inset axes from x = 0 to x = 1 and
             # from y = pos to y = pos + eps, in the parent's transAxes coords.
             bounds = [0, self._pos, 1., 1e-10]
+
+            # If a transformation is provided, use its y component rather than
+            # the parent's transAxes. This can be used to place axes in the data
+            # coords, for instance.
+            if transform is not None:
+                transform = transforms.blended_transform_factory(
+                    self._parent.transAxes, transform)
         else:  # 'y'
             bounds = [self._pos, 0, 1e-10, 1]
+            if transform is not None:
+                transform = transforms.blended_transform_factory(
+                    transform, self._parent.transAxes)  # Use provided x axis
+
+        # If no transform is provided, use the parent's transAxes
+        if transform is None:
+            transform = self._parent.transAxes
 
         # this locator lets the axes move in the parent axes coordinates.
         # so it never needs to know where the parent is explicitly in
         # figure coordinates.
         # it gets called in ax.apply_aspect() (of all places)
-        self.set_axes_locator(
-            _TransformedBoundsLocator(bounds, self._parent.transAxes))
+        self.set_axes_locator(_TransformedBoundsLocator(bounds, transform))
 
     def apply_aspect(self, position=None):
         # docstring inherited.
@@ -278,6 +301,14 @@ def set_color(self, color):
     See :doc:`/gallery/subplots_axes_and_figures/secondary_axis`
     for examples of making these conversions.
 
+transform : `.Transform`, optional
+    If specified, *location* will be
+    placed relative to this transform (in the direction of the axis)
+    rather than the parent's axis. i.e. a secondary x-axis will
+    use the provided y transform and the x transform of the parent.
+
+    .. versionadded:: 3.9
+
 Returns
 -------
 ax : axes._secondary_axes.SecondaryAxis
diff --git a/lib/matplotlib/axes/_secondary_axes.pyi b/lib/matplotlib/axes/_secondary_axes.pyi
@@ -18,13 +18,16 @@ class SecondaryAxis(_AxesBase):
             Callable[[ArrayLike], ArrayLike], Callable[[ArrayLike], ArrayLike]
         ]
         | Transform,
+        transform: Transform | None = ...,
         **kwargs
     ) -> None: ...
     def set_alignment(
         self, align: Literal["top", "bottom", "right", "left"]
     ) -> None: ...
     def set_location(
-        self, location: Literal["top", "bottom", "right", "left"] | float
+        self,
+        location: Literal["top", "bottom", "right", "left"] | float,
+        transform: Transform | None = ...
     ) -> None: ...
     def set_ticks(
         self,
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/secondary_xy.png b/lib/matplotlib/tests/baseline_images/test_axes/secondary_xy.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -7549,6 +7549,7 @@ def invert(x):
         secax(0.6, functions=(lambda x: x**2, lambda x: x**(1/2)))
         secax(0.8)
         secax("top" if nn == 0 else "right", functions=_Translation(2))
+        secax(6.25, transform=ax.transData)
 
 
 def test_secondary_fail():
@@ -7560,6 +7561,8 @@ def test_secondary_fail():
         ax.secondary_xaxis('right')
     with pytest.raises(ValueError):
         ax.secondary_yaxis('bottom')
+    with pytest.raises(TypeError):
+        ax.secondary_xaxis(0.2, transform='error')
 
 
 def test_secondary_resize():
