Allow unit-ful image data

dstansby · dstansby · commit d7f73abfeb4f · 2021-09-01T13:16:01.000+01:00
diff --git a/doc/api/next_api_changes/behavior/xxxxx-DS.rst b/doc/api/next_api_changes/behavior/xxxxx-DS.rst
@@ -0,0 +1,8 @@
+Unit converters can now support units on images
+-----------------------------------------------
+
+Matplotlib now supports using `~.axes.Axes.imshow` to plot data with units.
+For this to be supported by third-party `~.units.ConversionInterface`s,
+the `~.units.ConversionInterface.default_units` and
+`~.units.ConversionInterface.convert` must allow for the *axis* argument to be
+a ``matplotlib.images._ImageBase`` object.
diff --git a/doc/users/next_whats_new/imshow_units.rst b/doc/users/next_whats_new/imshow_units.rst
@@ -0,0 +1,4 @@
+`~.axes.Axes.imshow` can now be used with unit-ful data
+-------------------------------------------------------
+
+`~.axes.Axes.imshow` can now be used with data that has units attached to it.
diff --git a/examples/units/basic_units.py b/examples/units/basic_units.py
@@ -346,19 +346,16 @@ def convert(val, unit, axis):
         if np.iterable(val):
             if isinstance(val, np.ma.MaskedArray):
                 val = val.astype(float).filled(np.nan)
-            out = np.empty(len(val))
-            for i, thisval in enumerate(val):
-                if np.ma.is_masked(thisval):
-                    out[i] = np.nan
-                else:
-                    try:
-                        out[i] = thisval.convert_to(unit).get_value()
-                    except AttributeError:
-                        out[i] = thisval
+            out = np.empty(np.shape(val))
+            masked_mask = np.ma.getmaskarray(val)
+            out[masked_mask] = np.nan
+            out[~masked_mask] = val[~masked_mask].convert_to(unit).get_value()
             return out
+
         if np.ma.is_masked(val):
             return np.nan
         else:
+            # Scalar
             return val.convert_to(unit).get_value()
 
     @staticmethod
diff --git a/examples/units/units_image.py b/examples/units/units_image.py
@@ -0,0 +1,21 @@
+"""
+=================
+Images with units
+=================
+Plotting images with units.
+
+.. only:: builder_html
+
+   This example requires :download:`basic_units.py <basic_units.py>`
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+from basic_units import secs
+
+data = np.array([[1, 2],
+                 [3, 4]]) * secs
+
+fig, ax = plt.subplots()
+image = ax.imshow(data)
+fig.colorbar(image)
+plt.show()
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -470,8 +470,6 @@ def __init__(self, ax, mappable=None, *, cmap=None,
             linewidths=[0.5 * mpl.rcParams['axes.linewidth']])
         self.ax.add_collection(self.dividers)
 
-        self.locator = None
-        self.formatter = None
         self.__scale = None  # linear, log10 for now.  Hopefully more?
 
         if ticklocation == 'auto':
@@ -481,15 +479,23 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         self.set_label(label)
         self._reset_locator_formatter_scale()
 
+        self.locator = None
         if np.iterable(ticks):
             self.locator = ticker.FixedLocator(ticks, nbins=len(ticks))
         else:
             self.locator = ticks    # Handle default in _ticker()
 
+        self.formatter = None
         if isinstance(format, str):
             self.formatter = ticker.FormatStrFormatter(format)
-        else:
-            self.formatter = format  # Assume it is a Formatter or None
+        elif format is not None:
+            self.formatter = format  # Assume it is a Formatter
+        elif hasattr(mappable, 'converter') and hasattr(mappable, 'units'):
+            # Set from mappable if it has a converter and units
+            info = mappable.converter.axisinfo(mappable.units, self._long_axis)
+            if info is not None and info.majfmt is not None:
+                self.formatter = info.majfmt
+
         self.draw_all()
 
         if isinstance(mappable, contour.ContourSet) and not mappable.filled:
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -18,6 +18,7 @@
 import matplotlib.colors as mcolors
 import matplotlib.cm as cm
 import matplotlib.cbook as cbook
+import matplotlib.units as munits
 # For clarity, names from _image are given explicitly in this module:
 import matplotlib._image as _image
 # For user convenience, the names from _image are also imported into
@@ -696,6 +697,7 @@ 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._convert_units(A)
         self._A = cbook.safe_masked_invalid(A, copy=True)
 
         if (self._A.dtype != np.uint8 and
@@ -733,6 +735,30 @@ def set_data(self, A):
         self._rgbacache = None
         self.stale = True
 
+    def _convert_units(self, A):
+        # Take the first element since units expects a 1D sequence, not 2D
+        converter = munits.registry.get_converter(A[0])
+        if converter is None:
+            return A
+
+        try:
+            units = converter.default_units(A, self)
+        except Exception as e:
+            raise RuntimeError(
+                f'{converter} failed when trying to return the default units '
+                f'for this image. This may be because {converter} has not '
+                'implemented support for images in the default_units() method.'
+            ) from e
+
+        try:
+            return converter.convert(A, units, self)
+        except Exception as e:
+            raise RuntimeError(
+                f'{converter} failed when trying to convert the units '
+                f'for this image. This may be because {converter} has not '
+                'implemented support for images in the convert() method.'
+            ) from e
+
     def set_array(self, A):
         """
         Retained for backwards compatibility - use set_data instead.
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -118,7 +118,12 @@ 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* or ``None``.
+
+        *axis* can be either an `Axis` or an ``_ImageBase`` (if units of a 2D
+        image are being converted).
+        """
         return None
 
     @staticmethod
@@ -128,6 +133,9 @@ def convert(obj, unit, axis):
 
         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* can be either an `Axis` or an ``_ImageBase`` (if units of a 2D
+        image are being converted).
         """
         return obj
 

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +`~.axes.Axes.imshow` can now be used with unit-ful data
 +-------------------------------------------------------
++
 +`~.axes.Axes.imshow` can now be used with data that has units attached to it.