Plotting 3D array of values as coloured points on a grid

Question

I have an (X,Y,Z) numpy array that describes every point inside a box. I would like to do a 3D plot of this data where the colour of the point at [x,y,z] is the value of that point in the array I have so far tried something along the lines of:

fig = plt.figure()
ax = plt.axes(projection='3d')
data = np.random.rand(3,4,5)
xs = np.arange(0,data.shape[0])
ys = np.arange(0,data.shape[1])
zs = np.arange(0,data.shape[2])
for x in xs:
    for y in ys:
        for z in zs:
            ax.scatter(x, y, z, c = data[x,y,z])
plt.show()

This correctly plots a point at every index, but doesn't change the colours based on the value. I have seen a few methods using ravel/reshaping data into a 1D array but that throws an error due to the fact that this method is only plotting one point at a time.

Is there a more sensible way of doing this than separately plotting each point?

(PS this is for visualizing FDTD simulations of EM propagation, so if you know of a better method particular to that then that would also be helpful)

Answer 1

Try something like this

import numpy as np
import matplotlib.pyplot as plt

data = np.random.rand(3,4,5)
x = np.indices(data.shape)[0]
y = np.indices(data.shape)[1]
z = np.indices(data.shape)[2]
col = data.flatten()

# 3D Plot
fig = plt.figure()
ax3D = fig.add_subplot(projection='3d')
cm = plt.colormaps['brg']
p3d = ax3D.scatter(x, y, z, c=col, cmap=cm)
plt.colorbar(p3d)

plt.show()