More broadcasting in mplot3d.

anntzer · anntzer · commit b4296f4c8dea · 2017-07-29T16:41:34.000-07:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -149,13 +149,9 @@ def line_2d_to_3d(line, zs=0, zdir='z'):
 def path_to_3d_segment(path, zs=0, zdir='z'):
     '''Convert a path to a 3D segment.'''
 
-    if not iterable(zs):
-        zs = np.ones(len(path)) * zs
-
-    seg = []
+    zs = _backports_np.broadcast_to(zs, len(path))
     pathsegs = path.iter_segments(simplify=False, curves=False)
-    for (((x, y), code), z) in zip(pathsegs, zs):
-        seg.append((x, y, z))
+    seg = [(x, y, z) for (((x, y), code), z) in zip(pathsegs, zs)]
     seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg]
     return seg3d
 
@@ -165,21 +161,16 @@ def paths_to_3d_segments(paths, zs=0, zdir='z'):
     Convert paths from a collection object to 3D segments.
     '''
 
-    if not iterable(zs):
-        zs = np.ones(len(paths)) * zs
-
-    segments = []
-    for path, pathz in zip(paths, zs):
-        segments.append(path_to_3d_segment(path, pathz, zdir))
-    return segments
+    zs = _backports_np.broadcast_to(zs, len(paths))
+    segs = [path_to_3d_segment(path, pathz, zdir)
+            for path, pathz in zip(paths, zs)]
+    return segs
 
 
 def path_to_3d_segment_with_codes(path, zs=0, zdir='z'):
     '''Convert a path to a 3D segment with path codes.'''
 
-    if not iterable(zs):
-        zs = np.ones(len(path)) * zs
-
+    zs = _backports_np.broadcast_to(zs, len(path))
     seg = []
     codes = []
     pathsegs = path.iter_segments(simplify=False, curves=False)
@@ -195,9 +186,7 @@ def paths_to_3d_segments_with_codes(paths, zs=0, zdir='z'):
     Convert paths from a collection object to 3D segments with path codes.
     '''
 
-    if not iterable(zs):
-        zs = np.ones(len(paths)) * zs
-
+    zs = _backports_np.broadcast_to(zs, len(paths))
     segments = []
     codes_list = []
     for path, pathz in zip(paths, zs):
@@ -271,11 +260,9 @@ def __init__(self, *args, **kwargs):
         self.set_3d_properties(zs, zdir)
 
     def set_3d_properties(self, verts, zs=0, zdir='z'):
-        if not iterable(zs):
-            zs = np.ones(len(verts)) * zs
-
-        self._segment3d = [juggle_axes(x, y, z, zdir) \
-                for ((x, y), z) in zip(verts, zs)]
+        zs = _backports_np.broadcast_to(zs, len(verts))
+        self._segment3d = [juggle_axes(x, y, z, zdir)
+                           for ((x, y), z) in zip(verts, zs)]
         self._facecolor3d = Patch.get_facecolor(self)
 
     def get_path(self):
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -11,25 +11,23 @@
 """
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
-import math
 
 import six
 from six.moves import map, xrange, zip, reduce
 
+import math
 import warnings
 
 import numpy as np
-import matplotlib.axes as maxes
+
+from matplotlib import (
+    axes as maxes, cbook, collections as mcoll, colors as mcolors, docstring,
+    scale as mscale, transforms as mtransforms)
+from matplotlib._backports import numpy as _backports_np
 from matplotlib.axes import Axes, rcParams
-from matplotlib import cbook
-import matplotlib.transforms as mtransforms
+from matplotlib.colors import Normalize, LightSource
 from matplotlib.transforms import Bbox
-import matplotlib.collections as mcoll
-from matplotlib import docstring
-import matplotlib.scale as mscale
 from matplotlib.tri.triangulation import Triangulation
-from matplotlib import colors as mcolors
-from matplotlib.colors import Normalize, LightSource
 
 from . import art3d
 from . import proj3d
@@ -1536,8 +1534,7 @@ def plot(self, xs, ys, *args, **kwargs):
         zdir = kwargs.pop('zdir', 'z')
 
         # Match length
-        if not cbook.iterable(zs):
-            zs = np.ones(len(xs)) * zs
+        zs = _backports_np.broadcast_to(zs, len(xs))
 
         lines = super(Axes3D, self).plot(xs, ys, *args, **kwargs)
         for line in lines:
@@ -2332,18 +2329,16 @@ def scatter(self, xs, ys, zs=0, zdir='z', s=20, c=None, depthshade=True,
 
         had_data = self.has_data()
 
-        xs, ys, zs = np.broadcast_arrays(*map(np.ma.ravel, [xs, ys, zs]))
+        xs, ys, zs = np.broadcast_arrays(
+            *[np.ravel(np.ma.filled(t, np.nan)) for t in [xs, ys, zs]])
         s = np.ma.ravel(s)  # This doesn't have to match x, y in size.
 
         xs, ys, zs, s, c = cbook.delete_masked_points(xs, ys, zs, s, c)
 
-        patches = super(Axes3D, self).scatter(xs, ys, s=s, c=c, *args,
-                                              **kwargs)
-        if not cbook.iterable(zs):
-            is_2d = True
-            zs = np.ones(len(xs)) * zs
-        else:
-            is_2d = False
+        patches = super(Axes3D, self).scatter(
+            xs, ys, s=s, c=c, *args, **kwargs)
+        is_2d = not cbook.iterable(zs)
+        zs = _backports_np.broadcast_to(zs, len(xs))
         art3d.patch_collection_2d_to_3d(patches, zs=zs, zdir=zdir,
                                         depthshade=depthshade)
 
@@ -2382,8 +2377,7 @@ def bar(self, left, height, zs=0, zdir='z', *args, **kwargs):
 
         patches = super(Axes3D, self).bar(left, height, *args, **kwargs)
 
-        if not cbook.iterable(zs):
-            zs = np.ones(len(left)) * zs
+        zs = _backports_np.broadcast_to(zs, len(left))
 
         verts = []
         verts_zs = []
