Add support for x,y,z arguments

eric-wieser · eric-wieser · commit 50d9070674ce · 2017-07-18T23:09:43.000+01:00
diff --git a/examples/mplot3d/voxels_numpy_logo.py b/examples/mplot3d/voxels_numpy_logo.py
@@ -0,0 +1,47 @@
+'''
+===============================
+3D voxel plot of the numpy logo
+===============================
+
+Demonstrates using ``ax.voxels`` with uneven coordinates
+'''
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.mplot3d import Axes3D
+
+
+def explode(data):
+    size = np.array(data.shape)*2
+    data_e = np.zeros(size - 1, dtype=data.dtype)
+    data_e[::2, ::2, ::2] = data
+    return data_e
+
+# build up the numpy logo
+n_voxels = np.zeros((4, 3, 4), dtype=bool)
+n_voxels[0, 0, :] = True
+n_voxels[-1, 0, :] = True
+n_voxels[1, 0, 2] = True
+n_voxels[2, 0, 1] = True
+facecolors = np.where(n_voxels, '#FFD65DC0', '#7A88CCC0')
+edgecolors = np.where(n_voxels, '#BFAB6E', '#7D84A6')
+filled = np.ones(n_voxels.shape)
+
+# upscale the above voxel image, leaving gaps
+filled_2 = explode(filled)
+fcolors_2 = explode(facecolors)
+ecolors_2 = explode(edgecolors)
+
+# Shrink the gaps
+x, y, z = np.indices(np.array(filled_2.shape) + 1).astype(float) // 2
+x[0::2, :, :] += 0.05
+y[:, 0::2, :] += 0.05
+z[:, :, 0::2] += 0.05
+x[1::2, :, :] += 0.95
+y[:, 1::2, :] += 0.95
+z[:, :, 1::2] += 0.95
+
+fig = plt.figure()
+ax = fig.gca(projection='3d')
+ax.voxels(x, y, z, filled_2, facecolors=fcolors_2, edgecolors=ecolors_2)
+
+plt.show()
diff --git a/examples/mplot3d/voxels_rgb.py b/examples/mplot3d/voxels_rgb.py
@@ -10,27 +10,36 @@
 import numpy as np
 from mpl_toolkits.mplot3d import Axes3D
 
+
+def midpoints(x):
+    sl = ()
+    for i in range(x.ndim):
+        x = (x[sl + np.index_exp[:-1]] + x[sl + np.index_exp[1:]]) / 2.0
+        sl += np.index_exp[:]
+    return x
+
 # prepare some coordinates, and attach rgb values to each
-x, y, z = np.indices((16, 16, 16))
-r = (x + 0.5) / 16
-g = (y + 0.5) / 16
-b = (z + 0.5) / 16
+r, g, b = np.indices((17, 17, 17)) / 16.0
+rc = midpoints(r)
+gc = midpoints(g)
+bc = midpoints(b)
 
 # define a sphere about [0.5, 0.5, 0.5]
-sphere = (r - 0.5)**2 + (g - 0.5)**2 + (b - 0.5)**2 < 0.5**2
+sphere = (rc - 0.5)**2 + (gc - 0.5)**2 + (bc - 0.5)**2 < 0.5**2
 
 # combine the color components
 colors = np.zeros(sphere.shape + (3,))
-colors[..., 0] = r
-colors[..., 1] = g
-colors[..., 2] = b
+colors[..., 0] = rc
+colors[..., 1] = gc
+colors[..., 2] = bc
 
 # and plot everything
 fig = plt.figure()
 ax = fig.gca(projection='3d')
