Broadcasting for fun and profit.

anntzer · anntzer · commit ea968ed53999 · 2016-12-04T20:31:18.000-08:00
diff --git a/examples/pylab_examples/fancybox_demo2.py b/examples/pylab_examples/fancybox_demo2.py
@@ -4,7 +4,7 @@
 styles = mpatch.BoxStyle.get_styles()
 spacing = 1.2
 
-figheight = (spacing * len(styles) + .5)
+figheight = spacing * len(styles) + .5
 fig1 = plt.figure(1, (4/1.5, figheight/1.5))
 fontsize = 0.3 * 72
 
diff --git a/examples/pylab_examples/table_demo.py b/examples/pylab_examples/table_demo.py
@@ -25,7 +25,7 @@
 bar_width = 0.4
 
 # Initialize the vertical-offset for the stacked bar chart.
-y_offset = np.array([0.0] * len(columns))
+y_offset = np.zeros(len(columns))
 
 # Plot bars and create text labels for the table
 cell_text = []
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -14,6 +14,7 @@
 import matplotlib
 from matplotlib import _preprocess_data
 
+from matplotlib._backports import numpy as _backports_np
 import matplotlib.cbook as cbook
 from matplotlib.cbook import (mplDeprecation, STEP_LOOKUP_MAP,
                               iterable, is_string_like,
@@ -1997,20 +1998,14 @@ def bar(self, left, height, width=0.8, bottom=None, **kwargs):
         label = kwargs.pop('label', '')
         tick_labels = kwargs.pop('tick_label', None)
 
-        def make_iterable(x):
-            if not iterable(x):
-                return [x]
-            else:
-                return x
-
         # make them safe to take len() of
         _left = left
-        left = make_iterable(left)
-        height = make_iterable(height)
-        width = make_iterable(width)
+        left = np.atleast_1d(left)
+        height = np.atleast_1d(height)
+        width = np.atleast_1d(width)
         _bottom = bottom
-        bottom = make_iterable(bottom)
-        linewidth = make_iterable(linewidth)
+        bottom = np.atleast_1d(bottom)
+        linewidth = np.atleast_1d(linewidth)
 
         adjust_ylim = False
         adjust_xlim = False
@@ -2025,10 +2020,8 @@ def make_iterable(x):
                 bottom = [0]
 
             nbars = len(left)
-            if len(width) == 1:
-                width *= nbars
-            if len(bottom) == 1:
-                bottom *= nbars
+            width = _backports_np.broadcast_to(width, nbars)
+            bottom = _backports_np.broadcast_to(bottom, nbars)
 
             tick_label_axis = self.xaxis
             tick_label_position = left
@@ -2043,18 +2036,16 @@ def make_iterable(x):
                 left = [0]
 
             nbars = len(bottom)
-            if len(left) == 1:
-                left *= nbars
-            if len(height) == 1:
-                height *= nbars
+            left = _backports_np.broadcast_to(left, nbars)
+            height = _backports_np.broadcast_to(height, nbars)
 
             tick_label_axis = self.yaxis
             tick_label_position = bottom
         else:
             raise ValueError('invalid orientation: %s' % orientation)
 
         if len(linewidth) < nbars:
-            linewidth *= nbars
+            linewidth = np.tile(linewidth, nbars)
 
         color = list(mcolors.to_rgba_array(color))
         if len(color) == 0:  # until to_rgba_array is changed
@@ -2181,15 +2172,7 @@ def make_iterable(x):
         self.add_container(bar_container)
 
         if tick_labels is not None:
-            tick_labels = make_iterable(tick_labels)
-            if isinstance(tick_labels, six.string_types):
-                tick_labels = [tick_labels]
-            if len(tick_labels) == 1:
-                tick_labels *= nbars
-            if len(tick_labels) != nbars:
-                raise ValueError("incompatible sizes: argument 'tick_label' "
-                                 "must be length %d or string" % nbars)
-
+            tick_labels = _backports_np.broadcast_to(tick_labels, nbars)
             tick_label_axis.set_ticks(tick_label_position)
             tick_label_axis.set_ticklabels(tick_labels)
 
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -192,7 +192,8 @@ def draw_text(self, gc, x, y, s, prop, angle, ismath=False, mtext=None):
         flags = get_hinting_flag()
         font = self._get_agg_font(prop)
 
-        if font is None: return None
+        if font is None:
+            return None
         if len(s) == 1 and ord(s) > 127:
             font.load_char(ord(s), flags=flags)
         else:
diff --git a/lib/matplotlib/backends/backend_qt5.py b/lib/matplotlib/backends/backend_qt5.py
@@ -642,8 +642,8 @@ def edit_parameters(self):
             QtWidgets.QMessageBox.warning(
                 self.parent, "Error", "There are no axes to edit.")
             return
-        if len(allaxes) == 1:
-            axes = allaxes[0]
+        elif len(allaxes) == 1:
+            axes, = allaxes
         else:
             titles = []
             for axes in allaxes:
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -15,7 +15,6 @@
 import collections
 import datetime
 import errno
-import functools
 import glob
 import gzip
 import io
@@ -1799,7 +1798,7 @@ def delete_masked_points(*args):
             except:  # Fixme: put in tuple of possible exceptions?
                 pass
     if len(masks):
-        mask = functools.reduce(np.logical_and, masks)
+        mask = np.logical_and.reduce(masks)
         igood = mask.nonzero()[0]
         if len(igood) < nrecs:
             for i, x in enumerate(margs):
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -58,9 +58,12 @@
 
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
-import re
+
 import six
 from six.moves import zip
+
+import itertools
+import re
 import warnings
 
 import numpy as np
@@ -774,26 +777,23 @@ def __init__(self, colors, name='from_list', N=None):
 
             the list will be extended by repetition.
         """
-        self.colors = colors
         self.monochrome = False  # True only if all colors in map are
                                  # identical; needed for contouring.
         if N is None:
-            N = len(self.colors)
+            self.colors = colors
+            N = len(colors)
         else:
-            if (cbook.is_string_like(self.colors) and
-                    cbook.is_hashable(self.colors)):
-                self.colors = [self.colors] * N
+            if cbook.is_string_like(colors) and cbook.is_hashable(colors):
+                self.colors = [colors] * N
                 self.monochrome = True
-            elif cbook.iterable(self.colors):
-                self.colors = list(self.colors)  # in case it was a tuple
-                if len(self.colors) == 1:
+            elif cbook.iterable(colors):
+                if len(colors) == 1:
                     self.monochrome = True
-                if len(self.colors) < N:
-                    self.colors = list(self.colors) * N
-                del(self.colors[N:])
+                self.colors = list(
+                    itertools.islice(itertools.cycle(colors), N))
             else:
                 try:
-                    gray = float(self.colors)
+                    gray = float(colors)
                 except TypeError:
                     pass
                 else:
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -100,40 +100,32 @@ def get_grid_positions(self, fig):
         # calculate accumulated heights of columns
         cellH = totHeight/(nrows + hspace*(nrows-1))
         sepH = hspace*cellH
-
         if self._row_height_ratios is not None:
             netHeight = cellH * nrows
             tr = float(sum(self._row_height_ratios))
             cellHeights = [netHeight*r/tr for r in self._row_height_ratios]
         else:
             cellHeights = [cellH] * nrows
-
         sepHeights = [0] + ([sepH] * (nrows-1))
-        cellHs = np.add.accumulate(np.ravel(list(zip(sepHeights, cellHeights))))
-
+        cellHs = np.cumsum(np.column_stack([sepHeights, cellHeights]))
 
         # calculate accumulated widths of rows
         cellW = totWidth/(ncols + wspace*(ncols-1))
         sepW = wspace*cellW
-
         if self._col_width_ratios is not None:
             netWidth = cellW * ncols
             tr = float(sum(self._col_width_ratios))
             cellWidths = [netWidth*r/tr for r in self._col_width_ratios]
         else:
             cellWidths = [cellW] * ncols
-
         sepWidths = [0] + ([sepW] * (ncols-1))
-        cellWs = np.add.accumulate(np.ravel(list(zip(sepWidths, cellWidths))))
-
-
+        cellWs = np.cumsum(np.column_stack([sepWidths, cellWidths]))
 
         figTops = [top - cellHs[2*rowNum] for rowNum in range(nrows)]
         figBottoms = [top - cellHs[2*rowNum+1] for rowNum in range(nrows)]
         figLefts = [left + cellWs[2*colNum] for colNum in range(ncols)]
         figRights = [left + cellWs[2*colNum+1] for colNum in range(ncols)]
 
-
         return figBottoms, figTops, figLefts, figRights
 
     def __getitem__(self, key):
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -146,8 +146,7 @@ def flush_images():
             if len(image_group) == 1:
                 image_group[0].draw(renderer)
             elif len(image_group) > 1:
-                data, l, b = composite_images(
-                    image_group, renderer, mag)
+                data, l, b = composite_images(image_group, renderer, mag)
                 if data.size != 0:
                     gc = renderer.new_gc()
                     gc.set_clip_rectangle(parent.bbox)
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -294,7 +294,7 @@ def __init__(self, parent, handles, labels,
             self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.])
         else:
             self._scatteryoffsets = np.asarray(scatteryoffsets)
-        reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1
+        reps = self.scatterpoints // len(self._scatteryoffsets) + 1
         self._scatteryoffsets = np.tile(self._scatteryoffsets,
                                         reps)[:self.scatterpoints]
 
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -665,20 +665,8 @@ def recache(self, always=False):
         else:
             y = self._y
 
-        if len(x) == 1 and len(y) > 1:
-            x = x * np.ones(y.shape, float)
-        if len(y) == 1 and len(x) > 1:
-            y = y * np.ones(x.shape, float)
-
-        if len(x) != len(y):
-            raise RuntimeError('xdata and ydata must be the same length')
-
-        self._xy = np.empty((len(x), 2), dtype=float)
-        self._xy[:, 0] = x
-        self._xy[:, 1] = y
-
-        self._x = self._xy[:, 0]  # just a view
-        self._y = self._xy[:, 1]  # just a view
+        self._xy = np.column_stack(np.broadcast_arrays(x, y)).astype(float)
+        self._x, self._y = self._xy.T  # views
 
         self._subslice = False
         if (self.axes and len(x) > 1000 and self._is_sorted(x) and
diff --git a/lib/matplotlib/offsetbox.py b/lib/matplotlib/offsetbox.py
@@ -71,33 +71,28 @@ def _get_packed_offsets(wd_list, total, sep, mode="fixed"):
     # d_list is currently not used.
 
     if mode == "fixed":
-        offsets_ = np.add.accumulate([0] + [w + sep for w in w_list])
+        offsets_ = np.cumsum([0] + [w + sep for w in w_list])
         offsets = offsets_[:-1]
-
         if total is None:
             total = offsets_[-1] - sep
-
         return total, offsets
 
     elif mode == "expand":
         if len(w_list) > 1:
             sep = (total - sum(w_list)) / (len(w_list) - 1.)
         else:
-            sep = 0.
-        offsets_ = np.add.accumulate([0] + [w + sep for w in w_list])
+            sep = 0
+        offsets_ = np.cumsum([0] + [w + sep for w in w_list])
         offsets = offsets_[:-1]
-
         return total, offsets
 
     elif mode == "equal":
         maxh = max(w_list)
         if total is None:
             total = (maxh + sep) * len(w_list)
         else:
-            sep = float(total) / (len(w_list)) - maxh
-
-        offsets = np.array([(maxh + sep) * i for i in range(len(w_list))])
-
+            sep = total / len(w_list) - maxh
+        offsets = (maxh + sep) * np.arange(len(w_list))
         return total, offsets
 
     else:
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -1601,8 +1601,7 @@ def draw(self, renderer):
         # Get the width and height in pixels
         width = self.convert_xunits(self.width)
         height = self.convert_yunits(self.height)
-        width, height = self.get_transform().transform_point(
-            (width, height))
+        width, height = self.get_transform().transform_point((width, height))
         inv_error = (1.0 / 1.89818e-6) * 0.5
 
         if width < inv_error and height < inv_error:
diff --git a/lib/matplotlib/stackplot.py b/lib/matplotlib/stackplot.py
@@ -55,10 +55,7 @@ def stackplot(axes, x, *args, **kwargs):
     element in the stacked area plot.
     """
 
-    if len(args) == 1:
-        y = np.atleast_2d(*args)
-    elif len(args) > 1:
-        y = np.row_stack(args)
+    y = np.row_stack(args)
 
     labels = iter(kwargs.pop('labels', []))
 
diff --git a/lib/matplotlib/tests/test_backend_bases.py b/lib/matplotlib/tests/test_backend_bases.py
@@ -36,11 +36,10 @@ def check(master_transform, paths, all_transforms,
                                    [], 'data')]
         uses = rb._iter_collection_uses_per_path(
             paths, all_transforms, offsets, facecolors, edgecolors)
