ENH: allow pcolor to plot all data is x, y not N+1

jklymak · jklymak · commit d96916fa24e2 · 2019-12-30T11:17:04.000-08:00
diff --git a/doc/users/next_whats_new/2019-12-18-pcolordropdata.rst b/doc/users/next_whats_new/2019-12-18-pcolordropdata.rst
@@ -0,0 +1,17 @@
+Pcolor and Pcolormesh now have *dropdata* kwarg and rcParam
+-----------------------------------------------------------
+
+Previously `.axes.Axes.pcolor` and  `.axes.Axes.pcolormesh` handled
+the situation where *x* and *y* have the same (respective) size as *Z* by
+dropping the last row and column of *Z*, and *x* and *y* are regarded as the
+edges of the remaining rows and columns in *Z*.  However, most users probably
+really want *x* and *y* centered on the rows and columns of *Z*, so if
+they specify *dropdata* as True, both methods will now linearly interpolate to
+get the edges of the bins, and *x* and *y* will specify the (linear) center of
+each gridcell in the pseudocolor plot.
+
+Users can also specify this by the new :rc:`pcolor.dropdata` in their
+``.matplotlibrc`` or via `.rcParams`.
+
+See :doc:`pcolormesh </gallery/images_contours_and_fields/pcolormesh_levels>`
+for an example.  
diff --git a/examples/images_contours_and_fields/pcolormesh_levels.py b/examples/images_contours_and_fields/pcolormesh_levels.py
@@ -1,11 +1,11 @@
 """
 ==========
-pcolormesh
+Pcolormesh
 ==========
 
-Shows how to combine Normalization and Colormap instances to draw "levels" in
-`~.axes.Axes.pcolor`, `~.axes.Axes.pcolormesh` and `~.axes.Axes.imshow` type
-plots in a similar way to the levels keyword argument to contour/contourf.
+`.axes.Axes.pcolormesh` allows you to generate 2-D image-style plots.  Note it
+is somewhat faster than `~.axes.Axes.pcolor`.
+
 """
 
 import matplotlib
@@ -14,6 +14,66 @@
 from matplotlib.ticker import MaxNLocator
 import numpy as np
 
+###############################################################################
+# Basic Pcolormesh
+# ----------------
+#
+# We usually specify a pcolormesh by defining the edge of quadrilaterals and
+# the value of the quadrilateral.  Note that here *x* and *y* each have one
+# extra element than Z in the respective dimension.
+
+np.random.seed(19680801)
+Z = np.random.rand(6, 10)
+x = np.arange(-0.5, 10, 1)  # len = 11
+y = np.arange(4.5, 11, 1)  # len = 7
+
+fig, ax = plt.subplots()
+ax.pcolormesh(x, y, Z)
+
+###############################################################################
+# Non-rectilinear Pcolormesh
+# --------------------------
+#
+# Note that we can also specify matrices for *X* and *Y* and have
+# non-rectilinear quadrilaterals.
+
+x = np.arange(-0.5, 10, 1)  # len = 11
+y = np.arange(4.5, 11, 1)  # len = 7
+X, Y = np.meshgrid(x, y)
+X = X + 0.2 * Y  # tilt the co-ordinates.
+Y = Y + 0.3 * X
+
+fig, ax = plt.subplots()
+ax.pcolormesh(X, Y, Z)
+
+###############################################################################
+# Centered Co-ordinates
+# ---------------------
+#
+# Often a user wants to pass *X* and *Y* with the same sizes as *Z* to
+# `.axes.Axes.pcolormesh`.  Matplotlib will either drop the last row and
+# column of *Z* (default, h=istorically compatible with Matlab), or if
+# ``dropdata=True`` assume the user wanted *X* and *Y* centered on the
+# quadrilateral and linearly interpolate to get the edges.
+
+x = np.arange(10)  # len = 10
+y = np.arange(6)  # len = 6
+X, Y = np.meshgrid(x, y)
+
+fig, axs = plt.subplots(2, 1, sharex=True, sharey=True)
+axs[0].pcolormesh(X, Y, Z, vmin=np.min(Z), vmax=np.max(Z))
+axs[0].set_title('dropdata=True (default)')
+axs[1].pcolormesh(X, Y, Z, vmin=np.min(Z), vmax=np.max(Z), dropdata=False)
+axs[1].set_title('dropdata=False')
+
+###############################################################################
+# Making levels using Norms
+# -------------------------
+#
+# Shows how to combine Normalization and Colormap instances to draw
+# "levels" in `.axes.Axes.pcolor`, `.axes.Axes.pcolormesh`
+# and `.axes.Axes.imshow` type plots in a similar
+# way to the levels keyword argument to contour/contourf.
 
 # make these smaller to increase the resolution
 dx, dy = 0.05, 0.05
