matplotlib · QuLogic · Aug 12, 2021 · Jul 16, 2021 · Jul 23, 2021 · Jul 26, 2021
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1574,12 +1574,8 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
 
         if Z.ndim != 2:
             raise ValueError("Argument Z must be 2-dimensional.")
-        if np.any(np.isnan(Z)):
-            _api.warn_external(
-                "Z contains NaN values. This may result in rendering "
-                "artifacts.")
 
-        # TODO: Support masked arrays
+        Z = cbook._to_unmasked_float_array(Z)
         X, Y, Z = np.broadcast_arrays(X, Y, Z)
         rows, cols = Z.shape
 
@@ -1655,6 +1651,27 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
                     if fcolors is not None:
                         colset.append(fcolors[rs][cs])
 
+        # In cases where there are NaNs in the data (possibly from masked
+        # arrays), artifacts can be introduced. Here check whether NaNs exist
+        # and remove the entries if so
+        if not isinstance(polys, np.ndarray) or np.isnan(polys).any():
+            new_polys = []
+            new_colset = []
+
+            # Depending on fcolors, colset is either an empty list or has as
+            # many elements as polys. In the former case new_colset results in
+            # a list with None entries, that is discarded later.
+            for p, col in itertools.zip_longest(polys, colset):
+                new_poly = np.array(p)[~np.isnan(p).any(axis=1)]
+                if len(new_poly):
+                    new_polys.append(new_poly)
+                    new_colset.append(col)
+
+            # Replace previous polys and, if fcolors is not None, colset
+            polys = new_polys
+            if fcolors is not None:
+                colset = new_colset
+
         # note that the striding causes some polygons to have more coordinates
         # than others
         polyc = art3d.Poly3DCollection(polys, *args, **kwargs)

diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked.png
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked_strides.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked_strides.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -420,6 +420,45 @@ def test_surface3d_shaded():
     ax.set_zlim(-1.01, 1.01)
 
 
+@mpl3d_image_comparison(['surface3d_masked.png'])
+def test_surface3d_masked():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    y = [1, 2, 3, 4, 5, 6, 7, 8]
+
+    x, y = np.meshgrid(x, y)
+    matrix = np.array(
+        [
+            [-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+            [-1, 1, 2, 3, 4, 4, 4, 3, 2, 1, 1],
+            [-1, -1., 4, 5, 6, 8, 6, 5, 4, 3, -1.],
+            [-1, -1., 7, 8, 11, 12, 11, 8, 7, -1., -1.],
+            [-1, -1., 8, 9, 10, 16, 10, 9, 10, 7, -1.],
+            [-1, -1., -1., 12, 16, 20, 16, 12, 11, -1., -1.],
+            [-1, -1., -1., -1., 22, 24, 22, 20, 18, -1., -1.],
+            [-1, -1., -1., -1., -1., 28, 26, 25, -1., -1., -1.],
+        ]
+    )
+    z = np.ma.masked_less(matrix, 0)
+    norm = mcolors.Normalize(vmax=z.max(), vmin=z.min())
+    colors = plt.get_cmap("plasma")(norm(z))
+    ax.plot_surface(x, y, z, facecolors=colors)
+    ax.view_init(30, -80)
+
+
+@mpl3d_image_comparison(['surface3d_masked_strides.png'])
+def test_surface3d_masked_strides():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    x, y = np.mgrid[-6:6.1:1, -6:6.1:1]
+    z = np.ma.masked_less(x * y, 2)
+
+    ax.plot_surface(x, y, z, rstride=4, cstride=4)
+    ax.view_init(60, -45)
+
+
 @mpl3d_image_comparison(['text3d.png'], remove_text=False)
 def test_text3d():
     fig = plt.figure()