diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index dba6f9013df1..d4eef7ca77e6 100644 --- 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 new file mode 100644 index 000000000000..df7f1ebdf476 Binary files /dev/null and b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked.png differ 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 new file mode 100644 index 000000000000..5524fea4537b Binary files /dev/null and b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/surface3d_masked_strides.png differ diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py index ae8e91ddc24d..ebe801cb4622 100644 --- 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()