@@ -54,8 +114,6 @@
 # don't overlap
 fig.tight_layout()
 
-plt.show()
-
 #############################################################################
 #
 # ------------
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5603,7 +5603,7 @@ def imshow(self, X, cmap=None, norm=None, aspect=None,
         return im
 
     @staticmethod
-    def _pcolorargs(funcname, *args, allmatch=False):
+    def _pcolorargs(funcname, *args, allmatch=False, dropdata=True):
         # If allmatch is True, then the incoming X, Y, C must have matching
         # dimensions, taking into account that X and Y can be 1-D rather than
         # 2-D.  This perfect match is required for Gouraud shading.  For flat
@@ -5665,14 +5665,34 @@ def _pcolorargs(funcname, *args, allmatch=False):
                 raise TypeError('Dimensions of C %s are incompatible with'
                                 ' X (%d) and/or Y (%d); see help(%s)' % (
                                     C.shape, Nx, Ny, funcname))
-            C = C[:Ny - 1, :Nx - 1]
+
+            if dropdata:
+                C = C[:Ny - 1, :Nx - 1]
+            else:
+                def _interp_grid(X):
+                    # helper for below
+                    dX = np.diff(X, axis=1)/2.
+                    X = np.hstack((X[:, [0]] - dX[:, [0]],
+                                    X[:, :-1] + dX,
+                                    X[:, [-1]] + dX[:, [-1]])
+                                  )
+                    return X
+
+                if ncols == Nx:
+                    X = _interp_grid(X)
+                    Y = _interp_grid(Y)
+
+                if nrows == Ny:
+                    X = _interp_grid(X.T).T
+                    Y = _interp_grid(Y.T).T
+
         C = cbook.safe_masked_invalid(C)
         return X, Y, C
 
     @_preprocess_data()
     @docstring.dedent_interpd
     def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
-               vmax=None, **kwargs):
+               vmax=None, dropdata=None, **kwargs):
         r"""
         Create a pseudocolor plot with a non-regular rectangular grid.
 
@@ -5686,7 +5706,9 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
             ``pcolor()`` can be very slow for large arrays. In most
             cases you should use the similar but much faster
-            `~.Axes.pcolormesh` instead. See there for a discussion of the
+            `~.Axes.pcolormesh` instead. See
+            :ref:`Differences between pcolor() and pcolormesh()
+            <differences-pcolor-pcolormesh>` for a discussion of the
             differences.
 
         Parameters
@@ -5711,7 +5733,8 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
             The dimensions of *X* and *Y* should be one greater than those of
             *C*. Alternatively, *X*, *Y* and *C* may have equal dimensions, in
-            which case the last row and column of *C* will be ignored.
+            which case the last row and column of *C* will be ignored if
+            *dropdata* is True (see below).
 
             If *X* and/or *Y* are 1-D arrays or column vectors they will be
             expanded as needed into the appropriate 2-D arrays, making a
@@ -5751,6 +5774,13 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         snap : bool, default: False
             Whether to snap the mesh to pixel boundaries.
 
+        dropdata : bool, default: :rc:`pcolor.dropdata`
+            If True (default), and *X* and *Y* are the same size as C in their
+            respective dimensions, drop the last element of C in both
+            dimensions.  If False, *X* and *Y* are assumed to be the midpoints
+            of the quadrilaterals, and their edges calculated by linear
+            interpolation.
+
         Returns
         -------
         collection : `matplotlib.collections.Collection`
@@ -5801,12 +5831,19 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         ``pcolor()`` displays all columns of *C* if *X* and *Y* are not
         specified, or if *X* and *Y* have one more column than *C*.
         If *X* and *Y* have the same number of columns as *C* then the last
-        column of *C* is dropped. Similarly for the rows.
+        column of *C* is dropped if *dropdata* is True. Similarly for the rows.
+        If *dropdata* is False, then *X* and *Y* are assumed to be at the
+        middle of the quadrilaterals, and the edges of the quadrilaterals are
+        linearly interpolated.
 
-        Note: This behavior is different from MATLAB's ``pcolor()``, which
+        This behavior is different from MATLAB's ``pcolor()``, which
         always discards the last row and column of *C*.
         """
