FIX: add support for imshow extent to have units

pratimugale · pratimugale · commit 592bc4db5836 · 2022-01-31T02:04:32.000+05:30
diff --git a/doc/users/next_whats_new/imshow_extent_units.rst b/doc/users/next_whats_new/imshow_extent_units.rst
@@ -0,0 +1,30 @@
+The *extent* parameter of imshow and set_extent functions can now be expressed with Units
+-----------------------------------------------------------------------------------------
+The extent parameter of `~axes.Axes.imshow` can now be expressed
+with Units. It internally uses the `~AxesImage.set_extent` which contains
+the implemenation that enables extent being expressed with Units.
+If extent[0] is expressed in units, then so must extent[1].
+Similarly, this applies to extent[2] and extent[3] as well. The user is 
+responsible to ensure that the units are lined up.
+
+.. plot::
+    :include-source: true
+    
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from matplotlib.dates import HourLocator, DateFormatter
+    from matplotlib.ticker import (AutoMinorLocator)
+
+
+    fig, ax = plt.subplots()
+    dates = np.arange("2020-01-01","2020-01-13", dtype='datetime64[h]')
+
+    arr = [[i+j for i in range(10)] for j in range(10)]
+
+    ax.imshow(arr, origin='lower', extent=[dates[0], dates[10], dates[10], dates[0]])
+
+    ax.xaxis.set_major_formatter(DateFormatter('%d/%m:%H'))
+    ax.xaxis.set_major_locator(HourLocator(byhour=[0, 2, 4, 6, 8, 10]))
+    ax.xaxis.set_minor_locator(AutoMinorLocator())
+
+    plt.show()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5429,7 +5429,7 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
             See the :doc:`/tutorials/intermediate/imshow_extent` tutorial for
             examples and a more detailed description.
 
-        extent : floats (left, right, bottom, top), optional
+        extent : floats or units (left, right, bottom, top), optional
             The bounding box in data coordinates that the image will fill.
             The image is stretched individually along x and y to fill the box.
 
@@ -5506,6 +5506,7 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         if aspect is None:
             aspect = rcParams['image.aspect']
         self.set_aspect(aspect)
+
         im = mimage.AxesImage(self, cmap, norm, interpolation,
                               origin, extent, filternorm=filternorm,
                               filterrad=filterrad, resample=resample,
@@ -5522,7 +5523,7 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
 
         # update ax.dataLim, and, if autoscaling, set viewLim
         # to tightly fit the image, regardless of dataLim.
-        im.set_extent(im.get_extent())
+        im.set_extent(im.get_extent(), **kwargs)
 
         self.add_image(im)
         return im
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -949,7 +949,7 @@ def _check_unsampled_image(self):
         """Return whether the image would be better drawn unsampled."""
         return self.get_interpolation() == "none"
 
-    def set_extent(self, extent):
+    def set_extent(self, extent, **kwargs):
         """
         Set the image extent.
 
@@ -958,6 +958,10 @@ def set_extent(self, extent):
         extent : 4-tuple of float
             The position and size of the image as tuple
             ``(left, right, bottom, top)`` in data coordinates.
+        kwargs : dict
+            Other parameters from which unit info (i.e., the *xunits*,
+            *yunits*, *zunits* (for 3D axes), *runits* and *thetaunits* (for
+            polar axes) entries) is popped, if present.
 
         Notes
         -----
@@ -966,7 +970,21 @@ def set_extent(self, extent):
         state is not changed, so following this with ``ax.autoscale_view()``
         will redo the autoscaling in accord with ``dataLim``.
         """
-        self._extent = xmin, xmax, ymin, ymax = extent
+        (xmin, xmax), = self.axes._process_unit_info(
+            [("x", [extent[0], extent[1]])], kwargs)
+        (ymin, ymax), = self.axes._process_unit_info(
+            [("y", [extent[2], extent[3]])], kwargs)
+        xmin = self.axes._validate_converted_limits(
+            xmin, self.convert_xunits)
+        xmax = self.axes._validate_converted_limits(
+            xmax, self.convert_xunits)
+        ymin = self.axes._validate_converted_limits(
+            ymin, self.convert_yunits)
+        ymax = self.axes._validate_converted_limits(
+            ymax, self.convert_yunits)
+        extent = [xmin, xmax, ymin, ymax]
+
+        self._extent = extent
         corners = (xmin, ymin), (xmax, ymax)
         self.axes.update_datalim(corners)
         self.sticky_edges.x[:] = [xmin, xmax]
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/extent_units.pdf b/lib/matplotlib/tests/baseline_images/test_axes/extent_units.pdf
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -7539,3 +7539,13 @@ def test_clim():
         clim = (7, 8)
         norm = plot_method(clim=clim).norm
         assert (norm.vmin, norm.vmax) == clim
+
+
+@image_comparison(["extent_units.pdf"])
+def test_extent_units():
+    mpl.style.use("mpl20")
+    fig, ax = plt.subplots()
+    dates = np.arange("2020-01-01", "2020-01-13", dtype='datetime64')
+    arr = [[i+j for i in range(10)] for j in range(10)]
+    ax.imshow(arr, origin='lower', extent=[0, 10,
+              dates[10], dates[0]])
