Creating a 3D plot from a 3D numpy array

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.

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