Merge pull request matplotlib#14040 from ianthomas23/10167_tricontour_non_finite_z

efiring · web-flow · commit 765d452bdd11 · 2019-05-29T06:42:50.000-10:00
Gracefully handle non-finite z in tricontour (issue matplotlib#10167)
diff --git a/lib/matplotlib/tests/test_triangulation.py b/lib/matplotlib/tests/test_triangulation.py
@@ -1113,3 +1113,27 @@ def test_qhull_large_offset():
     triang = mtri.Triangulation(x, y)
     triang_offset = mtri.Triangulation(x + offset, y + offset)
     assert len(triang.triangles) == len(triang_offset.triangles)
+
+
+def test_tricontour_non_finite_z():
+    # github issue 10167.
+    x = [0, 1, 0, 1]
+    y = [0, 0, 1, 1]
+    triang = mtri.Triangulation(x, y)
+    plt.figure()
+
+    with pytest.raises(ValueError, match='z array must not contain non-finite '
+                                         'values within the triangulation'):
+        plt.tricontourf(triang, [0, 1, 2, np.inf])
+
+    with pytest.raises(ValueError, match='z array must not contain non-finite '
+                                         'values within the triangulation'):
+        plt.tricontourf(triang, [0, 1, 2, -np.inf])
+
+    with pytest.raises(ValueError, match='z array must not contain non-finite '
+                                         'values within the triangulation'):
+        plt.tricontourf(triang, [0, 1, 2, np.nan])
+
+    with pytest.raises(ValueError, match='z must not contain masked points '
+                                         'within the triangulation'):
+        plt.tricontourf(triang, np.ma.array([0, 1, 2, 3], mask=[1, 0, 0, 0]))
diff --git a/lib/matplotlib/tri/tricontour.py b/lib/matplotlib/tri/tricontour.py
@@ -79,12 +79,24 @@ def _contour_args(self, args, kwargs):
             fn = 'contour'
         tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args,
                                                                    **kwargs)
-        z = np.asarray(args[0])
+        z = np.ma.asarray(args[0])
         if z.shape != tri.x.shape:
             raise ValueError('z array must have same length as triangulation x'
                              ' and y arrays')
-        self.zmax = z.max()
-        self.zmin = z.min()
+
+        # z values must be finite, only need to check points that are included
+        # in the triangulation.
+        z_check = z[np.unique(tri.get_masked_triangles())]
+        if np.ma.is_masked(z_check):
+            raise ValueError('z must not contain masked points within the '
+                             'triangulation')
+        if not np.isfinite(z_check).all():
+            raise ValueError('z array must not contain non-finite values '
+                             'within the triangulation')
+
+        z = np.ma.masked_invalid(z, copy=False)
+        self.zmax = float(z_check.max())
+        self.zmin = float(z_check.min())
         if self.logscale and self.zmin <= 0:
             raise ValueError('Cannot %s log of negative values.' % fn)
         self._contour_level_args(z, args[1:])
