Merge pull request #18549 from jklymak/fix-pcolorarg-convert-before-interp

tacaswell · web-flow · commit a010c54a5943 · 2020-09-23T13:22:26.000-04:00
FIX: unit-convert pcolorargs before interpolating
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5422,8 +5422,7 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         self.add_image(im)
         return im
 
-    @staticmethod
-    def _pcolorargs(funcname, *args, shading='flat'):
+    def _pcolorargs(self, funcname, *args, shading='flat', **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;
@@ -5454,6 +5453,9 @@ def _pcolorargs(funcname, *args, shading='flat'):
             # Check x and y for bad data...
             C = np.asanyarray(args[2])
             X, Y = [cbook.safe_masked_invalid(a) for a in args[:2]]
+            # unit conversion allows e.g. datetime objects as axis values
+            X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
+
             if funcname == 'pcolormesh':
                 if np.ma.is_masked(X) or np.ma.is_masked(Y):
                     raise ValueError(
@@ -5702,12 +5704,10 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         if shading is None:
             shading = rcParams['pcolor.shading']
         shading = shading.lower()
-        X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading)
+        X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading,
+                                            kwargs=kwargs)
         Ny, Nx = X.shape
 
-        # unit conversion allows e.g. datetime objects as axis values
-        X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
-
         # convert to MA, if necessary.
         C = ma.asarray(C)
         X = ma.asarray(X)
@@ -5976,12 +5976,10 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         kwargs.setdefault('edgecolors', 'None')
 
         X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
-                                            shading=shading)
+                                            shading=shading, kwargs=kwargs)
         Ny, Nx = X.shape
         X = X.ravel()
         Y = Y.ravel()
-        # unit conversion allows e.g. datetime objects as axis values
-        X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
 
         # convert to one dimensional arrays
         C = C.ravel()
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1201,6 +1201,22 @@ def test_pcolornearest(fig_test, fig_ref):
     ax.pcolormesh(x2, y2, Z, shading='nearest')
 
 
+@check_figures_equal(extensions=["png"])
+def test_pcolornearestunits(fig_test, fig_ref):
+    ax = fig_test.subplots()
+    x = [datetime.datetime.fromtimestamp(x * 3600) for x in range(10)]
+    y = np.arange(0, 3)
+    np.random.seed(19680801)
+    Z = np.random.randn(2, 9)
+    ax.pcolormesh(x, y, Z, shading='flat')
+
+    ax = fig_ref.subplots()
+    # specify the centers
+    x2 = [datetime.datetime.fromtimestamp((x + 0.5) * 3600) for x in range(9)]
+    y2 = y[:-1] + np.diff(y) / 2
+    ax.pcolormesh(x2, y2, Z, shading='nearest')
+
+
 @check_figures_equal(extensions=["png"])
 def test_pcolordropdata(fig_test, fig_ref):
     ax = fig_test.subplots()