-        X, Y, C = self._pcolorargs('pcolor', *args, allmatch=False)
+        if dropdata is None:
+            dropdata = rcParams['pcolor.dropdata']
+
+        X, Y, C = self._pcolorargs('pcolor', *args, dropdata=dropdata,
+                                   allmatch=False)
         Ny, Nx = X.shape
 
         # unit conversion allows e.g. datetime objects as axis values
@@ -5903,7 +5940,8 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
     @_preprocess_data()
     @docstring.dedent_interpd
     def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
-                   vmax=None, shading='flat', antialiased=False, **kwargs):
+                   vmax=None, shading='flat', antialiased=False,
+                   dropdata=None, **kwargs):
         """
         Create a pseudocolor plot with a non-regular rectangular grid.
 
@@ -5913,9 +5951,9 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
         *X* and *Y* can be used to specify the corners of the quadrilaterals.
 
-        .. note::
+        .. hint::
 
-           `~Axes.pcolormesh` is similar to `~Axes.pcolor`. It's much faster
+           `~Axes.pcolormesh` is similar to `~Axes.pcolor`. It is much faster
            and preferred in most cases. For a detailed discussion on the
            differences see :ref:`Differences between pcolor() and pcolormesh()
            <differences-pcolor-pcolormesh>`.
@@ -5942,7 +5980,8 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
             The dimensions of *X* and *Y* should be one greater than those of
             *C*. Alternatively, *X*, *Y* and *C* may have equal dimensions, in
-            which case the last row and column of *C* will be ignored.
+            which case the last row and column of *C* will be ignored, unless
+            *dropdata* is *False* (see below).
 
             If *X* and/or *Y* are 1-D arrays or column vectors they will be
             expanded as needed into the appropriate 2-D arrays, making a
@@ -5991,6 +6030,13 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         snap : bool, default: False
             Whether to snap the mesh to pixel boundaries.
 
+        dropdata : bool, default: :rc:`pcolor.dropdata`
+            If True (default), and *X* and *Y* are the same size as C in their
+            respective dimensions, drop the last element of C in both
+            dimensions.  If False, *X* and *Y* are assumed to be the midpoints
+            of the quadrilaterals, and their edges calculated by linear
+            interpolation.
+
         Returns
         -------
         mesh : `matplotlib.collections.QuadMesh`
@@ -6062,7 +6108,11 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
         allmatch = (shading == 'gouraud')
 
-        X, Y, C = self._pcolorargs('pcolormesh', *args, allmatch=allmatch)
+        if dropdata is None:
+            dropdata = rcParams['pcolor.dropdata']
+
+        X, Y, C = self._pcolorargs('pcolormesh', *args,
+                allmatch=allmatch, dropdata=dropdata)
         Ny, Nx = X.shape
         X = X.ravel()
         Y = Y.ravel()
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2515,11 +2515,11 @@ def minorticks_on():
 @docstring.copy(Axes.pcolor)
 def pcolor(
         *args, alpha=None, norm=None, cmap=None, vmin=None,
-        vmax=None, data=None, **kwargs):
+        vmax=None, dropdata=None, data=None, **kwargs):
     __ret = gca().pcolor(
         *args, alpha=alpha, norm=norm, cmap=cmap, vmin=vmin,
-        vmax=vmax, **({"data": data} if data is not None else {}),
-        **kwargs)
+        vmax=vmax, dropdata=dropdata, **({"data": data} if data is not
+        None else {}), **kwargs)
     sci(__ret)
     return __ret
 
@@ -2528,12 +2528,13 @@ def pcolor(
 @docstring.copy(Axes.pcolormesh)
 def pcolormesh(
         *args, alpha=None, norm=None, cmap=None, vmin=None,
-        vmax=None, shading='flat', antialiased=False, data=None,
-        **kwargs):
+        vmax=None, shading='flat', antialiased=False, dropdata=None,
+        data=None, **kwargs):
     __ret = gca().pcolormesh(
         *args, alpha=alpha, norm=norm, cmap=cmap, vmin=vmin,
         vmax=vmax, shading=shading, antialiased=antialiased,
-        **({"data": data} if data is not None else {}), **kwargs)
+        dropdata=dropdata, **({"data": data} if data is not None else
+        {}), **kwargs)
     sci(__ret)
     return __ret
 
diff --git a/lib/matplotlib/rcsetup.py b/lib/matplotlib/rcsetup.py
@@ -1030,6 +1030,9 @@ def validate_webagg_address(s):
     # marker props
     'markers.fillstyle': ['full', validate_fillstyle],
 
+    ## pcolor(mesh) props:
+    'pcolor.dropdata': [True, validate_bool],
+
     ## patch props
     'patch.linewidth':   [1.0, validate_float],     # line width in points
     'patch.edgecolor':   ['black', validate_color],
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1268,6 +1268,52 @@ def test_pcolorargs():
         ax.pcolormesh(x, y, Z[:-1, :-1])
 
 
+@check_figures_equal(extensions=["png"])
+def test_pcolornodropdata(fig_test, fig_ref):
+    ax = fig_test.subplots()
+    x = np.arange(0, 10)
+    y = np.arange(0, 3)
+    np.random.seed(19680801)
+    Z = np.random.randn(2, 9)
+    ax.pcolormesh(x, y, Z)
+
+    ax = fig_ref.subplots()
+    x2 = x[:-1] + np.diff(x) / 2
+    y2 = y[:-1] + np.diff(y) / 2
+    ax.pcolormesh(x2, y2, Z, dropdata=False)
+
+
+@check_figures_equal(extensions=["png"])
+def test_pcolornodropdatarc(fig_test, fig_ref):
+    matplotlib.rcParams['pcolor.dropdata'] = False
+    ax = fig_test.subplots()
+    x = np.arange(0, 10)
+    y = np.arange(0, 3)
+    np.random.seed(19680801)
+    Z = np.random.randn(2, 9)
+    ax.pcolormesh(x, y, Z)
+
+    ax = fig_ref.subplots()
+    x2 = x[:-1] + np.diff(x) / 2
+    y2 = y[:-1] + np.diff(y) / 2
+    ax.pcolormesh(x2, y2, Z)
+
+
+@check_figures_equal(extensions=["png"])
+def test_pcolordropdata(fig_test, fig_ref):
+    ax = fig_test.subplots()
+    x = np.arange(0, 10)
+    y = np.arange(0, 4)
+    np.random.seed(19680801)
+    Z = np.random.randn(3, 9)
+    ax.pcolormesh(x[:-1], y[:-1], Z[:-1, :-1])
+
+    ax = fig_ref.subplots()
+    x2 = x[:-1]
+    y2 = y[:-1]
+    ax.pcolormesh(x2, y2, Z)
+
+
 @image_comparison(['canonical'])
 def test_canonical():
     fig, ax = plt.subplots()
diff --git a/matplotlibrc.template b/matplotlibrc.template
@@ -135,6 +135,7 @@
 
 #markers.fillstyle : full  ## {full, left, right, bottom, top, none}
 
+# pcolor.dropdata : True
 
 ## ***************************************************************************
 ## * PATCHES                                                                 *
