Place 3D contourfs midway between levels

dstansby · dstansby · commit 44e04d8d8c79 · 2021-09-27T14:25:19.000+01:00
diff --git a/doc/api/next_api_changes/behavior/21026-DS.rst b/doc/api/next_api_changes/behavior/21026-DS.rst
@@ -0,0 +1,5 @@
+3D contourf polygons placed between levels
+------------------------------------------
+The polygons used in a 3D `~mpl_toolkits.mplot3d.Axes3D.contourf` plot are
+now placed halfway between the contour levels, as each polygon represents the
+location of values that lie between two levels.
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -813,6 +813,8 @@ def __init__(self, ax, *args,
         kwargs = self._process_args(*args, **kwargs)
         self._process_levels()
 
+        self._extend_min = self.extend in ['min', 'both']
+        self._extend_max = self.extend in ['max', 'both']
         if self.colors is not None:
             ncolors = len(self.levels)
             if self.filled:
@@ -821,25 +823,27 @@ def __init__(self, ax, *args,
 
             # Handle the case where colors are given for the extended
             # parts of the contour.
-            extend_min = self.extend in ['min', 'both']
-            extend_max = self.extend in ['max', 'both']
+
             use_set_under_over = False
             # if we are extending the lower end, and we've been given enough
             # colors then skip the first color in the resulting cmap. For the
             # extend_max case we don't need to worry about passing more colors
             # than ncolors as ListedColormap will clip.
-            total_levels = ncolors + int(extend_min) + int(extend_max)
-            if len(self.colors) == total_levels and (extend_min or extend_max):
+            total_levels = (ncolors +
+                            int(self._extend_min) +
+                            int(self._extend_max))
+            if (len(self.colors) == total_levels and
+                    (self._extend_min or self._extend_max)):
                 use_set_under_over = True
-                if extend_min:
+                if self._extend_min:
                     i0 = 1
 
             cmap = mcolors.ListedColormap(self.colors[i0:None], N=ncolors)
 
             if use_set_under_over:
-                if extend_min:
+                if self._extend_min:
                     cmap.set_under(self.colors[0])
-                if extend_max:
+                if self._extend_max:
                     cmap.set_over(self.colors[-1])
 
         self.collections = cbook.silent_list(None)
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2068,12 +2068,32 @@ def add_contour_set(
                 art3d.line_collection_2d_to_3d(linec, z, zdir=zdir)
 
     def add_contourf_set(self, cset, zdir='z', offset=None):
+        self._add_contourf_set(cset, zdir=zdir, offset=offset)
+
+    def _add_contourf_set(self, cset, zdir='z', offset=None):
+        """
+        Returns
+        -------
+        levels : numpy.ndarray
+            Levels at which the filled contours are added.
+        """
         zdir = '-' + zdir
-        for z, linec in zip(cset.levels, cset.collections):
+
+        midpoints = cset.levels[:-1] + np.diff(cset.levels) / 2
+        # Linearly interpolate to get levels for any extensions
+        if cset._extend_min:
+            min_level = cset.levels[0] - np.diff(cset.levels[:2]) / 2
+            midpoints = np.insert(midpoints, 0, min_level)
+        if cset._extend_max:
+            max_level = cset.levels[-1] + np.diff(cset.levels[-2:]) / 2
+            midpoints = np.append(midpoints, max_level)
+
+        for z, linec in zip(midpoints, cset.collections):
             if offset is not None:
                 z = offset
             art3d.poly_collection_2d_to_3d(linec, z, zdir=zdir)
             linec.set_sort_zpos(z)
+        return midpoints
 
     @_preprocess_data()
     def contour(self, X, Y, Z, *args,
@@ -2168,6 +2188,16 @@ def tricontour(self, *args,
         self.auto_scale_xyz(X, Y, Z, had_data)
         return cset
 
+    def _auto_scale_contourf(self, X, Y, Z, zdir, levels, had_data):
+        # Autoscale in the zdir based on the levels added, which are
+        # different from data range if any contour extensions are present
+        dim_vals = {'x': X, 'y': Y, 'z': Z, zdir: levels}
+        # Input data and levels have different sizes, but auto_scale_xyz
+        # expected same-size input, so manually take min/max limits
+        limits = [(np.nanmin(dim_vals[dim]), np.nanmax(dim_vals[dim]))
+                  for dim in ['x', 'y', 'z']]
+        self.auto_scale_xyz(*limits, had_data)
+
     @_preprocess_data()
     def contourf(self, X, Y, Z, *args, zdir='z', offset=None, **kwargs):
         """
@@ -2195,9 +2225,9 @@ def contourf(self, X, Y, Z, *args, zdir='z', offset=None, **kwargs):
 
         jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir)
         cset = super().contourf(jX, jY, jZ, *args, **kwargs)
-        self.add_contourf_set(cset, zdir, offset)
+        levels = self._add_contourf_set(cset, zdir, offset)
 
-        self.auto_scale_xyz(X, Y, Z, had_data)
+        self._auto_scale_contourf(X, Y, Z, zdir, levels, had_data)
         return cset
 
     contourf3D = contourf
@@ -2246,9 +2276,9 @@ def tricontourf(self, *args, zdir='z', offset=None, **kwargs):
         tri = Triangulation(jX, jY, tri.triangles, tri.mask)
 
         cset = super().tricontourf(tri, jZ, *args, **kwargs)
-        self.add_contourf_set(cset, zdir, offset)
+        levels = self._add_contourf_set(cset, zdir, offset)
 
-        self.auto_scale_xyz(X, Y, Z, had_data)
+        self._auto_scale_contourf(X, Y, Z, zdir, levels, had_data)
         return cset
 
     def add_collection3d(self, col, zs=0, zdir='z'):
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/tricontour.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/tricontour.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -155,6 +155,34 @@ def test_contourf3d_fill():
     ax.set_zlim(-1, 1)
 
 
+@pytest.mark.parametrize('extend, levels', [['both', [2, 4, 6]],
+                                            ['min', [2, 4, 6, 8]],
+                                            ['max', [0, 2, 4, 6]]])
+@check_figures_equal(extensions=["png"])
+def test_contourf3d_extend(fig_test, fig_ref, extend, levels):
+    X, Y = np.meshgrid(np.arange(-2, 2, 0.25), np.arange(-2, 2, 0.25))
+    # Z is in the range [0, 8]
+    Z = X**2 + Y**2
+
+    # Manually set the over/under colors to be the end of the colormap
+    cmap = plt.get_cmap('viridis').copy()
+    cmap.set_under(cmap(0))
+    cmap.set_over(cmap(255))
+    # Set vmin/max to be the min/max values plotted on the reference image
+    kwargs = {'vmin': 1, 'vmax': 7, 'cmap': cmap}
+
+    ax_ref = fig_ref.add_subplot(projection='3d')
+    ax_ref.contourf(X, Y, Z, levels=[0, 2, 4, 6, 8], **kwargs)
+
+    ax_test = fig_test.add_subplot(projection='3d')
+    ax_test.contourf(X, Y, Z, levels, extend=extend, **kwargs)
+
+    for ax in [ax_ref, ax_test]:
+        ax.set_xlim(-2, 2)
+        ax.set_ylim(-2, 2)
+        ax.set_zlim(-10, 10)
+
+
 @mpl3d_image_comparison(['tricontour.png'], tol=0.02)
 def test_tricontour():
     fig = plt.figure()