-        seen = [0] * len(raw_paths)
-        for i in ids:
-            seen[i] += 1
-        for n in seen:
-            assert n in (uses-1, uses)
+        if raw_paths:
+            seen = np.bincount(ids, minlength=len(raw_paths))
+            for n in seen:
+                assert n in [uses - 1, uses]
 
     check(id, paths, tforms, offsets, facecolors, edgecolors)
     check(id, paths[0:1], tforms, offsets, facecolors, edgecolors)
diff --git a/lib/matplotlib/tests/test_bbox_tight.py b/lib/matplotlib/tests/test_bbox_tight.py
@@ -30,7 +30,7 @@ def test_bbox_inches_tight():
     ind = np.arange(len(colLabels)) + 0.3  # the x locations for the groups
     cellText = []
     width = 0.4     # the width of the bars
-    yoff = np.array([0.0] * len(colLabels))
+    yoff = np.zeros(len(colLabels))
     # the bottom values for stacked bar chart
     fig, ax = plt.subplots(1, 1)
     for row in xrange(rows):
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -326,10 +326,8 @@ def test_image_edges():
 
     data = np.tile(np.arange(12), 15).reshape(20, 9)
 
-    im = ax.imshow(data, origin='upper',
-                   extent=[-10, 10, -10, 10], interpolation='none',
-                   cmap='gray'
-                   )
+    im = ax.imshow(data, origin='upper', extent=[-10, 10, -10, 10],
+                   interpolation='none', cmap='gray')
 
     x = y = 2
     ax.set_xlim([-x, x])
diff --git a/lib/matplotlib/tests/test_table.py b/lib/matplotlib/tests/test_table.py
@@ -30,7 +30,7 @@ def test_zorder():
     rowLabels = ['%d year' % x for x in (100, 50)]
 
     cellText = []
-    yoff = np.array([0.0] * len(colLabels))
+    yoff = np.zeros(len(colLabels))
     for row in reversed(data):
         yoff += row
         cellText.append(['%1.1f' % (x/1000.0) for x in yoff])
diff --git a/lib/matplotlib/textpath.py b/lib/matplotlib/textpath.py
@@ -225,7 +225,7 @@ def get_glyphs_with_font(self, font, s, glyph_map=None,
         rects = []
 
         return (list(zip(glyph_ids, xpositions, ypositions, sizes)),
-                     glyph_map_new, rects)
+                glyph_map_new, rects)
 
     def get_glyphs_mathtext(self, prop, s, glyph_map=None,
                             return_new_glyphs_only=False):
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
