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.

Asked By: pter

||

Source

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

Answered By: DSM

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.

Answered By: Mr Squid

Creating a 3D plot from a 3D numpy array

Question:

Answers: