Basic support for units on ScalarMappable

dstansby · dstansby · commit b786a555181f · 2024-02-20T11:44:54.000Z
diff --git a/doc/api/next_api_changes/behavior/27721-DS.rst b/doc/api/next_api_changes/behavior/27721-DS.rst
@@ -0,0 +1,16 @@
+Unit converters can now support units in images
+-----------------------------------------------
+
+`~.cm.ScalarMappable` can now contain data with units. This adds support for
+unit-ful data to be plotted using - `~.axes.Axes.imshow`, `~.axes.Axes.pcolor`,
+and `~.axes.Axes.pcolormesh`
+
+For this to be supported by third-party `~.units.ConversionInterface`,
+the `~.units.ConversionInterface.default_units` and
+`~.units.ConversionInterface.convert` methods must allow for the *axis*
+argument to be ``None``, and `~.units.ConversionInterface.convert` must be able to
+convert data of more than one dimension (e.g. when plotting images the data is 2D).
+
+If a conversion interface raises an error when given ``None`` or 2D data as described
+above, this error will be re-raised when a user tries to use one of the newly supported
+plotting methods with unit-ful data.
diff --git a/doc/users/next_whats_new/image_units.rst b/doc/users/next_whats_new/image_units.rst
@@ -0,0 +1,8 @@
+Unit support for images
+-----------------------
+This adds support for image data with units has been added to the following plotting
+methods:
+
+- `~.axes.Axes.imshow`
+- `~.axes.Axes.pcolor`
+- `~.axes.Axes.pcolormesh`
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -24,6 +24,7 @@
 from matplotlib import _api, colors, cbook, scale
 from matplotlib._cm import datad
 from matplotlib._cm_listed import cmaps as cmaps_listed
+import matplotlib.units as munits
 
 
 _LUTSIZE = mpl.rcParams['image.lut']
@@ -283,6 +284,8 @@ def __init__(self, norm=None, cmap=None):
             The colormap used to map normalized data values to RGBA colors.
         """
         self._A = None
+        self._units = None
+        self._converter = None
         self._norm = None  # So that the setter knows we're initializing.
         self.set_norm(norm)  # The Normalize instance of this ScalarMappable.
         self.cmap = None  # So that the setter knows we're initializing.
@@ -387,6 +390,35 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
         rgba = self.cmap(x, alpha=alpha, bytes=bytes)
         return rgba
 
+    def _strip_units(self, A):
+        """
+        Remove units from A, and save the units and converter used to do the conversion.
+        """
+        self._converter = munits.registry.get_converter(A)
+        if self._converter is None:
+            self._units = None
+            return A
+
+        try:
+            self._units = self._converter.default_units(A, None)
+        except Exception as e:
+            raise RuntimeError(
+                f'{self._converter} failed when trying to return the default units for '
+                'this image. This may be because support has not been '
+                'implemented for `axis=None` in the default_units() method.'
+            ) from e
+
+        try:
+            A = self._converter.convert(A, self._units, None)
+        except Exception as e:
+            raise munits.ConversionError(
+                f'{self._converter} failed when trying to convert the units for this '
+                'image. This may be because support has not been implemented '
+                'for `axis=None` in the convert() method.'
+            ) from e
+
+        return A
+
     def set_array(self, A):
         """
         Set the value array from array-like *A*.
@@ -402,7 +434,7 @@ def set_array(self, A):
         if A is None:
             self._A = None
             return
-
+        A = self._strip_units(A)
         A = cbook.safe_masked_invalid(A, copy=True)
         if not np.can_cast(A.dtype, float, "same_kind"):
             raise TypeError(f"Image data of dtype {A.dtype} cannot be "
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -728,6 +728,8 @@ def set_data(self, A):
         """
         if isinstance(A, PIL.Image.Image):
             A = pil_to_array(A)  # Needed e.g. to apply png palette.
+
+        A = self._strip_units(A)
         self._A = self._normalize_image_array(A)
         self._imcache = None
         self.stale = True
@@ -1142,6 +1144,7 @@ def set_data(self, x, y, A):
             (M, N) `~numpy.ndarray` or masked array of values to be
             colormapped, or (M, N, 3) RGB array, or (M, N, 4) RGBA array.
         """
+        A = self._strip_units(A)
         A = self._normalize_image_array(A)
         x = np.array(x, np.float32)
         y = np.array(y, np.float32)
@@ -1302,6 +1305,7 @@ def set_data(self, x, y, A):
             - (M, N, 3): RGB array
             - (M, N, 4): RGBA array
         """
+        A = self._strip_units(A)
         A = self._normalize_image_array(A)
         x = np.arange(0., A.shape[1] + 1) if x is None else np.array(x, float).ravel()
         y = np.arange(0., A.shape[0] + 1) if y is None else np.array(y, float).ravel()
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -41,6 +41,9 @@ def __getitem__(self, item):
     def __array__(self):
         return np.asarray(self.magnitude)
 
+    def __len__(self):
+        return len(self.__array__())
+
 
 @pytest.fixture
 def quantity_converter():
@@ -294,3 +297,30 @@ def test_plot_kernel():
     # just a smoketest that fail
     kernel = Kernel([1, 2, 3, 4, 5])
     plt.plot(kernel)
+
+
+@image_comparison(['mappable_units.png'], style="mpl20")
+def test_mappable_units(quantity_converter):
+    # Check that showing an image with units works
+    munits.registry[Quantity] = quantity_converter
+    x, y = np.meshgrid([0, 1], [0, 1])
+    data = Quantity(np.arange(4).reshape(2, 2), 'hours')
+
+    fig, axs = plt.subplots(nrows=2, ncols=2)
+
+    # imshow
+    ax = axs[0, 0]
+    mappable = ax.imshow(data, origin='lower')
+    cbar = fig.colorbar(mappable, ax=ax)
+
+    # pcolor
+    ax = axs[0, 1]
+    mappable = ax.pcolor(x, y, data)
+    fig.colorbar(mappable, ax=ax)
+
+    # pcolormesh + horizontal colorbar
+    ax = axs[1, 0]
+    mappable = ax.pcolormesh(x, y, data)
+    fig.colorbar(mappable, ax=ax, orientation="horizontal")
+
+    axs[1, 1].axis("off")
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -118,16 +118,22 @@ def axisinfo(unit, axis):
 
     @staticmethod
     def default_units(x, axis):
-        """Return the default unit for *x* or ``None`` for the given axis."""
+        """
+        Return the default unit for *x*.
+
+        *axis* may be an `~.axis.Axis` or ``None``.
+        """
         return None
 
     @staticmethod
     def convert(obj, unit, axis):
         """
-        Convert *obj* using *unit* for the specified *axis*.
+        Convert *obj* using *unit*.
 
         If *obj* is a sequence, return the converted sequence.  The output must
         be a sequence of scalars that can be used by the numpy array layer.
+
+        *axis* may be an `~.axis.Axis` or ``None``.
         """
         return obj
 
