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