Replace np.compress by boolean indexing.

anntzer · anntzer · commit d3cd57a20cfe · 2019-02-08T00:06:06.000+01:00
It's shorter, better known, and works directly on 2D inputs without
fiddling with ravel()ing (for the changes for pcolor).
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5814,23 +5814,21 @@ def pcolor(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         # don't plot if C or any of the surrounding vertices are masked.
         mask = ma.getmaskarray(C) + xymask
 
-        compress = np.compress
-
-        ravelmask = (mask == 0).ravel()
-        X1 = compress(ravelmask, ma.filled(X[:-1, :-1]).ravel())
-        Y1 = compress(ravelmask, ma.filled(Y[:-1, :-1]).ravel())
-        X2 = compress(ravelmask, ma.filled(X[1:, :-1]).ravel())
-        Y2 = compress(ravelmask, ma.filled(Y[1:, :-1]).ravel())
-        X3 = compress(ravelmask, ma.filled(X[1:, 1:]).ravel())
-        Y3 = compress(ravelmask, ma.filled(Y[1:, 1:]).ravel())
-        X4 = compress(ravelmask, ma.filled(X[:-1, 1:]).ravel())
-        Y4 = compress(ravelmask, ma.filled(Y[:-1, 1:]).ravel())
+        unmask = ~mask
+        X1 = ma.filled(X[:-1, :-1])[unmask]
+        Y1 = ma.filled(Y[:-1, :-1])[unmask]
+        X2 = ma.filled(X[1:, :-1])[unmask]
+        Y2 = ma.filled(Y[1:, :-1])[unmask]
+        X3 = ma.filled(X[1:, 1:])[unmask]
+        Y3 = ma.filled(Y[1:, 1:])[unmask]
+        X4 = ma.filled(X[:-1, 1:])[unmask]
+        Y4 = ma.filled(Y[:-1, 1:])[unmask]
         npoly = len(X1)
 
         xy = np.stack([X1, Y1, X2, Y2, X3, Y3, X4, Y4, X1, Y1], axis=-1)
         verts = xy.reshape((npoly, 5, 2))
 
-        C = compress(ravelmask, ma.filled(C[0:Ny - 1, 0:Nx - 1]).ravel())
+        C = ma.filled(C[:Ny - 1, :Nx - 1])[unmask]
 
         linewidths = (0.25,)
         if 'linewidth' in kwargs:
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1311,23 +1311,23 @@ def _compute_conf_interval(data, med, iqr, bootstrap):
             hival = np.percentile(x, whis[1])
 
         # get high extreme
-        wiskhi = np.compress(x <= hival, x)
+        wiskhi = x[x <= hival]
         if len(wiskhi) == 0 or np.max(wiskhi) < q3:
             stats['whishi'] = q3
         else:
             stats['whishi'] = np.max(wiskhi)
 
         # get low extreme
-        wisklo = np.compress(x >= loval, x)
+        wisklo = x[x >= loval]
         if len(wisklo) == 0 or np.min(wisklo) > q1:
             stats['whislo'] = q1
         else:
             stats['whislo'] = np.min(wisklo)
 
         # compute a single array of outliers
         stats['fliers'] = np.hstack([
-            np.compress(x < stats['whislo'], x),
-            np.compress(x > stats['whishi'], x)
+            x[x < stats['whislo']],
+            x[x > stats['whishi']],
         ])
 
         # add in the remaining stats
diff --git a/lib/matplotlib/ticker.py b/lib/matplotlib/ticker.py
@@ -2641,9 +2641,9 @@ def __call__(self):
         mod = np.abs((locs - t0) % majorstep)
         cond1 = mod > minorstep / 10.0
         cond2 = ~np.isclose(mod, majorstep, atol=0)
-        locs = locs.compress(cond1 & cond2)
+        locs = locs[cond1 & cond2]
 
-        return self.raise_if_exceeds(np.array(locs))
+        return self.raise_if_exceeds(locs)
 
     def tick_values(self, vmin, vmax):
         raise NotImplementedError('Cannot get tick locations for a '
diff --git a/lib/matplotlib/tri/triangulation.py b/lib/matplotlib/tri/triangulation.py
@@ -117,7 +117,7 @@ def get_masked_triangles(self):
         Return an array of triangles that are not masked.
         """
         if self.mask is not None:
-            return self.triangles.compress(1 - self.mask, axis=0)
+            return self.triangles[~self.mask]
         else:
             return self.triangles
 
diff --git a/lib/matplotlib/tri/tripcolor.py b/lib/matplotlib/tri/tripcolor.py
@@ -112,7 +112,7 @@ def tripcolor(ax, *args, alpha=1.0, norm=None, cmap=None, vmin=None,
             C = C[maskedTris].mean(axis=1)
         elif tri.mask is not None:
             # Remove color values of masked triangles.
-            C = C.compress(1-tri.mask)
+            C = C[~tri.mask]
 
         collection = PolyCollection(verts, **kwargs)
 
diff --git a/lib/matplotlib/tri/tritools.py b/lib/matplotlib/tri/tritools.py
@@ -293,7 +293,7 @@ def _total_to_compress_renum(mask, n=None):
             n = np.size(mask)
         if mask is not None:
             renum = np.full(n, -1, dtype=np.int32)  # Default num is -1
-            valid = np.arange(n, dtype=np.int32).compress(~mask, axis=0)
+            valid = np.arange(n, dtype=np.int32)[~mask]
             renum[valid] = np.arange(np.size(valid, 0), dtype=np.int32)
             return renum
         else:
