Creating a 3D plot from a 3D numpy array

Question

Ok, so I feel like there should be an easy way to create a 3-dimensional scatter plot using matplotlib. I have a 3D numpy array (dset) with 0's where I don't want a point and 1's where I do, basically to plot it now I have to step through three for: loops as such:

for i in range(30):
 for x in range(60):
 for y in range(60):
 if dset[i, x, y] == 1:
 ax.scatter(x, y, -i, zdir='z', c= 'red')

Any suggestions on how I could accomplish this more efficiently? Any ideas would be greatly appreciated.

Answer 1

If you have a dset like that, and you want to just get the 1 values, you could use nonzero, which "returns a tuple of arrays, one for each dimension of a, containing the indices of the non-zero elements in that dimension.".

For example, we can make a simple 3d array:

>>> import numpy
>>> numpy.random.seed(29)
>>> d = numpy.random.randint(0, 2, size=(3,3,3))
>>> d
array([[[1, 1, 0],
 [1, 0, 0],
 [0, 1, 1]],
 [[0, 1, 1],
 [1, 0, 0],
 [0, 1, 1]],
 [[1, 1, 0],
 [0, 1, 0],
 [0, 0, 1]]])

and find where the nonzero elements are located:

>>> d.nonzero()
(array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2]), array([0, 0, 1, 2, 2, 0, 0, 1, 2, 2, 0, 0, 1, 2]), array([0, 1, 0, 1, 2, 1, 2, 0, 1, 2, 0, 1, 1, 2]))
>>> z,x,y = d.nonzero()

If we wanted a more complicated cut, we could have done something like (d > 3.4).nonzero() or something, as True has an integer value of 1 and counts as nonzero.

Finally, we plot:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x, y, -z, zdir='z', c= 'red')
plt.savefig("demo.png")

giving

demo 3d image

• pter: Beautiful. I have a a question about the (d > 3.4).nonzero() part though, would that return just the places where d has a value greater than 3.4?
• DSM: @pter: exactly right. (d > 3.4) gives a boolean array the same shape as d with True where the entry is > 3.4 and False elsewhere.
• Joran Beasley: sweet I didnt know about this (dont usually need 3d plots... but if i do this will be awesome!) thanks DSM
• MSS: Is there a way to draw surface plots in matplotlib using 3d numpy array.

Answer 2

If you wanted to avoid using the nonzero option (for example, if you had a 3D numpy array whose values were supposed to be the color values of the data points), you could do what you do, but save some lines of code by using ndenumerate.

Your example might become:

for index, x in np.ndenumerate(dset):
 if x == 1:
 ax.scatter(*index, c = 'red')

I guess the point is just that you dont need to have nested for loops to iterate through multidimensional numpy arrays.