Remove/rework uses of np.where where possible.

anntzer · anntzer · commit 128ee1764695 · 2019-02-17T13:26:58.000+01:00
In a few places, np.where can be replaced by direct use of boolean
arrays, or other constructs.

In colors.py, the use of putmask may be slightly obscure (so I left a
comment explaining it), but avoids repeating the modified array twice,
so I think is more legible.

In polar.py, preallocating xy is just pointless given that we're
constructing intermediate arrays for x and y and filling them into xy
afterwards.
diff --git a/examples/images_contours_and_fields/barcode_demo.py b/examples/images_contours_and_fields/barcode_demo.py
@@ -11,12 +11,11 @@
 # Fixing random state for reproducibility
 np.random.seed(19680801)
 
-
 # the bar
-x = np.where(np.random.rand(500) > 0.7, 1.0, 0.0)
+x = np.random.rand(500) > 0.7
 
 axprops = dict(xticks=[], yticks=[])
-barprops = dict(aspect='auto', cmap=plt.cm.binary, interpolation='nearest')
+barprops = dict(aspect='auto', cmap='binary', interpolation='nearest')
 
 fig = plt.figure()
 
@@ -28,7 +27,6 @@
 ax2 = fig.add_axes([0.3, 0.1, 0.6, 0.1], **axprops)
 ax2.imshow(x.reshape((1, -1)), **barprops)
 
-
 plt.show()
 
 #############################################################################
diff --git a/examples/images_contours_and_fields/specgram_demo.py b/examples/images_contours_and_fields/specgram_demo.py
@@ -3,8 +3,7 @@
 Spectrogram Demo
 ================
 
-Demo of a spectrogram plot
-(:meth:`~.axes.Axes.specgram`).
+Demo of a spectrogram plot (`~.axes.Axes.specgram`).
 """
 import matplotlib.pyplot as plt
 import numpy as np
@@ -18,8 +17,7 @@
 s2 = 2 * np.sin(2 * np.pi * 400 * t)
 
 # create a transient "chirp"
-mask = np.where(np.logical_and(t > 10, t < 12), 1.0, 0.0)
-s2 = s2 * mask
+s2[t <= 10] = s2[12 <= t] = 0
 
 # add some noise into the mix
 nse = 0.01 * np.random.random(size=len(t))
diff --git a/examples/mplot3d/trisurf3d_2.py b/examples/mplot3d/trisurf3d_2.py
@@ -1,4 +1,4 @@
-'''
+"""
 ===========================
 More triangular 3D surfaces
 ===========================
@@ -8,7 +8,7 @@
 The first demonstrates use of plot_trisurf's triangles argument, and the
 second sets a Triangulation object's mask and passes the object directly
 to plot_trisurf.
-'''
+"""
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -71,7 +71,7 @@
 # Mask off unwanted triangles.
 xmid = x[triang.triangles].mean(axis=1)
 ymid = y[triang.triangles].mean(axis=1)
-mask = np.where(xmid**2 + ymid**2 < min_radius**2, 1, 0)
+mask = xmid**2 + ymid**2 < min_radius**2
 triang.set_mask(mask)
 
 # Plot the surface.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1775,28 +1775,19 @@ def blend_hsv(self, rgb, intensity, hsv_max_sat=None, hsv_max_val=None,
         hsv = rgb_to_hsv(rgb[:, :, 0:3])
 
         # modify hsv values to simulate illumination.
-        hsv[:, :, 1] = np.where(np.logical_and(np.abs(hsv[:, :, 1]) > 1.e-10,
-                                               intensity > 0),
-                                ((1. - intensity) * hsv[:, :, 1] +
-                                 intensity * hsv_max_sat),
-                                hsv[:, :, 1])
-
-        hsv[:, :, 2] = np.where(intensity > 0,
-                                ((1. - intensity) * hsv[:, :, 2] +
-                                 intensity * hsv_max_val),
-                                hsv[:, :, 2])
-
-        hsv[:, :, 1] = np.where(np.logical_and(np.abs(hsv[:, :, 1]) > 1.e-10,
-                                               intensity < 0),
-                                ((1. + intensity) * hsv[:, :, 1] -
-                                 intensity * hsv_min_sat),
-                                hsv[:, :, 1])
-        hsv[:, :, 2] = np.where(intensity < 0,
-                                ((1. + intensity) * hsv[:, :, 2] -
-                                 intensity * hsv_min_val),
-                                hsv[:, :, 2])
-        hsv[:, :, 1:] = np.where(hsv[:, :, 1:] < 0., 0, hsv[:, :, 1:])
-        hsv[:, :, 1:] = np.where(hsv[:, :, 1:] > 1., 1, hsv[:, :, 1:])
+        np.putmask(hsv[:, :, 1],  # i.e. A[mask] = B[mask].
+                   (np.abs(hsv[:, :, 1]) > 1.e-10) & (intensity > 0),
+                   (1 - intensity) * hsv[:, :, 1] + intensity * hsv_max_sat)
+        np.putmask(hsv[:, :, 2],
+                   intensity > 0,
+                   (1 - intensity) * hsv[:, :, 2] + intensity * hsv_max_val)
+        np.putmask(hsv[:, :, 1],
+                   (np.abs(hsv[:, :, 1]) > 1.e-10) & (intensity < 0),
+                   (1 + intensity) * hsv[:, :, 1] - intensity * hsv_min_sat)
+        np.putmask(hsv[:, :, 2],
+                   intensity < 0,
+                   (1 + intensity) * hsv[:, :, 2] - intensity * hsv_min_val)
+        np.clip(hsv[:, :, 1:], 0, 1, out=hsv[:, :, 1:])
         # convert modified hsv back to rgb.
         return hsv_to_rgb(hsv)
 
diff --git a/lib/matplotlib/projections/polar.py b/lib/matplotlib/projections/polar.py
@@ -44,27 +44,16 @@ def __str__(self):
 
     def transform_non_affine(self, tr):
         # docstring inherited
-        xy = np.empty(tr.shape, float)
-
-        t = tr[:, 0:1]
-        r = tr[:, 1:2]
-        x = xy[:, 0:1]
-        y = xy[:, 1:2]
-
+        t, r = np.transpose(tr)
         # PolarAxes does not use the theta transforms here, but apply them for
         # backwards-compatibility if not being used by it.
         if self._apply_theta_transforms and self._axis is not None:
             t *= self._axis.get_theta_direction()
             t += self._axis.get_theta_offset()
-
         if self._use_rmin and self._axis is not None:
             r = (r - self._axis.get_rorigin()) * self._axis.get_rsign()
-
-        mask = r < 0
-        x[:] = np.where(mask, np.nan, r * np.cos(t))
-        y[:] = np.where(mask, np.nan, r * np.sin(t))
-
-        return xy
+        r = np.where(r >= 0, r, np.nan)
+        return np.column_stack([r * np.cos(t), r * np.sin(t)])
 
     def transform_path_non_affine(self, path):
         # docstring inherited
diff --git a/lib/matplotlib/tri/triinterpolate.py b/lib/matplotlib/tri/triinterpolate.py
@@ -1341,7 +1341,7 @@ def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000):
     # Jacobi pre-conditioner
     kvec = A.diag
     # For diag elem < 1e-6 we keep 1e-6.
-    kvec = np.where(kvec > 1.e-6, kvec, 1.e-6)
+    kvec = np.maximum(kvec, 1e-6)
 
     # Initial guess
     if x0 is None:
