Use np.stack / np.column_stack where possible.

anntzer · anntzer · commit 8c7fab7cf135 · 2018-03-03T11:58:43.000-08:00
diff --git a/examples/api/custom_projection_example.py b/examples/api/custom_projection_example.py
@@ -395,18 +395,17 @@ def __init__(self, resolution):
             self._resolution = resolution
 
         def transform_non_affine(self, ll):
-            longitude = ll[:, 0:1]
-            latitude = ll[:, 1:2]
+            longitude, latitude = ll.T
 
             # Pre-compute some values
-            half_long = longitude / 2.0
+            half_long = longitude / 2
             cos_latitude = np.cos(latitude)
-            sqrt2 = np.sqrt(2.0)
+            sqrt2 = np.sqrt(2)
 
-            alpha = np.sqrt(1.0 + cos_latitude * np.cos(half_long))
-            x = (2.0 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha
+            alpha = np.sqrt(1 + cos_latitude * np.cos(half_long))
+            x = (2 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha
             y = (sqrt2 * np.sin(latitude)) / alpha
-            return np.concatenate((x, y), 1)
+            return np.column_stack([x, y])
         transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__
 
         def transform_path_non_affine(self, path):
diff --git a/examples/misc/demo_ribbon_box.py b/examples/misc/demo_ribbon_box.py
@@ -30,7 +30,7 @@ def __init__(self, color):
                       self.original_image.dtype)
 
         im[:, :, :3] = self.b_and_h[:, :, np.newaxis]
-        im[:, :, :3] -= self.color[:, :, np.newaxis]*(1. - np.array(rgb))
+        im[:, :, :3] -= self.color[:, :, np.newaxis] * (1 - np.array(rgb))
         im[:, :, 3] = self.alpha
 
         self.im = im
diff --git a/examples/mplot3d/trisurf3d.py b/examples/mplot3d/trisurf3d.py
@@ -16,10 +16,7 @@
 
 # Make radii and angles spaces (radius r=0 omitted to eliminate duplication).
 radii = np.linspace(0.125, 1.0, n_radii)
-angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
-
-# Repeat all angles for each radius.
-angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
+angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)[..., np.newaxis]
 
 # Convert polar (radii, angles) coords to cartesian (x, y) coords.
 # (0, 0) is manually added at this stage,  so there will be no duplicate
diff --git a/examples/pyplots/boxplot_demo_pyplot.py b/examples/pyplots/boxplot_demo_pyplot.py
@@ -17,7 +17,7 @@
 center = np.ones(25) * 50
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
-data = np.concatenate((spread, center, flier_high, flier_low), 0)
+data = np.concatenate((spread, center, flier_high, flier_low))
 
 ###############################################################################
 
@@ -64,7 +64,7 @@
 center = np.ones(25) * 40
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
-d2 = np.concatenate((spread, center, flier_high, flier_low), 0)
+d2 = np.concatenate((spread, center, flier_high, flier_low))
 data.shape = (-1, 1)
 d2.shape = (-1, 1)
 
diff --git a/examples/shapes_and_collections/ellipse_collection.py b/examples/shapes_and_collections/ellipse_collection.py
@@ -12,7 +12,7 @@
 y = np.arange(15)
 X, Y = np.meshgrid(x, y)
 
-XY = np.hstack((X.ravel()[:, np.newaxis], Y.ravel()[:, np.newaxis]))
+XY = np.column_stack((X.ravel(), Y.ravel()))
 
 ww = X / 10.0
 hh = Y / 15.0
diff --git a/examples/shapes_and_collections/line_collection.py b/examples/shapes_and_collections/line_collection.py
@@ -22,7 +22,7 @@
 # Here are many sets of y to plot vs x
 ys = x[:50, np.newaxis] + x[np.newaxis, :]
 
-segs = np.zeros((50, 100, 2), float)
+segs = np.zeros((50, 100, 2))
 segs[:, :, 1] = ys
 segs[:, :, 0] = x
 
diff --git a/examples/specialty_plots/mri_with_eeg.py b/examples/specialty_plots/mri_with_eeg.py
@@ -40,11 +40,10 @@
 ax1.set_ylabel('MRI density')
 
 # Load the EEG data
-numSamples, numRows = 800, 4
+n_samples, n_rows = 800, 4
 with cbook.get_sample_data('eeg.dat') as eegfile:
-    data = np.fromfile(eegfile, dtype=float)
-data.shape = (numSamples, numRows)
-t = 10.0 * np.arange(numSamples) / numSamples
+    data = np.fromfile(eegfile, dtype=float).reshape((n_samples, n_rows))
+t = 10 * np.arange(n_samples) / n_samples
 
 # Plot the EEG
 ticklocs = []
@@ -55,15 +54,15 @@
 dmax = data.max()
 dr = (dmax - dmin) * 0.7  # Crowd them a bit.
 y0 = dmin
-y1 = (numRows - 1) * dr + dmax
+y1 = (n_rows - 1) * dr + dmax
 ax2.set_ylim(y0, y1)
 
 segs = []
-for i in range(numRows):
-    segs.append(np.hstack((t[:, np.newaxis], data[:, i, np.newaxis])))
+for i in range(n_rows):
+    segs.append(np.column_stack((t, data[:, i])))
     ticklocs.append(i * dr)
 
-offsets = np.zeros((numRows, 2), dtype=float)
+offsets = np.zeros((n_rows, 2), dtype=float)
 offsets[:, 1] = ticklocs
 
 lines = LineCollection(segs, offsets=offsets, transOffset=None)
diff --git a/examples/statistics/boxplot_demo.py b/examples/statistics/boxplot_demo.py
@@ -23,7 +23,7 @@
 center = np.ones(25) * 50
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
-data = np.concatenate((spread, center, flier_high, flier_low), 0)
+data = np.concatenate((spread, center, flier_high, flier_low))
 
 fig, axs = plt.subplots(2, 3)
 
@@ -59,7 +59,7 @@
 center = np.ones(25) * 40
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
-d2 = np.concatenate((spread, center, flier_high, flier_low), 0)
+d2 = np.concatenate((spread, center, flier_high, flier_low))
 data.shape = (-1, 1)
 d2.shape = (-1, 1)
 # Making a 2-D array only works if all the columns are the
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5732,26 +5732,20 @@ def pcolor(self, *args, **kwargs):
         # don't plot if C or any of the surrounding vertices are masked.
         mask = ma.getmaskarray(C) + xymask
 
-        newaxis = np.newaxis
         compress = np.compress
 
         ravelmask = (mask == 0).ravel()
-        X1 = compress(ravelmask, ma.filled(X[0:-1, 0:-1]).ravel())
-        Y1 = compress(ravelmask, ma.filled(Y[0:-1, 0:-1]).ravel())
-        X2 = compress(ravelmask, ma.filled(X[1:, 0:-1]).ravel())
-        Y2 = compress(ravelmask, ma.filled(Y[1:, 0:-1]).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[0:-1, 1:]).ravel())
-        Y4 = compress(ravelmask, ma.filled(Y[0:-1, 1:]).ravel())
+        X4 = compress(ravelmask, ma.filled(X[:-1, 1:]).ravel())
+        Y4 = compress(ravelmask, ma.filled(Y[:-1, 1:]).ravel())
         npoly = len(X1)
 
-        xy = np.concatenate((X1[:, newaxis], Y1[:, newaxis],
-                             X2[:, newaxis], Y2[:, newaxis],
-                             X3[:, newaxis], Y3[:, newaxis],
-                             X4[:, newaxis], Y4[:, newaxis],
-                             X1[:, newaxis], Y1[:, newaxis]),
-                            axis=1)
+        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())
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -341,8 +341,7 @@ def _makefill(self, x, y, kw, kwargs):
         # modify the kwargs dictionary.
         self._setdefaults(default_dict, kwargs)
 
-        seg = mpatches.Polygon(np.hstack((x[:, np.newaxis],
-                                          y[:, np.newaxis])),
+        seg = mpatches.Polygon(np.column_stack((x, y)),
                                facecolor=facecolor,
                                fill=kwargs.get('fill', True),
                                closed=kw['closed'])
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1781,11 +1781,9 @@ def convert_mesh_to_paths(tri):
 
         This function is primarily of use to backend implementers.
         """
-        Path = mpath.Path
         triangles = tri.get_masked_triangles()
-        verts = np.concatenate((tri.x[triangles][..., np.newaxis],
-                                tri.y[triangles][..., np.newaxis]), axis=2)
-        return [Path(x) for x in verts]
+        verts = np.stack((tri.x[triangles], tri.y[triangles]), axis=-1)
+        return [mpath.Path(x) for x in verts]
 
     @artist.allow_rasterization
     def draw(self, renderer):
@@ -1798,8 +1796,7 @@ def draw(self, renderer):
         tri = self._triangulation
         triangles = tri.get_masked_triangles()
 
-        verts = np.concatenate((tri.x[triangles][..., np.newaxis],
-                                tri.y[triangles][..., np.newaxis]), axis=2)
+        verts = np.stack((tri.x[triangles], tri.y[triangles]), axis=-1)
 
         self.update_scalarmappable()
         colors = self._facecolors[triangles]
@@ -1880,22 +1877,19 @@ def convert_mesh_to_paths(meshWidth, meshHeight, coordinates):
 
         This function is primarily of use to backend implementers.
         """
-        Path = mpath.Path
-
         if isinstance(coordinates, np.ma.MaskedArray):
             c = coordinates.data
         else:
             c = coordinates
-
         points = np.concatenate((
-                    c[0:-1, 0:-1],
-                    c[0:-1, 1:],
+                    c[:-1, :-1],
+                    c[:-1, 1:],
                     c[1:, 1:],
-                    c[1:, 0:-1],
-                    c[0:-1, 0:-1]
+                    c[1:, :-1],
+                    c[:-1, :-1]
                 ), axis=2)
         points = points.reshape((meshWidth * meshHeight, 5, 2))
-        return [Path(x) for x in points]
+        return [mpath.Path(x) for x in points]
 
     def convert_mesh_to_triangles(self, meshWidth, meshHeight, coordinates):
         """
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -162,49 +162,38 @@ def _slice_or_none(in_v, slc):
                         _slice_or_none(codes, slice(start, None, step)))
 
         elif isinstance(step, float):
-            if not (isinstance(start, int) or
-                    isinstance(start, float)):
-                raise ValueError('`markevery` is a tuple with '
-                    'len 2 and second element is a float, but '
-                    'the first element is not a float or an '
-                    'int; '
+            if not isinstance(start, (int, float)):
+                raise ValueError(
+                    '`markevery` is a tuple with len 2 and second element is '
+                    'a float, but the first element is not a float or an int; '
                     'markevery=%s' % (markevery,))
-            #calc cumulative distance along path (in display
-            # coords):
+            # calc cumulative distance along path (in display coords):
             disp_coords = affine.transform(tpath.vertices)
-            delta = np.empty((len(disp_coords), 2),
-                             dtype=float)
-            delta[0, :] = 0.0
-            delta[1:, :] = (disp_coords[1:, :] -
-                                disp_coords[:-1, :])
+            delta = np.empty((len(disp_coords), 2))
+            delta[0, :] = 0
+            delta[1:, :] = disp_coords[1:, :] - disp_coords[:-1, :]
             delta = np.sum(delta**2, axis=1)
             delta = np.sqrt(delta)
             delta = np.cumsum(delta)
-            #calc distance between markers along path based on
-            # the axes bounding box diagonal being a distance
-            # of unity:
-            scale = ax_transform.transform(
-                np.array([[0, 0], [1, 1]]))
+            # calc distance between markers along path based on the axes
+            # bounding box diagonal being a distance of unity:
+            scale = ax_transform.transform(np.array([[0, 0], [1, 1]]))
             scale = np.diff(scale, axis=0)
             scale = np.sum(scale**2)
             scale = np.sqrt(scale)
-            marker_delta = np.arange(start * scale,
-                                     delta[-1],
-                                     step * scale)
-            #find closest actual data point that is closest to
+            marker_delta = np.arange(start * scale, delta[-1], step * scale)
+            # find closest actual data point that is closest to
             # the theoretical distance along the path:
-            inds = np.abs(delta[np.newaxis, :] -
-                            marker_delta[:, np.newaxis])
+            inds = np.abs(delta[np.newaxis, :] - marker_delta[:, np.newaxis])
             inds = inds.argmin(axis=1)
             inds = np.unique(inds)
             # return, we are done here
             return Path(verts[inds],
                         _slice_or_none(codes, inds))
         else:
-            raise ValueError('`markevery` is a tuple with '
-                'len 2, but its second element is not an int '
-                'or a float; '
-                'markevery=%s' % (markevery,))
+            raise ValueError(
+                '`markevery` is a tuple with len 2, but its second element is '
+                'not an int or a float; markevery=%s' % (markevery,))
 
     elif isinstance(markevery, slice):
         # mazol tov, it's already a slice, just return
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -644,22 +644,20 @@ def unit_rectangle(cls):
     @classmethod
     def unit_regular_polygon(cls, numVertices):
         """
-        Return a :class:`Path` instance for a unit regular
-        polygon with the given *numVertices* and radius of 1.0,
-        centered at (0, 0).
+        Return a :class:`Path` instance for a unit regular polygon with the
+        given *numVertices* and radius of 1.0, centered at (0, 0).
         """
         if numVertices <= 16:
             path = cls._unit_regular_polygons.get(numVertices)
         else:
             path = None
         if path is None:
-            theta = (2*np.pi/numVertices *
-                     np.arange(numVertices + 1).reshape((numVertices + 1, 1)))
-            # This initial rotation is to make sure the polygon always
-            # "points-up"
-            theta += np.pi / 2.0
-            verts = np.concatenate((np.cos(theta), np.sin(theta)), 1)
-            codes = np.empty((numVertices + 1,))
+            theta = ((2 * np.pi / numVertices) * np.arange(numVertices + 1)
+                     # This initial rotation is to make sure the polygon always
+                     # "points-up".
+                     + np.pi / 2)
+            verts = np.column_stack((np.cos(theta), np.sin(theta)))
+            codes = np.empty(numVertices + 1)
             codes[0] = cls.MOVETO
             codes[1:-1] = cls.LINETO
             codes[-1] = cls.CLOSEPOLY
@@ -673,9 +671,8 @@ def unit_regular_polygon(cls, numVertices):
     @classmethod
     def unit_regular_star(cls, numVertices, innerCircle=0.5):
         """
-        Return a :class:`Path` for a unit regular star
-        with the given numVertices and radius of 1.0, centered at (0,
-        0).
+        Return a :class:`Path` for a unit regular star with the given
+        numVertices and radius of 1.0, centered at (0, 0).
         """
         if numVertices <= 16:
             path = cls._unit_regular_stars.get((numVertices, innerCircle))
@@ -702,9 +699,8 @@ def unit_regular_star(cls, numVertices, innerCircle=0.5):
     @classmethod
     def unit_regular_asterisk(cls, numVertices):
         """
-        Return a :class:`Path` for a unit regular
-        asterisk with the given numVertices and radius of 1.0,
-        centered at (0, 0).
+        Return a :class:`Path` for a unit regular asterisk with the given
+        numVertices and radius of 1.0, centered at (0, 0).
         """
         return cls.unit_regular_star(numVertices, 0.0)
 
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
diff --git a/lib/matplotlib/tri/triinterpolate.py b/lib/matplotlib/tri/triinterpolate.py
diff --git a/lib/matplotlib/tri/tripcolor.py b/lib/matplotlib/tri/tripcolor.py
diff --git a/lib/mpl_toolkits/axisartist/grid_finder.py b/lib/mpl_toolkits/axisartist/grid_finder.py
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
diff --git a/lib/mpl_toolkits/tests/test_axisartist_grid_helper_curvelinear.py b/lib/mpl_toolkits/tests/test_axisartist_grid_helper_curvelinear.py
