Add unit conversion for pcolor methods

dstansby · dstansby · commit d9b709d29c6d · 2024-01-30T18:30:10.000Z
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -8,6 +8,7 @@
 from numpy import ma
 
 import matplotlib as mpl
+import matplotlib.cm as cm
 import matplotlib.category  # Register category unit converter as side effect.
 import matplotlib.cbook as cbook
 import matplotlib.collections as mcoll
@@ -5795,12 +5796,28 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
         self.add_image(im)
         return im
 
+    @staticmethod
+    def _convert_C_units(C):
+        """
+        Remove any units attached to C, and return the units and converter used to do
+        the conversion.
+        """
+        sm = cm.ScalarMappable()
+        C = sm._strip_units(C)
+        converter = sm._converter
+        units = sm._units
+
+        C = np.asanyarray(C)
+        C = cbook.safe_masked_invalid(C, copy=True)
+        return C, units, converter
+
     def _pcolorargs(self, funcname, *args, shading='auto', **kwargs):
         # - create X and Y if not present;
         # - reshape X and Y as needed if they are 1-D;
         # - check for proper sizes based on `shading` kwarg;
         # - reset shading if shading='auto' to flat or nearest
         #   depending on size;
+        # - if C has units, get the converter
 
         _valid_shading = ['gouraud', 'nearest', 'flat', 'auto']
         try:
@@ -5812,19 +5829,19 @@ def _pcolorargs(self, funcname, *args, shading='auto', **kwargs):
             shading = 'auto'
 
         if len(args) == 1:
-            C = np.asanyarray(args[0])
+            C, units, converter = self._convert_C_units(args[0])
             nrows, ncols = C.shape[:2]
             if shading in ['gouraud', 'nearest']:
                 X, Y = np.meshgrid(np.arange(ncols), np.arange(nrows))
             else:
                 X, Y = np.meshgrid(np.arange(ncols + 1), np.arange(nrows + 1))
                 shading = 'flat'
             C = cbook.safe_masked_invalid(C, copy=True)
-            return X, Y, C, shading
+            return X, Y, C, shading, units, converter
 
         if len(args) == 3:
             # Check x and y for bad data...
-            C = np.asanyarray(args[2])
+            C, units, converter = self._convert_C_units(args[2])
             # unit conversion allows e.g. datetime objects as axis values
             X, Y = args[:2]
             X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
@@ -5905,7 +5922,7 @@ def _interp_grid(X):
                 shading = 'flat'
 
         C = cbook.safe_masked_invalid(C, copy=True)
-        return X, Y, C, shading
+        return X, Y, C, shading, units, converter
 
     @_preprocess_data()
     @_docstring.dedent_interpd
@@ -6057,8 +6074,9 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         if shading is None:
             shading = mpl.rcParams['pcolor.shading']
         shading = shading.lower()
-        X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading,
-                                            kwargs=kwargs)
+        X, Y, C, shading, units, converter = self._pcolorargs(
+            'pcolor', *args, shading=shading, kwargs=kwargs
+        )
         linewidths = (0.25,)
         if 'linewidth' in kwargs:
             kwargs['linewidths'] = kwargs.pop('linewidth')
@@ -6094,6 +6112,8 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
 
         collection = mcoll.PolyQuadMesh(
             coords, array=C, cmap=cmap, norm=norm, alpha=alpha, **kwargs)
+        collection._units = units
+        collection._converter = converter
         collection._scale_norm(norm, vmin, vmax)
 
         # Transform from native to data coordinates?
@@ -6313,15 +6333,18 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         shading = shading.lower()
         kwargs.setdefault('edgecolors', 'none')
 
-        X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
-                                            shading=shading, kwargs=kwargs)
+        X, Y, C, shading, units, converter = self._pcolorargs(
+            'pcolormesh', *args, shading=shading, kwargs=kwargs
+        )
         coords = np.stack([X, Y], axis=-1)
 
         kwargs.setdefault('snap', mpl.rcParams['pcolormesh.snap'])
 
         collection = mcoll.QuadMesh(
             coords, antialiased=antialiased, shading=shading,
             array=C, cmap=cmap, norm=norm, alpha=alpha, **kwargs)
+        collection._units = units
+        collection._converter = converter
         collection._scale_norm(norm, vmin, vmax)
 
         coords = coords.reshape(-1, 2)  # flatten the grid structure; keep x, y
