Merge pull request #10614 from anntzer/stack

timhoffm · web-flow · commit 3f3065d3a07f · 2018-02-28T01:45:14.000+01:00
Use np.stack instead of list(zip()) in colorbar.py.
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -557,13 +557,11 @@ def add_lines(self, levels, colors, linewidths, erase=True):
             colors = np.asarray(colors)[igood]
         if cbook.iterable(linewidths):
             linewidths = np.asarray(linewidths)[igood]
-        N = len(y)
-        x = np.array([0.0, 1.0])
-        X, Y = np.meshgrid(x, y)
+        X, Y = np.meshgrid([0, 1], y)
         if self.orientation == 'vertical':
-            xy = [list(zip(X[i], Y[i])) for i in xrange(N)]
+            xy = np.stack([X, Y], axis=-1)
         else:
-            xy = [list(zip(Y[i], X[i])) for i in xrange(N)]
+            xy = np.stack([Y, X], axis=-1)
         col = collections.LineCollection(xy, linewidths=linewidths)
 
         if erase and self.lines:
diff --git a/lib/mpl_toolkits/axes_grid1/colorbar.py b/lib/mpl_toolkits/axes_grid1/colorbar.py
@@ -565,10 +565,11 @@ def _add_solids(self, X, Y, C):
 
         self.solids = col
         if self.drawedges:
-            self.dividers = collections.LineCollection(self._edges(X,Y),
-                              colors=(mpl.rcParams['axes.edgecolor'],),
-                              linewidths=(0.5*mpl.rcParams['axes.linewidth'],),
-                              )
+            self.dividers = collections.LineCollection(
+                self._edges(X,Y),
+                colors=(mpl.rcParams['axes.edgecolor'],),
+                linewidths=(0.5*mpl.rcParams['axes.linewidth'],),
+            )
             self.ax.add_collection(self.dividers)
         else:
             self.dividers = None
@@ -577,22 +578,16 @@ def add_lines(self, levels, colors, linewidths):
         '''
         Draw lines on the colorbar. It deletes preexisting lines.
         '''
-        del self.lines
-
-        N = len(levels)
-        x = np.array([1.0, 2.0])
-        X, Y = np.meshgrid(x,levels)
+        X, Y = np.meshgrid([1, 2], levels)
         if self.orientation == 'vertical':
-            xy = [list(zip(X[i], Y[i])) for i in range(N)]
+            xy = np.stack([X, Y], axis=-1)
         else:
-            xy = [list(zip(Y[i], X[i])) for i in range(N)]
-        col = collections.LineCollection(xy, linewidths=linewidths,
-                                         )
+            xy = np.stack([Y, X], axis=-1)
+        col = collections.LineCollection(xy, linewidths=linewidths)
         self.lines = col
         col.set_color(colors)
         self.ax.add_collection(col)
 
-
     def _select_locator(self, formatter):
         '''
         select a suitable locator