-ax.voxels(sphere,
+ax.voxels(r, g, b, sphere,
           facecolors=colors,
           edgecolors=np.clip(2*colors - 0.5, 0, 1),  # brighter
           linewidth=0.5)
+ax.set(xlabel='r', ylabel='g', zlabel='b')
 
 plt.show()
diff --git a/examples/mplot3d/voxels_torus.py b/examples/mplot3d/voxels_torus.py
@@ -0,0 +1,47 @@
+'''
+=======================================================
+3D voxel / volumetric plot with cylindrical coordinates
+=======================================================
+
+Demonstrates using the ``x, y, z`` arguments of ``ax.voxels``.
+'''
+
+import matplotlib.pyplot as plt
+import matplotlib.colors
+import numpy as np
+from mpl_toolkits.mplot3d import Axes3D
+
+
+def midpoints(x):
+    sl = ()
+    for i in range(x.ndim):
+        x = (x[sl + np.index_exp[:-1]] + x[sl + np.index_exp[1:]]) / 2.0
+        sl += np.index_exp[:]
+    return x
+
+# prepare some coordinates, and attach rgb values to each
+r, theta, z = np.mgrid[0:1:11j, 0:np.pi*2:25j, -0.5:0.5:11j]
+x = r*np.cos(theta)
+y = r*np.sin(theta)
+
+rc, thetac, zc = midpoints(r), midpoints(theta), midpoints(z)
+
+# define a wobbly torus about [0.7, *, 0]
+sphere = (rc - 0.7)**2 + (zc + 0.2*np.cos(thetac*2))**2 < 0.2**2
+
+# combine the color components
+hsv = np.zeros(sphere.shape + (3,))
+hsv[..., 0] = thetac / (np.pi*2)
+hsv[..., 1] = rc
+hsv[..., 2] = zc + 0.5
+colors = matplotlib.colors.hsv_to_rgb(hsv)
+
+# and plot everything
+fig = plt.figure()
+ax = fig.gca(projection='3d')
+ax.voxels(x, y, z, sphere,
+          facecolors=colors,
+          edgecolors=np.clip(2*colors - 0.5, 0, 1),  # brighter
+          linewidth=0.5)
+
+plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2746,8 +2746,10 @@ def calc_arrow(uvw, angle=15):
 
     quiver3D = quiver
 
-    def voxels(self, filled, **kwargs):
+    def voxels(self, *args, **kwargs):
         """
+        ax.voxels([x, y, z,] /, filled, **kwargs)
+
         Plot a set of filled voxels
 
         All voxels are plotted as 1x1x1 cubes on the axis, with filled[0,0,0]
@@ -2758,10 +2760,16 @@ def voxels(self, filled, **kwargs):
 
         Parameters
         ----------
-        filled : np.array
+        filled : 3D np.array of bool
             A 3d array of values, with truthy values indicating which voxels
             to fill
 
+        x, y, z : 3D np.array, optional
+            The coordinates of the corners of the voxels. This should broadcast
+            to a shape one larger in every dimension than the shape of `filled`.
+            These arguments can be used to plot non-cubic voxels.
+            If not specified, defaults to increasing integers along each axis.
+
         facecolors, edgecolors : array_like
             The color to draw the faces and edges of the voxels. This parameter
             can be:
@@ -2791,10 +2799,33 @@ def voxels(self, filled, **kwargs):
 
         .. plot:: gallery/mplot3d/voxels.py
         .. plot:: gallery/mplot3d/voxels_rgb.py
+        .. plot:: gallery/mplot3d/voxels_torus.py
+        .. plot:: gallery/mplot3d/voxels_numpy_logo.py
         """
+
+        # work out which signature we should be using, and use it to parse
+        # the arguments. Name must be voxels for the correct error message
+        if len(args) >= 3:
+            # underscores indicate position only
+            def voxels(__x, __y, __z, filled, **kwargs):
+                return (__x, __y, __z), filled, kwargs
+        else:
+            def voxels(filled, **kwargs):
+                return None, filled, kwargs
+
+        xyz, filled, kwargs = voxels(*args, **kwargs)
+
         # check dimensions
         if filled.ndim != 3:
             raise ValueError("Argument filled must be 3-dimensional")
+        size = np.array(filled.shape, dtype=np.intp)
+
+        # check xyz coordinates, which are one larger than the filled shape
+        coord_shape = tuple(size + 1)
+        if xyz is None:
+            x, y, z = np.indices(coord_shape)
+        else:
+            x, y, z = (broadcast_to(c, coord_shape) for c in xyz)
 
         def _broadcast_color_arg(color, name):
             if np.ndim(color) in (0, 1):
@@ -2821,25 +2852,21 @@ def _broadcast_color_arg(color, name):
         edgecolors = _broadcast_color_arg(edgecolors, 'edgecolors')
 
         # always scale to the full array, even if the data is only in the center
-        self.auto_scale_xyz(
-            [0, filled.shape[0]],
-            [0, filled.shape[1]],
-            [0, filled.shape[2]]
-        )
+        self.auto_scale_xyz(x, y, z)
 
         # points lying on corners of a square
         square = np.array([
             [0, 0, 0],
             [0, 1, 0],
             [1, 1, 0],
             [1, 0, 0]
-        ])
+        ], dtype=np.intp)
 
         voxel_faces = defaultdict(list)
 
         def permutation_matrices(n):
             """ Generator of cyclic permutation matices """
