Cleanup hexbin.

anntzer · anntzer · commit b5fd5d2c2beb · 2019-03-18T08:52:22.000+01:00
Use `&amp;` instead of `*` for boolean arrays; vectorize accumulation
instead of iterating over individual entries by hand; misc. style
cleanups.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4675,26 +4675,14 @@ def hexbin(self, x, y, C=None, gridsize=100, bins=None,
         if C is None:
             lattice1 = np.zeros((nx1, ny1))
             lattice2 = np.zeros((nx2, ny2))
-
-            cond1 = (0 <= ix1) * (ix1 < nx1) * (0 <= iy1) * (iy1 < ny1)
-            cond2 = (0 <= ix2) * (ix2 < nx2) * (0 <= iy2) * (iy2 < ny2)
-
-            cond1 *= bdist
-            cond2 *= np.logical_not(bdist)
-            ix1, iy1 = ix1[cond1], iy1[cond1]
-            ix2, iy2 = ix2[cond2], iy2[cond2]
-
-            for ix, iy in zip(ix1, iy1):
-                lattice1[ix, iy] += 1
-            for ix, iy in zip(ix2, iy2):
-                lattice2[ix, iy] += 1
-
-            # threshold
+            c1 = (0 <= ix1) & (ix1 < nx1) & (0 <= iy1) & (iy1 < ny1) & bdist
+            c2 = (0 <= ix2) & (ix2 < nx2) & (0 <= iy2) & (iy2 < ny2) & ~bdist
+            np.add.at(lattice1, (ix1[c1], iy1[c1]), 1)
+            np.add.at(lattice2, (ix2[c2], iy2[c2]), 1)
             if mincnt is not None:
                 lattice1[lattice1 < mincnt] = np.nan
                 lattice2[lattice2 < mincnt] = np.nan
-            accum = np.hstack((lattice1.ravel(),
-                               lattice2.ravel()))
+            accum = np.concatenate([lattice1.ravel(), lattice2.ravel()])
             good_idxs = ~np.isnan(accum)
 
         else:
@@ -4751,9 +4739,8 @@ def hexbin(self, x, y, C=None, gridsize=100, bins=None,
         offsets = offsets[good_idxs, :]
         accum = accum[good_idxs]
 
-        polygon = np.zeros((6, 2), float)
-        polygon[:, 0] = sx * np.array([0.5, 0.5, 0.0, -0.5, -0.5, 0.0])
-        polygon[:, 1] = sy * np.array([-0.5, 0.5, 1.0, 0.5, -0.5, -1.0]) / 3.0
+        polygon = [sx, sy / 3] * np.array(
+            [[.5, -.5], [.5, .5], [0., 1.], [-.5, .5], [-.5, -.5], [0., -1.]])
 
         if linewidths is None:
             linewidths = [1.0]
