Merge pull request #23561 from oscargus/tricleanups

QuLogic · web-flow · commit fa8988965a1d · 2022-08-09T16:47:41.000-04:00
Clean up code in tri
diff --git a/lib/matplotlib/tri/triangulation.py b/lib/matplotlib/tri/triangulation.py
@@ -80,18 +80,14 @@ def __init__(self, x, y, triangles=None, mask=None):
                     f'range 0 <= i < {len(self.x)} but found value '
                     f'{self.triangles.min()}')
 
-        if mask is not None:
-            self.mask = np.asarray(mask, dtype=bool)
-            if self.mask.shape != (self.triangles.shape[0],):
-                raise ValueError('mask array must have same length as '
-                                 'triangles array')
-
         # Underlying C++ object is not created until first needed.
         self._cpp_triangulation = None
 
         # Default TriFinder not created until needed.
         self._trifinder = None
 
+        self.set_mask(mask)
+
     def calculate_plane_coefficients(self, z):
         """
         Calculate plane equation coefficients for all unmasked triangles from
diff --git a/lib/matplotlib/tri/tricontour.py b/lib/matplotlib/tri/tricontour.py
@@ -19,7 +19,7 @@ class TriContourSet(ContourSet):
     def __init__(self, ax, *args, **kwargs):
         """
         Draw triangular grid contour lines or filled regions,
-        depending on whether keyword arg 'filled' is False
+        depending on whether keyword arg *filled* is False
         (default) or True.
 
         The first argument of the initializer must be an `~.axes.Axes`
@@ -51,10 +51,6 @@ def _process_args(self, *args, **kwargs):
         return kwargs
 
     def _contour_args(self, args, kwargs):
-        if self.filled:
-            fn = 'contourf'
-        else:
-            fn = 'contour'
         tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args,
                                                                    **kwargs)
         z = np.ma.asarray(args[0])
@@ -76,7 +72,8 @@ def _contour_args(self, args, kwargs):
         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)
+            func = 'contourf' if self.filled else 'contour'
+            raise ValueError(f'Cannot {func} log of negative values.')
         self._process_contour_level_args(args[1:])
         return (tri, z)
 
diff --git a/lib/matplotlib/tri/triinterpolate.py b/lib/matplotlib/tri/triinterpolate.py
@@ -272,16 +272,15 @@ def gradient(self, x, y):
     gradient.__doc__ = TriInterpolator._docstringgradient
 
     def _interpolate_single_key(self, return_key, tri_index, x, y):
+        _api.check_in_list(['z', 'dzdx', 'dzdy'], return_key=return_key)
         if return_key == 'z':
             return (self._plane_coefficients[tri_index, 0]*x +
                     self._plane_coefficients[tri_index, 1]*y +
                     self._plane_coefficients[tri_index, 2])
         elif return_key == 'dzdx':
             return self._plane_coefficients[tri_index, 0]
-        elif return_key == 'dzdy':
+        else:  # 'dzdy'
             return self._plane_coefficients[tri_index, 1]
-        else:
-            raise ValueError("Invalid return_key: " + return_key)
 
 
 class CubicTriInterpolator(TriInterpolator):
@@ -413,6 +412,7 @@ def __init__(self, triangulation, z, kind='min_E', trifinder=None,
         # Computing eccentricities
         self._eccs = self._compute_tri_eccentricities(self._tris_pts)
         # Computing dof estimations for HCT triangle shape function
+        _api.check_in_list(['user', 'geom', 'min_E'], kind=kind)
         self._dof = self._compute_dof(kind, dz=dz)
         # Loading HCT element
         self._ReferenceElement = _ReducedHCT_Element()
@@ -428,23 +428,22 @@ def gradient(self, x, y):
     gradient.__doc__ = TriInterpolator._docstringgradient
 
     def _interpolate_single_key(self, return_key, tri_index, x, y):
+        _api.check_in_list(['z', 'dzdx', 'dzdy'], return_key=return_key)
         tris_pts = self._tris_pts[tri_index]
         alpha = self._get_alpha_vec(x, y, tris_pts)
         ecc = self._eccs[tri_index]
         dof = np.expand_dims(self._dof[tri_index], axis=1)
         if return_key == 'z':
             return self._ReferenceElement.get_function_values(
                 alpha, ecc, dof)
-        elif return_key in ['dzdx', 'dzdy']:
+        else:  # 'dzdx', 'dzdy'
             J = self._get_jacobian(tris_pts)
             dzdx = self._ReferenceElement.get_function_derivatives(
                 alpha, J, ecc, dof)
             if return_key == 'dzdx':
                 return dzdx[:, 0, 0]
             else:
                 return dzdx[:, 1, 0]
-        else:
-            raise ValueError("Invalid return_key: " + return_key)
 
     def _compute_dof(self, kind, dz=None):
         """
@@ -472,10 +471,8 @@ def _compute_dof(self, kind, dz=None):
             TE = _DOF_estimator_user(self, dz=dz)
         elif kind == 'geom':
             TE = _DOF_estimator_geom(self)
-        elif kind == 'min_E':
+        else:  # 'min_E', checked in __init__
             TE = _DOF_estimator_min_E(self)
-        else:
-            _api.check_in_list(['user', 'geom', 'min_E'], kind=kind)
         return TE.compute_dof_from_df()
 
     @staticmethod
@@ -1510,7 +1507,7 @@ def _roll_vectorized(M, roll_indices, axis):
         for ir in range(r):
             for ic in range(c):
                 M_roll[:, ir, ic] = M[vec_indices, (-roll_indices+ir) % r, ic]
-    elif axis == 1:
+    else:  # 1
         for ir in range(r):
             for ic in range(c):
                 M_roll[:, ir, ic] = M[vec_indices, ir, (-roll_indices+ic) % c]
@@ -1562,15 +1559,15 @@ def _extract_submatrices(M, block_indices, block_size, axis):
     r, c = M.shape
     if axis == 0:
         sh = [block_indices.shape[0], block_size, c]
-    elif axis == 1:
+    else:  # 1
         sh = [block_indices.shape[0], r, block_size]
 
     dt = M.dtype
     M_res = np.empty(sh, dtype=dt)
     if axis == 0:
         for ir in range(block_size):
             M_res[:, ir, :] = M[(block_indices*block_size+ir), :]
-    elif axis == 1:
+    else:  # 1
         for ic in range(block_size):
             M_res[:, :, ic] = M[:, (block_indices*block_size+ic)]
 
diff --git a/lib/matplotlib/tri/tripcolor.py b/lib/matplotlib/tri/tripcolor.py
@@ -123,7 +123,7 @@ def tripcolor(ax, *args, alpha=1.0, norm=None, cmap=None, vmin=None,
                 "are specified at the points, not at the faces.")
         collection = TriMesh(tri, alpha=alpha, array=point_colors,
                              cmap=cmap, norm=norm, **kwargs)
-    else:
+    else:  # 'flat'
         # Vertices of triangles.
         maskedTris = tri.get_masked_triangles()
         verts = np.stack((tri.x[maskedTris], tri.y[maskedTris]), axis=-1)