-            mat = np.eye(n, dtype=int)
+            mat = np.eye(n, dtype=np.intp)
             for i in range(n):
                 yield mat
                 mat = np.roll(mat, 1, axis=0)
@@ -2848,7 +2875,7 @@ def permutation_matrices(n):
         # render
         for permute in permutation_matrices(3):
             # find the set of ranges to iterate over
-            pc, qc, rc = permute.T.dot(filled.shape[:3])
+            pc, qc, rc = permute.T.dot(size)
             pinds = np.arange(pc)
             qinds = np.arange(qc)
             rinds = np.arange(rc)
@@ -2890,7 +2917,20 @@ def permutation_matrices(n):
 
         # iterate over the faces, and generate a Poly3DCollection for each voxel
         polygons = {}
-        for coord, faces in voxel_faces.items():
+        for coord, faces_inds in voxel_faces.items():
+            # convert indices into 3D positions
+            if xyz is None:
+                faces = faces_inds
+            else:
+                faces = []
+                for face_inds in faces_inds:
+                    ind = face_inds[:, 0], face_inds[:, 1], face_inds[:, 2]
+                    face = np.empty(face_inds.shape)
+                    face[:, 0] = x[ind]
+                    face[:, 1] = y[ind]
+                    face[:, 2] = z[ind]
+                    faces.append(face)
+
             poly = art3d.Poly3DCollection(faces,
                 facecolors=facecolors[coord],
                 edgecolors=edgecolors[coord],
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/voxels-xyz.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/voxels-xyz.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -655,6 +655,43 @@ def test_alpha(self):
             assert voxels[coord], "faces returned for absent voxel"
             assert isinstance(poly, art3d.Poly3DCollection)
 
+    @image_comparison(
+        baseline_images=['voxels-xyz'],
+        extensions=['png'],
+        tol=0.01
+    )
+    def test_xyz(self):
+        fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
+
+        def midpoints(x):
+            sl = ()
+            for i in range(x.ndim):
+                x = (x[sl + np.index_exp[:-1]] +
+                     x[sl + np.index_exp[1:]]) / 2.0
+                sl += np.index_exp[:]
+            return x
+
+        # prepare some coordinates, and attach rgb values to each
+        r, g, b = np.indices((17, 17, 17)) / 16.0
+        rc = midpoints(r)
+        gc = midpoints(g)
+        bc = midpoints(b)
+
+        # define a sphere about [0.5, 0.5, 0.5]
+        sphere = (rc - 0.5)**2 + (gc - 0.5)**2 + (bc - 0.5)**2 < 0.5**2
+
+        # combine the color components
+        colors = np.zeros(sphere.shape + (3,))
+        colors[..., 0] = rc
+        colors[..., 1] = gc
+        colors[..., 2] = bc
+
+        # and plot everything
+        ax.voxels(r, g, b, sphere,
+                  facecolors=colors,
+                  edgecolors=np.clip(2*colors - 0.5, 0, 1),  # brighter
+                  linewidth=0.5)
+
 
 def test_inverted_cla():
     # Github PR #5450. Setting autoscale should reset
