Add support for units in vmin/vmax

dstansby · dstansby · commit 9fd22ac9566c · 2024-02-20T11:44:54.000Z
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -484,10 +484,14 @@ def set_clim(self, vmin=None, vmax=None):
                 vmin, vmax = vmin
             except (TypeError, ValueError):
                 pass
-        if vmin is not None:
-            self.norm.vmin = colors._sanitize_extrema(vmin)
-        if vmax is not None:
-            self.norm.vmax = colors._sanitize_extrema(vmax)
+
+        def _process_lim(lim):
+            if self._converter is not None:
+                lim = self._converter.convert(lim, self._units, axis=None)
+            return colors._sanitize_extrema(lim)
+
+        self.norm.vmin = _process_lim(vmin)
+        self.norm.vmax = _process_lim(vmax)
 
     def get_alpha(self):
         """
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -305,22 +305,24 @@ def test_mappable_units(quantity_converter):
     munits.registry[Quantity] = quantity_converter
     x, y = np.meshgrid([0, 1], [0, 1])
     data = Quantity(np.arange(4).reshape(2, 2), 'hours')
+    vmin = Quantity(1, "hours")  # Test a limit different from min of the data
+    vmax = Quantity(3 * 60, "minutes")  # Test a different unit to the data
 
-    fig, axs = plt.subplots(nrows=2, ncols=2)
+    fig, axs = plt.subplots(nrows=2, ncols=2, constrained_layout=True)
 
     # imshow
     ax = axs[0, 0]
-    mappable = ax.imshow(data, origin='lower')
-    cbar = fig.colorbar(mappable, ax=ax)
+    mappable = ax.imshow(data, origin='lower', vmin=vmin, vmax=vmax)
+    cbar = fig.colorbar(mappable, ax=ax, extend="min")
 
     # pcolor
     ax = axs[0, 1]
-    mappable = ax.pcolor(x, y, data)
-    fig.colorbar(mappable, ax=ax)
+    mappable = ax.pcolor(x, y, data, vmin=vmin, vmax=vmax)
+    fig.colorbar(mappable, ax=ax, extend="min")
 
     # pcolormesh + horizontal colorbar
     ax = axs[1, 0]
-    mappable = ax.pcolormesh(x, y, data)
-    fig.colorbar(mappable, ax=ax, orientation="horizontal")
+    mappable = ax.pcolormesh(x, y, data, vmin=vmin, vmax=vmax)
+    fig.colorbar(mappable, ax=ax, orientation="horizontal", extend="min")
 
     axs[1, 1].axis("off")
