Merge pull request #26253 from LemonBoy/inf-surf

tacaswell · web-flow · commit 83b3fef409e2 · 2023-08-03T18:56:21.000-04:00
Filter out inf values in plot_surface
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1708,18 +1708,18 @@ def plot_surface(self, X, Y, Z, *, 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():
+        # In cases where there are non-finite values in the data (possibly NaNs from
+        # masked arrays), artifacts can be introduced. Here check whether such values
+        # are present and remove them.
+        if not isinstance(polys, np.ndarray) or not np.isfinite(polys).all():
             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)]
+                new_poly = np.array(p)[np.isfinite(p).all(axis=1)]
                 if len(new_poly):
                     new_polys.append(new_poly)
                     new_colset.append(col)
diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_zsort_inf.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_zsort_inf.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -2244,3 +2244,16 @@ def test_scatter_masked_color():
     # Assert sizes' equality
     assert len(path3d.get_offsets()) ==\
            len(super(type(path3d), path3d).get_facecolors())
+
+
+@mpl3d_image_comparison(['surface3d_zsort_inf.png'], style='mpl20')
+def test_surface3d_zsort_inf():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    x, y = np.mgrid[-2:2:0.1, -2:2:0.1]
+    z = np.sin(x)**2 + np.cos(y)**2
+    z[x.shape[0] // 2:, x.shape[1] // 2:] = np.inf
+
+    ax.plot_surface(x, y, z, cmap='jet')
+    ax.view_init(elev=45, azim=145)
